add tests for set_siae_found_list with semantic search

lemarche/tenders/tests.py

@@ -2,6 +2,7 @@
 from datetime import timedelta
 from importlib import import_module
 from random import randint
+from unittest import mock
 
 from django.apps import apps
 from django.contrib.gis.geos import Point
@@ -164,6 +165,99 @@ def test_set_slug(self):
     #     self.assertRaises(IntegrityError, tender.save)
 
 
+class TenderModelMatchingTest(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        cls.sector = SectorFactory()
+
+        # siae found by presta_type
+        cls.siae_one = SiaeFactory(
+            is_active=True,
+            kind=siae_constants.KIND_AI,
+            presta_type=[siae_constants.PRESTA_PREST, siae_constants.PRESTA_BUILD],
+            geo_range=siae_constants.GEO_RANGE_COUNTRY,
+        )
+        cls.siae_one.sectors.add(cls.sector)
+
+        # siae found by presta_type and semantic search
+        cls.siae_two = SiaeFactory(
+            is_active=True,
+            kind=siae_constants.KIND_ESAT,
+            presta_type=[siae_constants.PRESTA_BUILD],
+            geo_range=siae_constants.GEO_RANGE_COUNTRY,
+        )
+        cls.siae_two.sectors.add(cls.sector)
+
+        # siaes found by mocked semantic search
+        cls.siae_three = SiaeFactory()
+        cls.siae_four = SiaeFactory()
+
+        # siae not found
+        cls.siae_five = SiaeFactory()
+
+    def test_set_siae_found_list_without_semantic_search(self):
+        with mock.patch(
+            "lemarche.tenders.models.api_elasticsearch.siaes_similarity_search"
+        ) as mock_siaes_similarity_search:
+            tender = TenderFactory(
+                presta_type=[siae_constants.PRESTA_BUILD],
+                sectors=[self.sector],
+                is_country_area=True,
+                validated_at=None,
+            )
+
+            siaes_found = Siae.objects.filter_with_tender(tender)
+            tender.set_siae_found_list()
+            tender.refresh_from_db()
+            self.assertEqual(list(siaes_found), list(tender.siaes.all()))
+
+        mock_siaes_similarity_search.assert_not_called()
+
+    def test_set_siae_found_list_with_semantic_search(self):
+        with mock.patch(
+            "lemarche.tenders.models.api_elasticsearch.siaes_similarity_search"
+        ) as mock_siaes_similarity_search:
+            mock_siaes_similarity_search.return_value = [self.siae_two.pk, self.siae_three.pk, self.siae_four.pk]
+            tender = TenderFactory(
+                presta_type=[siae_constants.PRESTA_BUILD],
+                sectors=[self.sector],
+                is_country_area=True,
+                with_semantic_matching=True,
+                validated_at=None,
+            )
+
+            siaes_found = Siae.objects.filter_with_tender(tender)
+            tender.set_siae_found_list()
+            tender.refresh_from_db()
+
+            self.assertEqual(tender.siaes.count(), 4)
+            for siae in list(siaes_found) + [self.siae_three, self.siae_four]:
+                with self.subTest(siae=siae):
+                    self.assertIn(siae, tender.siaes.all())
+
+            # test ai_had_found_it_too field value
+            self.assertEqual(TenderSiae.objects.get(tender=tender, siae=self.siae_one).ai_had_found_it_too, False)
+            tender_siae_two = TenderSiae.objects.get(tender=tender, siae=self.siae_two)
+            self.assertEqual(tender_siae_two.ai_had_found_it_too, True)
+            for siae in [self.siae_two, self.siae_three, self.siae_four]:
+                with self.subTest(siae=siae):
+                    self.assertEqual(TenderSiae.objects.get(tender=tender, siae=siae).ai_had_found_it_too, True)
+
+            # test source
+            for siae in [self.siae_one, self.siae_two]:
+                with self.subTest(siae=siae):
+                    self.assertEqual(
+                        TenderSiae.objects.get(tender=tender, siae=siae).source,
+                        tender_constants.TENDER_SIAE_SOURCE_EMAIL,
+                    )
+            for siae in [self.siae_three, self.siae_four]:
+                with self.subTest(siae=siae):
+                    self.assertEqual(
+                        TenderSiae.objects.get(tender=tender, siae=siae).source, tender_constants.TENDER_SIAE_SOURCE_AI
+                    )
+        mock_siaes_similarity_search.assert_called_once()
+
+
 class TenderModelQuerysetTest(TestCase):
     @classmethod
     def setUpTestData(cls):