From 58c7b33cbd9bf1763bf426f5af4c6abce387f624 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=A1s=20Fantone?= <nicolas.fantone@klarna.com>
Date: Wed, 4 Sep 2024 10:27:34 +0100
Subject: [PATCH] Make annotate return instance of actual mockset class

---
 django_mock_queries/query.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/django_mock_queries/query.py b/django_mock_queries/query.py
index ca39759..d58f1ed 100644
--- a/django_mock_queries/query.py
+++ b/django_mock_queries/query.py
@@ -127,7 +127,7 @@ def annotate(self, **kwargs):
                 row._annotated_fields.append(key)
                 setattr(row, key, get_attribute(row, value)[0])
 
-        return MockSet(*results, clone=self)
+        return self._mockset_class()(*results, clone=self)
 
     def aggregate(self, *args, **kwargs):
         result = {}