From 263ede6e674f33763670b20074fafc4185c7adf7 Mon Sep 17 00:00:00 2001
From: Matthew Elwell <matthew.elwell@flagsmith.com>
Date: Fri, 10 May 2024 14:49:01 +0100
Subject: [PATCH] Add some additional checks on number of queries

---
 .../identities/test_edge_identity_models.py   |  9 +++++---
 ...test_unit_versioning_versioning_service.py | 21 +++++++++++++------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/api/tests/unit/edge_api/identities/test_edge_identity_models.py b/api/tests/unit/edge_api/identities/test_edge_identity_models.py
index b10367febfde..92b1f89320bf 100644
--- a/api/tests/unit/edge_api/identities/test_edge_identity_models.py
+++ b/api/tests/unit/edge_api/identities/test_edge_identity_models.py
@@ -6,6 +6,7 @@
 from django.utils import timezone
 from flag_engine.features.models import FeatureModel, FeatureStateModel
 from freezegun import freeze_time
+from pytest_django import DjangoAssertNumQueries
 from pytest_mock import MockerFixture
 
 from edge_api.identities.models import EdgeIdentity
@@ -484,6 +485,7 @@ def test_get_all_feature_states_post_v2_versioning_migration(
     segment_featurestate: FeatureState,
     edge_identity_model: EdgeIdentity,
     mocker: MockerFixture,
+    django_assert_num_queries: DjangoAssertNumQueries,
 ) -> None:
     """
     Specific test to reproduce an issue seen after migrating our staging environment to
@@ -507,9 +509,10 @@ def test_get_all_feature_states_post_v2_versioning_migration(
     edge_identity_dynamo_wrapper_mock.get_segment_ids.return_value = [segment.id]
 
     # When
-    feature_states, identity_override_feature_names = (
-        edge_identity_model.get_all_feature_states()
-    )
+    with django_assert_num_queries(4):
+        feature_states, identity_override_feature_names = (
+            edge_identity_model.get_all_feature_states()
+        )
 
     # Then
     assert len(feature_states) == 1
diff --git a/api/tests/unit/features/versioning/test_unit_versioning_versioning_service.py b/api/tests/unit/features/versioning/test_unit_versioning_versioning_service.py
index 422fde91f79b..fb54f9588d07 100644
--- a/api/tests/unit/features/versioning/test_unit_versioning_versioning_service.py
+++ b/api/tests/unit/features/versioning/test_unit_versioning_versioning_service.py
@@ -2,6 +2,7 @@
 
 from django.db.models import Q
 from django.utils import timezone
+from pytest_django import DjangoAssertNumQueries
 
 from environments.identities.models import Identity
 from environments.models import Environment
@@ -18,7 +19,9 @@
 
 
 def test_get_environment_flags_queryset_returns_only_latest_versions(
-    feature, environment
+    feature: Feature,
+    environment: Environment,
+    django_assert_num_queries: DjangoAssertNumQueries,
 ):
     # Given
     feature_state_v1 = FeatureState.objects.get(
@@ -31,7 +34,11 @@ def test_get_environment_flags_queryset_returns_only_latest_versions(
     feature_state_v1.clone(env=environment, as_draft=True)  # draft feature state
 
     # When
-    feature_states = get_environment_flags_queryset(environment=environment)
+    with django_assert_num_queries(2):
+        feature_states = get_environment_flags_queryset(environment=environment)
+
+        # trigger the queryset to execute and ensure the number of queries is correct
+        list(feature_states)
 
     # Then
     assert feature_states.count() == 1
@@ -115,6 +122,7 @@ def test_get_environment_flags_v2_versioning_returns_latest_live_versions_of_fea
     environment_v2_versioning: Environment,
     feature: Feature,
     admin_user: FFAdminUser,
+    django_assert_num_queries: DjangoAssertNumQueries,
 ) -> None:
     # Given
     # a second feature with its corresponding environment feature version
@@ -139,10 +147,11 @@ def test_get_environment_flags_v2_versioning_returns_latest_live_versions_of_fea
     environment_feature_1_version_2.publish(admin_user)
 
     # When
-    environment_feature_states = get_environment_flags_list(
-        environment=environment_v2_versioning,
-        additional_filters=Q(feature_segment=None, identity=None),
-    )
+    with django_assert_num_queries(2):
+        environment_feature_states = get_environment_flags_list(
+            environment=environment_v2_versioning,
+            additional_filters=Q(feature_segment=None, identity=None),
+        )
 
     # Then
     assert set(environment_feature_states) == {