fix: v1-override-limit-includes-all-feature-segments-in-count

api/features/feature_segments/limits.py

@@ -32,16 +32,14 @@ def _check(left: int, right: int) -> bool:
         )
         q = q & Q(environment_feature_version__in=latest_versions)
 
-    existing_overridden_segment_ids = set(
-        environment.feature_segments.filter(q).values_list("segment_id", flat=True)
-    )
-    segment_override_count = len(existing_overridden_segment_ids)
+    existing_overrides = environment.feature_segments.filter(q)
+    segment_override_count = existing_overrides.count()
 
-    extra = len(segment_ids_to_create_overrides) - len(
-        set(segment_ids_to_delete_overrides).intersection(
-            existing_overridden_segment_ids
-        )
+    existing_segment_ids = set(existing_overrides.values_list("segment_id", flat=True))
+    to_delete_count = len(
+        set(segment_ids_to_delete_overrides).intersection(existing_segment_ids)
     )
+    extra = len(segment_ids_to_create_overrides) - to_delete_count
 
     return _check(
         segment_override_count + extra,

api/tests/unit/features/feature_segments/test_unit_feature_segments_limits.py

@@ -1,10 +1,100 @@
 from environments.models import Environment
 from features.feature_segments.limits import exceeds_segment_override_limit
-from features.models import Feature
+from features.models import Feature, FeatureSegment
+from features.versioning.models import EnvironmentFeatureVersion
 from projects.models import Project
 from segments.models import Segment
 
 
+def test_segment_override_limit_includes_all_feature_segments_states(
+    project: Project,
+    environment: Environment,
+) -> None:
+    # Given
+    project.max_segment_overrides_allowed = 3
+    project.save()
+
+    features = [
+        Feature.objects.create(name=f"feature_{i}", project=project) for i in range(5)
+    ]
+    segment = Segment.objects.create(name="shared_segment", project=project)
+
+    for feature in features:
+        FeatureSegment.objects.create(
+            feature=feature, segment=segment, environment=environment
+        )
+
+    # When
+    result = exceeds_segment_override_limit(environment=environment)
+
+    # Then
+    assert result is True
+
+
+def test_segment_override_limit_is_triggered_with_distinct_segments_per_feature_states(
+    project: Project,
+    environment: Environment,
+) -> None:
+    # Given
+    project.max_segment_overrides_allowed = 3
+    project.save()
+
+    features = [
+        Feature.objects.create(name=f"feature_{i}", project=project) for i in range(5)
+    ]
+    segments = [
+        Segment.objects.create(name=f"segment_{i}", project=project) for i in range(5)
+    ]
+
+    for feature, segment in zip(features, segments):
+        FeatureSegment.objects.create(
+            feature=feature, segment=segment, environment=environment
+        )
+
+    # When
+    result = exceeds_segment_override_limit(environment=environment)
+
+    # Then
+    assert result is True
+
+
+def test_segment_override_limit_v2_delete_count_uses_unique_segment_ids(
+    project: Project,
+    environment_v2_versioning: Environment,
+) -> None:
+    # Given
+    project.max_segment_overrides_allowed = 5
+    project.save()
+
+    features = [
+        Feature.objects.create(name=f"feature_{i}", project=project) for i in range(5)
+    ]
+    shared_segment = Segment.objects.create(name="shared_segment", project=project)
+    new_segment = Segment.objects.create(name="new_segment", project=project)
+
+    for feature in features:
+        version = EnvironmentFeatureVersion.objects.get_or_create(
+            environment=environment_v2_versioning,
+            feature=feature,
+        )[0]
+        FeatureSegment.objects.create(
+            feature=feature,
+            segment=shared_segment,
+            environment=environment_v2_versioning,
+            environment_feature_version=version,
+        )
+
+    # When
+    result = exceeds_segment_override_limit(
+        environment=environment_v2_versioning,
+        segment_ids_to_create_overrides=[new_segment.id],
+        segment_ids_to_delete_overrides=[shared_segment.id],
+    )
+
+    # Then
+    assert result is True
+
+
 def test_segment_override_limit_does_not_exclude_invalid_overrides_being_deleted(
     feature: Feature,
     segment: Segment,