From 55a9ef7bcfc62b9752a7407fbfe545b2df6bb72c Mon Sep 17 00:00:00 2001
From: Gagan Trivedi <gagandeeptrivedi47@gmail.com>
Date: Tue, 5 Sep 2023 14:42:10 +0530
Subject: [PATCH] fix(model/featurestate): make environment not null (#2708)

---
 ...0_set_featurestate_environment_not_null.py | 36 +++++++++++
 api/features/models.py                        |  1 -
 api/tests/unit/features/test_migrations.py    | 60 +++++++++++++++++++
 3 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 api/features/migrations/0060_set_featurestate_environment_not_null.py

diff --git a/api/features/migrations/0060_set_featurestate_environment_not_null.py b/api/features/migrations/0060_set_featurestate_environment_not_null.py
new file mode 100644
index 000000000000..7096489de69c
--- /dev/null
+++ b/api/features/migrations/0060_set_featurestate_environment_not_null.py
@@ -0,0 +1,36 @@
+# Generated by Django 3.2.20 on 2023-08-30 13:14
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+def delete_feature_states_without_environment(apps, schema_editor):
+    FeatureState = apps.get_model("features", "FeatureState")
+    FeatureState.objects.filter(environment__isnull=True).delete()
+
+
+class Migration(migrations.Migration):
+    atomic = False
+    dependencies = [
+        (
+            "environments",
+            "0032_rename_use_mv_v2_evaluation_to_use_in_percentage_split_evaluation",
+        ),
+        ("features", "0059_fix_feature_type"),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            delete_feature_states_without_environment,
+            reverse_code=migrations.RunPython.noop,
+        ),
+        migrations.AlterField(
+            model_name="featurestate",
+            name="environment",
+            field=models.ForeignKey(
+                on_delete=django.db.models.deletion.CASCADE,
+                related_name="feature_states",
+                to="environments.environment",
+            ),
+        ),
+    ]
diff --git a/api/features/models.py b/api/features/models.py
index 3b8c0a350592..3b174286e818 100644
--- a/api/features/models.py
+++ b/api/features/models.py
@@ -400,7 +400,6 @@ class FeatureState(
     environment = models.ForeignKey(
         "environments.Environment",
         related_name="feature_states",
-        null=True,
         on_delete=models.CASCADE,
     )
     identity = models.ForeignKey(
diff --git a/api/tests/unit/features/test_migrations.py b/api/tests/unit/features/test_migrations.py
index 6a6e7a14d47f..92ea1954f443 100644
--- a/api/tests/unit/features/test_migrations.py
+++ b/api/tests/unit/features/test_migrations.py
@@ -2,6 +2,7 @@
 
 import pytest
 from django.conf import settings
+from django.db.utils import IntegrityError
 from django.utils import timezone
 
 from features.feature_types import MULTIVARIATE, STANDARD
@@ -241,3 +242,62 @@ def test_fix_feature_type_migration(migrator):
     )
     assert NewFeature.objects.get(id=standard_feature.id).type == STANDARD
     assert NewFeature.objects.get(id=mv_feature.id).type == MULTIVARIATE
+
+
+@pytest.mark.skipif(
+    settings.SKIP_MIGRATION_TESTS is True,
+    reason="Skip migration tests to speed up tests where necessary",
+)
+def test_migrate_set_featurestate_environment_not_null(migrator):
+    # Given
+    old_state = migrator.apply_initial_migration(("features", "0059_fix_feature_type"))
+
+    Organisation = old_state.apps.get_model("organisations", "Organisation")
+    Project = old_state.apps.get_model("projects", "Project")
+    Environment = old_state.apps.get_model("environments", "Environment")
+    Feature = old_state.apps.get_model("features", "Feature")
+    FeatureState = old_state.apps.get_model("features", "FeatureState")
+
+    organisation = Organisation.objects.create(name="test org")
+    project = Project.objects.create(
+        name="test project", organisation_id=organisation.id
+    )
+    environment = Environment.objects.create(
+        name="test environment", project_id=project.id
+    )
+    feature = Feature.objects.create(name="test_feature", project_id=project.id)
+
+    # mimick the creation of the feature states that would have happened when save is called on the model (but doesn't
+    # happen because we're using the migrator models)
+    FeatureState.objects.create(feature=feature, environment=environment, enabled=True)
+
+    # Let's create a feature state with a null environment
+    fs_without_environment = FeatureState.objects.create(
+        feature_id=feature.id,
+        environment_id=None,
+        enabled=True,
+    )
+
+    # When
+    new_state = migrator.apply_tested_migration(
+        ("features", "0060_set_featurestate_environment_not_null")
+    )
+
+    # Then
+    NewFeatureState = new_state.apps.get_model("features", "FeatureState")
+
+    # The feature state should have been deleted
+    assert not NewFeatureState.objects.filter(id=fs_without_environment.id).exists()
+
+    # feature state with environment still exists
+    assert NewFeatureState.objects.filter(
+        feature_id=feature.id, environment_id=environment.id
+    ).exists()
+
+    # and we cant create a feature state without an environment
+    with pytest.raises(IntegrityError):
+        NewFeatureState.objects.create(
+            feature_id=feature.id,
+            environment_id=None,
+            enabled=True,
+        )