LITE-31232 Units for bulk_relate_cqrs_serialization

dj_cqrs/utils.py

@@ -88,19 +88,19 @@ def apply_query_timeouts(model_cls):  # pragma: no cover
 class _BulkRelateCM(ContextDecorator):
     def __init__(self, cqrs_id=None):
         self._cqrs_id = cqrs_id
-        self._mapping = defaultdict(lambda: defaultdict(list))
+        self._mapping = defaultdict(lambda: defaultdict(set))
         self._cache = {}
 
-    def register(self, instance, using):
+    def register(self, instance, using=None):
         instance_cqrs_id = getattr(instance, 'CQRS_ID', None)
-        if self._cqrs_id and instance_cqrs_id != self._cqrs_id:
+        if (not instance_cqrs_id) or (self._cqrs_id and instance_cqrs_id != self._cqrs_id):
             return
 
-        self._mapping[instance_cqrs_id][using].append(instance.pk)
+        self._mapping[instance_cqrs_id][using].add(instance.pk)
 
-    def get_cached_instance(self, instance, using):
+    def get_cached_instance(self, instance, using=None):
         instance_cqrs_id = getattr(instance, 'CQRS_ID', None)
-        if self._cqrs_id and instance_cqrs_id != self._cqrs_id:
+        if (not instance_cqrs_id) or (self._cqrs_id and instance_cqrs_id != self._cqrs_id):
             return
 
         instance_pk = instance.pk
@@ -115,15 +115,20 @@ def get_cached_instance(self, instance, using):
         qs = instance.__class__._default_manager.using(using)
         instances_cache = {
             instance.pk: instance
-            for instance in instance.__class__.relate_cqrs_serialization(qs).filter(
+            for instance in instance.__class__.relate_cqrs_serialization(qs)
+            .filter(
                 pk__in=cached_pks,
-            ).order_by().all()
+            )
+            .order_by()
+            .all()
         }
-        self._cache.update({
-            instance_cqrs_id: {
-                using: instances_cache,
-            },
-        })
+        self._cache.update(
+            {
+                instance_cqrs_id: {
+                    using: instances_cache,
+                },
+            }
+        )
         return instances_cache.get(instance_pk)
 
     def __enter__(self):

tests/test_utils.py

@@ -10,14 +10,18 @@
 from uuid import UUID
 
 import pytest
+from django.db import transaction
 
+from dj_cqrs.state import cqrs_state
 from dj_cqrs.utils import (
     apply_query_timeouts,
+    bulk_relate_cqrs_serialization,
     get_delay_queue_max_size,
     get_json_valid_value,
     get_message_expiration_dt,
     get_messages_prefetch_count_per_worker,
 )
+from tests.dj_master import models as master_models
 from tests.dj_replica import models
 
 
@@ -109,3 +113,127 @@ def test_apply_query_timeouts(settings, engine, p_count):
         assert apply_query_timeouts(models.BasicFieldsModelRef) is None
 
     assert p.call_count == p_count
