diff --git a/src/modelsubquery/functions.py b/src/modelsubquery/functions.py index 33652bc..e9d2c93 100644 --- a/src/modelsubquery/functions.py +++ b/src/modelsubquery/functions.py @@ -13,7 +13,7 @@ def _model_fields(model, fields): then return all the declared fields. """ # TODO: the pk/id field should always be returned - declared = {f.name for f in model._meta.get_fields()} + declared = {f.column for f in model._meta.local_concrete_fields} if fields is None: return declared diff --git a/testproject/testapp/models.py b/testproject/testapp/models.py index 8451ee8..56a2fed 100644 --- a/testproject/testapp/models.py +++ b/testproject/testapp/models.py @@ -11,6 +11,9 @@ class Book(models.Model): rating = models.IntegerField(blank=True, null=True) published = models.DateField(default=date.today) has_cover = models.BooleanField(default=True) + author = models.ForeignKey( + "Person", on_delete=models.CASCADE, related_name="books", null=True + ) class PersonQuerySet(models.QuerySet): @@ -29,3 +32,8 @@ class Person(models.Model): birth = models.DateField(default=date.today) objects = PersonQuerySet.as_manager() + + +class Shelve(models.Model): + title = models.CharField(max_length=100) + books = models.ManyToManyField("Book", related_name="shelves") diff --git a/testproject/testapp/tests.py b/testproject/testapp/tests.py index 1954b6b..450572a 100644 --- a/testproject/testapp/tests.py +++ b/testproject/testapp/tests.py @@ -4,7 +4,21 @@ from django.test import TestCase from model_bakery import baker -from .models import Book, Person +from modelsubquery.functions import _model_fields + +from .models import Book, Person, Shelve + + +class ModelFieldTest(TestCase): + def test_model_fields_book(self): + fields = _model_fields(Book, None) + self.assertEqual( + fields, {"id", "title", "author_id", "published", "rating", "has_cover"} + ) + + def test_model_fields_shelve(self): + fields = _model_fields(Shelve, None) + self.assertEqual(fields, {"id", "title"}) @time_machine.travel("2000-01-01") @@ -79,3 +93,16 @@ def test_deffered_fields(self): self.assertEqual(person.book_of_year.title, "test") with self.assertNumQueries(1): self.assertEqual(person.book_of_year.rating, 5) + + def test_with_related_field(self): + author = baker.make(Person) + book = baker.make(Book, author=author) + + # Yes he wrote a book in his birthyear! ;-) + + person = Person.objects.with_book_of_the_year().get() + self.assertEqual(person.book_of_year, book) + + # The relation is there, but it's lazy! + with self.assertNumQueries(1): + self.assertEqual(person.book_of_year.author, author)