Skip to content

Commit

Permalink
add subqueries tests
Browse files Browse the repository at this point in the history
  • Loading branch information
timgraham authored and WaVEV committed Jan 18, 2025
1 parent fd57ce8 commit 27ec33c
Show file tree
Hide file tree
Showing 3 changed files with 75 additions and 2 deletions.
4 changes: 2 additions & 2 deletions django_mongodb_backend/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ def _get_column_from_expression(self, expr, alias):
Create a column named `alias` from the given expression to hold the
aggregate value.
"""
column_target = expr.output_field.__class__()
column_target = expr.output_field.clone()
column_target.db_column = alias
column_target.set_attributes_from_name(alias)
return Col(self.collection_name, column_target)
Expand Down Expand Up @@ -81,7 +81,7 @@ def _prepare_expressions_for_pipeline(self, expression, target, annotation_group
alias = (
f"__aggregation{next(annotation_group_idx)}" if sub_expr != expression else target
)
column_target = sub_expr.output_field.__class__()
column_target = sub_expr.output_field.clone()
column_target.db_column = alias
column_target.set_attributes_from_name(alias)
inner_column = Col(self.collection_name, column_target)
Expand Down
8 changes: 8 additions & 0 deletions tests/model_fields_/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ class Data(models.Model):
integer = models.IntegerField(db_column="custom_column")
auto_now = models.DateTimeField(auto_now=True)
auto_now_add = models.DateTimeField(auto_now_add=True)
json_value = models.JSONField(default=dict)


class Address(models.Model):
Expand All @@ -119,3 +120,10 @@ class Author(models.Model):
class Book(models.Model):
name = models.CharField(max_length=100)
author = EmbeddedModelField(Author)


class Library(models.Model):
name = models.CharField(max_length=100)
books = models.ManyToManyField("Book", related_name="libraries")
location = models.CharField(max_length=100, null=True, blank=True)
best_seller = models.CharField(max_length=100, null=True, blank=True)
65 changes: 65 additions & 0 deletions tests/model_fields_/test_embedded_model.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
from django.core.exceptions import ValidationError
from django.db import models
from django.db.models import (
Exists,
OuterRef,
Subquery,
)
from django.test import SimpleTestCase, TestCase
from django.test.utils import isolate_apps

Expand All @@ -11,6 +16,7 @@
Book,
Data,
Holder,
Library,
)


Expand Down Expand Up @@ -123,3 +129,62 @@ class MyModel(models.Model):
self.assertEqual(
msg, "Embedded models cannot have relational fields (Target.key is a ForeignKey)."
)


class SubqueryExistsTest(TestCase):
def setUp(self):
# Create test data
address1 = Address.objects.create(city="New York", state="NY", zip_code=10001)
address2 = Address.objects.create(city="Boston", state="MA", zip_code=20002)
author1 = Author.objects.create(name="Alice", age=30, address=address1)
author2 = Author.objects.create(name="Bob", age=40, address=address2)
book1 = Book.objects.create(name="Book A", author=author1)
book2 = Book.objects.create(name="Book B", author=author2)
Book.objects.create(name="Book C", author=author2)
Book.objects.create(name="Book D", author=author2)
Book.objects.create(name="Book E", author=author1)

library1 = Library.objects.create(
name="Central Library", location="Downtown", best_seller="Book A"
)
library2 = Library.objects.create(
name="Community Library", location="Suburbs", best_seller="Book A"
)

# Add books to libraries
library1.books.add(book1, book2)
library2.books.add(book2)

def test_exists_subquery(self):
subquery = Book.objects.filter(
author__name=OuterRef("name"), author__address__city="Boston"
)
queryset = Author.objects.filter(Exists(subquery))

self.assertEqual(queryset.count(), 1)

def test_in_subquery(self):
subquery = Author.objects.filter(age__gt=35).values("name")
queryset = Book.objects.filter(author__name__in=Subquery(subquery)).order_by("name")

self.assertEqual(queryset.count(), 3)
self.assertQuerySetEqual(queryset, ["Book B", "Book C", "Book D"], lambda book: book.name)

def test_range_query(self):
queryset = Author.objects.filter(age__range=(25, 45)).order_by("name")

self.assertEqual(queryset.count(), 2)
self.assertQuerySetEqual(queryset, ["Alice", "Bob"], lambda author: author.name)

def test_exists_with_foreign_object(self):
subquery = Library.objects.filter(best_seller=OuterRef("name"))
queryset = Book.objects.filter(Exists(subquery))

self.assertEqual(queryset.count(), 1)
self.assertEqual(queryset.first().name, "Book A")

def test_foreign_field_with_ranges(self):
queryset = Library.objects.filter(books__author__age__range=(25, 35))

self.assertEqual(queryset.count(), 1)
self.assertEqual(queryset.first().name, "Central Library")

0 comments on commit 27ec33c

Please sign in to comment.