From a39a04fa522d0d68efafcc07513351c1045fe58b Mon Sep 17 00:00:00 2001
From: M1ha <work_shvein_mihail@mail.ru>
Date: Fri, 4 Oct 2019 17:54:47 +0500
Subject: [PATCH] Added tests and fixes for bulk_create_returning and
 save_returning

---
 src/django_clickhouse/models.py |  5 ++++-
 tests/models.py                 | 20 +++++++++++++++++---
 tests/test_models.py            | 23 +++++++++++++++++++++++
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/src/django_clickhouse/models.py b/src/django_clickhouse/models.py
index 52c979e..97815f0 100644
--- a/src/django_clickhouse/models.py
+++ b/src/django_clickhouse/models.py
@@ -102,7 +102,7 @@ class ClickHouseSyncQuerySetMixin(ClickHouseSyncRegisterMixin):
 
 
 # I add library dependant mixins to base classes only if libraries are installed
-qs_bases = [ClickHouseSyncQuerySetMixin, DjangoQuerySet]
+qs_bases = [ClickHouseSyncQuerySetMixin]
 
 if not getattr(UpdateReturningMixin, 'fake', False):
     qs_bases.append(ClickHouseSyncUpdateReturningQuerySetMixin)
@@ -110,6 +110,9 @@ if not getattr(UpdateReturningMixin, 'fake', False):
 if not getattr(BulkUpdateManagerMixin, 'fake', False):
     qs_bases.append(ClickHouseSyncBulkUpdateQuerySetMixin)
 
+
+# QuerySet must be the last one, so it can be redeclared in mixins
+qs_bases.append(DjangoQuerySet)
 ClickHouseSyncQuerySet = type('ClickHouseSyncModelQuerySet', tuple(qs_bases), {})
 
 
diff --git a/tests/models.py b/tests/models.py
index b27b271..78c8e21 100644
--- a/tests/models.py
+++ b/tests/models.py
@@ -2,17 +2,31 @@
 This file contains sample models to use in tests
 """
 from django.db import models
+from django.db.models.manager import BaseManager
+from django_pg_returning.models import UpdateReturningModel
 
-from django_clickhouse.models import ClickHouseSyncModel
+from django_clickhouse.models import ClickHouseSyncModel, ClickHouseSyncQuerySet
 
 
-class TestModel(ClickHouseSyncModel):
+class TestQuerySet(ClickHouseSyncQuerySet):
+    pass
+
+
+class TestManager(BaseManager.from_queryset(TestQuerySet)):
+    pass
+
+
+class TestModel(UpdateReturningModel, ClickHouseSyncModel):
+    objects = TestManager()
+
     value = models.IntegerField()
     created_date = models.DateField()
     created = models.DateTimeField()
 
 
-class SecondaryTestModel(ClickHouseSyncModel):
+class SecondaryTestModel(UpdateReturningModel, ClickHouseSyncModel):
+    objects = TestManager()
+
     value = models.IntegerField()
     created_date = models.DateField()
     created = models.DateTimeField()
diff --git a/tests/test_models.py b/tests/test_models.py
index 0de5c92..978cdc2 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -102,6 +102,16 @@ class TestOperations(TransactionTestCase):
         self.assertListEqual([('update', "%s.%d" % (self.db_alias, item.id)) for item in self.before_op_items],
                              self.storage.get_operations(self.clickhouse_model.get_import_key(), 10))
 
+    def test_bulk_create_returning(self):
+        items = [
+            self.django_model(created_date=datetime.date.today(), created=datetime.datetime.now(), value=i)
+            for i in range(5)
+        ]
+        items = self.django_model.objects.bulk_create_returning(items)
+        self.assertEqual(5, len(items))
+        self.assertSetEqual({('insert', "%s.%d" % (self.db_alias, instance.pk)) for instance in items},
+                            set(self.storage.get_operations(self.clickhouse_model.get_import_key(), 10)))
+
     def test_qs_update_returning(self):
         self.django_model.objects.filter(pk=1).update_returning(created_date=datetime.date.today())
         self.assertListEqual([('update', "%s.1" % self.db_alias)],
@@ -123,6 +133,19 @@ class TestOperations(TransactionTestCase):
         self.assertListEqual([('delete', "%s.%d" % (self.db_alias, item.id)) for item in self.before_op_items],
                              self.storage.get_operations(self.clickhouse_model.get_import_key(), 10))
 
+    def test_save_returning(self):
+        # INSERT operation
+        instance = self.django_model(created_date=datetime.date.today(), created=datetime.datetime.now(), value=2)
+        instance.save_returning()
+        self.assertListEqual([('insert', "%s.%d" % (self.db_alias, instance.pk))],
+                             self.storage.get_operations(self.clickhouse_model.get_import_key(), 10))
+
+        # UPDATE operation
+        instance.save_returning()
+        self.assertListEqual([('insert', "%s.%d" % (self.db_alias, instance.pk)),
+                              ('update', "%s.%d" % (self.db_alias, instance.pk))],
+                             self.storage.get_operations(self.clickhouse_model.get_import_key(), 10))
+
     def test_delete(self):
         instance = self.django_model.objects.get(pk=1)
         instance.delete()