+
+
+@pytest.mark.django_db(transaction=True)
+def test_bulk_relate_cqrs_serialization_simple_model(mocker):
+    produce_mock = mocker.patch('dj_cqrs.controller.producer.produce')
+
+    @bulk_relate_cqrs_serialization()
+    def func():
+        assert cqrs_state.bulk_relate_cm
+
+        instance = master_models.SimplestModel(id=1)
+        instance.save()
+
+    assert cqrs_state.bulk_relate_cm is None
+    func()
+
+    assert master_models.SimplestModel.objects.count() == 1
+    assert produce_mock.call_count == 1
+    assert cqrs_state.bulk_relate_cm is None
+
+
+@pytest.mark.django_db(transaction=True)
+def test_bulk_relate_cqrs_serialization_serialized_model(mocker):
+    produce_mock = mocker.patch('dj_cqrs.controller.producer.produce')
+
+    assert cqrs_state.bulk_relate_cm is None
+    with bulk_relate_cqrs_serialization(cqrs_id=master_models.Author.CQRS_ID):
+        bulk_relate_cm = cqrs_state.bulk_relate_cm
+
+        with transaction.atomic(savepoint=False):
+            master_models.Author.objects.create(id=1)
+
+            assert bulk_relate_cm
+            assert bulk_relate_cm._mapping
+            assert not bulk_relate_cm._cache
+
+        assert bulk_relate_cm._cache
+
+    assert master_models.Author.objects.count() == 1
+    assert produce_mock.call_count == 1
+    assert cqrs_state.bulk_relate_cm is None
+
+
+def test_bulk_relate_cqrs_serialization_error():
+    assert cqrs_state.bulk_relate_cm is None
+
+    try:
+        with bulk_relate_cqrs_serialization(cqrs_id=master_models.Author.CQRS_ID):
+            assert cqrs_state.bulk_relate_cm
+            raise ValueError
+    except ValueError:
+        pass
+
+    assert cqrs_state.bulk_relate_cm is None
+
+
+@pytest.mark.django_db(transaction=True)
+def test_bulk_relate_cqrs_serialization_register():
+    author1 = master_models.Author(id=1)
+    author2 = master_models.Author(id=2)
+
+    with bulk_relate_cqrs_serialization(cqrs_id=master_models.Author.CQRS_ID):
+        bulk_relate_cm = cqrs_state.bulk_relate_cm
+        bulk_relate_cm.register(ValueError)
+        bulk_relate_cm.register(master_models.FilteredSimplestModel())
+        bulk_relate_cm.register(author1, 'default')
+        bulk_relate_cm.register(author1, 'default')
+        bulk_relate_cm.register(author1, 'other')
+        bulk_relate_cm.register(author2, 'other')
+        bulk_relate_cm.register(author2)
+
+        assert bulk_relate_cm._mapping == {
+            master_models.Author.CQRS_ID: {
+                'default': {1},
+                'other': {1, 2},
+                None: {2},
+            },
+        }
+
+    assert cqrs_state.bulk_relate_cm is None
+
+
+@pytest.mark.django_db(transaction=True)
+def test_bulk_relate_cqrs_serialization_get_cached_instance(mocker, django_assert_num_queries):
+    produce_mock = mocker.patch('dj_cqrs.controller.producer.produce')
+
+    simple = master_models.SimplestModel.objects.create(id=1)
+
+    with bulk_relate_cqrs_serialization():
+        bulk_relate_cm = cqrs_state.bulk_relate_cm
+
+        with transaction.atomic():
+            author1 = master_models.Author.objects.create(id=1)
+            author1.name = 'new'
+            author1.save()
+            author2 = master_models.Author.objects.create(id=2)
+
+            af = master_models.AutoFieldsModel.objects.using('default').create()
+            publisher = master_models.Publisher.objects.create(id=3)
+
+        assert produce_mock.call_count == 4
+        assert bulk_relate_cm._cache == {
+            master_models.Author.CQRS_ID: {
+                'default': {
+                    1: author1,
+                    2: author2,
+                },
+            },
+        }
+
+        assert bulk_relate_cm.get_cached_instance(publisher) is None
+        assert bulk_relate_cm.get_cached_instance(ValueError, 'test') is None
+
+        with django_assert_num_queries(0):
+            assert bulk_relate_cm.get_cached_instance(simple) is None
+            assert bulk_relate_cm.get_cached_instance(author1, 'default') == author1
+            assert bulk_relate_cm.get_cached_instance(author1, 'default') == author1
+            assert bulk_relate_cm.get_cached_instance(author1, 'other') is None
+            assert bulk_relate_cm.get_cached_instance(author2, 'default') == author2
+            assert bulk_relate_cm.get_cached_instance(author2) is None
+            assert bulk_relate_cm.get_cached_instance(master_models.Author(id=3)) is None
+            assert bulk_relate_cm.get_cached_instance(af) is None
+
+    assert cqrs_state.bulk_relate_cm is None