From 686e1ab81aae938a4a85680a6ed581d90b7ff11d Mon Sep 17 00:00:00 2001
From: Zach Aysan <zachaysan@gmail.com>
Date: Fri, 26 Apr 2024 10:25:00 -0400
Subject: [PATCH] fix: Filter versioned features (#3756)

---
 api/features/versioning/versioning_service.py | 10 +++----
 api/features/views.py                         | 10 +++++--
 .../unit/features/test_unit_features_views.py | 30 +++++++++++++++++--
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/api/features/versioning/versioning_service.py b/api/features/versioning/versioning_service.py
index bcac998faebe..abeb3a50c778 100644
--- a/api/features/versioning/versioning_service.py
+++ b/api/features/versioning/versioning_service.py
@@ -3,15 +3,13 @@
 from django.db.models import Prefetch, Q, QuerySet
 from django.utils import timezone
 
+from environments.models import Environment
 from features.models import FeatureState
 from features.versioning.models import EnvironmentFeatureVersion
 
-if typing.TYPE_CHECKING:
-    from environments.models import Environment
-
 
 def get_environment_flags_queryset(
-    environment: "Environment", feature_name: str = None
+    environment: Environment, feature_name: str = None
 ) -> QuerySet[FeatureState]:
     """
     Get a queryset of the latest live versions of an environments' feature states
@@ -21,14 +19,14 @@ def get_environment_flags_queryset(
 
 
 def get_environment_flags_list(
-    environment: "Environment",
+    environment: Environment,
     feature_name: str = None,
     additional_filters: Q = None,
     additional_select_related_args: typing.Iterable[str] = None,
     additional_prefetch_related_args: typing.Iterable[
         typing.Union[str, Prefetch]
     ] = None,
-) -> typing.List["FeatureState"]:
+) -> list[FeatureState]:
     """
     Get a list of the latest committed versions of FeatureState objects that are
     associated with the given environment. Can be filtered to remove segment /
diff --git a/api/features/views.py b/api/features/views.py
index 4d99588258d8..323bef8dd175 100644
--- a/api/features/views.py
+++ b/api/features/views.py
@@ -258,12 +258,16 @@ def apply_state_to_queryset(
         if not getattr(self, "environment", None):
             self.environment = Environment.objects.get(id=environment_id)
 
-        feature_states = FeatureState.objects.get_live_feature_states(
+        feature_states = get_environment_flags_list(
             environment=self.environment,
-            additional_filters=base_q & filter_search_q & filter_enabled_q,
+            additional_filters=base_q,
         )
 
-        feature_ids = {fs.feature_id for fs in feature_states}
+        feature_ids = FeatureState.objects.filter(
+            filter_search_q & filter_enabled_q,
+            id__in=[fs.id for fs in feature_states],
+        ).values_list("feature_id", flat=True)
+
         return queryset.filter(id__in=feature_ids)
 
     @swagger_auto_schema(
diff --git a/api/tests/unit/features/test_unit_features_views.py b/api/tests/unit/features/test_unit_features_views.py
index fbdd1f1117c6..7c175fa43e0c 100644
--- a/api/tests/unit/features/test_unit_features_views.py
+++ b/api/tests/unit/features/test_unit_features_views.py
@@ -2632,13 +2632,16 @@ def test_list_features_with_feature_state(
 
 def test_list_features_with_filter_by_value_search_string_and_int(
     staff_client: APIClient,
+    staff_user: FFAdminUser,
     project: Project,
     feature: Feature,
     with_project_permissions: WithProjectPermissionsCallable,
-    environment: Environment,
+    environment_v2_versioning: Environment,
 ) -> None:
     # Given
     with_project_permissions([VIEW_PROJECT])
+    environment = environment_v2_versioning
+
     feature2 = Feature.objects.create(
         name="another_feature", project=project, initial_value="initial_value"
     )
@@ -2654,10 +2657,33 @@ def test_list_features_with_filter_by_value_search_string_and_int(
         project=project,
     )
 
-    feature_state1 = feature.feature_states.filter(environment=environment).first()
+    environment_feature_version1 = EnvironmentFeatureVersion.objects.create(
+        environment=environment,
+        feature=feature,
+    )
+
+    feature_state1 = FeatureState.objects.filter(
+        environment_feature_version=environment_feature_version1
+    ).first()
     feature_state1.enabled = True
     feature_state1.save()
 
+    # Create a secondary feature state that will be versioned in the past.
+    environment_feature_version1b = EnvironmentFeatureVersion.objects.create(
+        environment=environment,
+        feature=feature,
+    )
+
+    feature_state1b = FeatureState.objects.filter(
+        environment_feature_version=environment_feature_version1b
+    ).first()
+    feature_state1b.enabled = False
+    feature_state1b.save()
+
+    environment_feature_version1b.publish(staff_user)
+
+    environment_feature_version1.publish(staff_user)
+
     feature_state_value1 = feature_state1.feature_state_value
     feature_state_value1.string_value = None
     feature_state_value1.integer_value = 1945