From 8cd27318bd9de08cd45785868e7d77541c9e3a61 Mon Sep 17 00:00:00 2001
From: Matthew Elwell <matthew.elwell@flagsmith.com>
Date: Fri, 10 Jan 2025 12:20:46 +0000
Subject: [PATCH] fix(export): fixes issue exporting identity feature states
 with missing attributes (#4987)

---
 api/edge_api/identities/export.py             | 16 ++++++-----
 .../test_unit_import_export_export.py         | 27 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/api/edge_api/identities/export.py b/api/edge_api/identities/export.py
index ebfebb85f498..5db66340949e 100644
--- a/api/edge_api/identities/export.py
+++ b/api/edge_api/identities/export.py
@@ -65,13 +65,15 @@ def export_edge_identity_and_overrides(  # noqa: C901
                         override["enabled"],
                     )
                 )
-                featurestate_value = override["feature_state_value"]
-                if featurestate_value is not None:
-                    # export feature state value
-                    identity_override_export.append(
-                        export_featurestate_value(featurestate_value, featurestate_uuid)
-                    )
-                if mvfsv_overrides := override["multivariate_feature_state_values"]:
+
+                # We always want to create the FeatureStateValue, but if there is none in the
+                # dynamo object, we just create a default object with a value of null.
+                featurestate_value = override.get("feature_state_value")
+                identity_override_export.append(
+                    export_featurestate_value(featurestate_value, featurestate_uuid)
+                )
+
+                if mvfsv_overrides := override.get("multivariate_feature_state_values"):
                     for mvfsv_override in mvfsv_overrides:
                         mv_feature_option_id = mvfsv_override[
                             "multivariate_feature_option"
diff --git a/api/tests/unit/import_export/test_unit_import_export_export.py b/api/tests/unit/import_export/test_unit_import_export_export.py
index 8cb374d28666..ab5ec47139c2 100644
--- a/api/tests/unit/import_export/test_unit_import_export_export.py
+++ b/api/tests/unit/import_export/test_unit_import_export_export.py
@@ -346,6 +346,11 @@ def test_export_edge_identities(
         project=project, name="mv_feature_with_deleted_option", type=MULTIVARIATE
     )
 
+    # and a feature that we will override with missing attributes
+    feature_to_override_with_missing_attrs = Feature.objects.create(
+        project=project, name="feature_to_override_with_missing_attrs"
+    )
+
     mv_option = multivariate_options[0]
 
     identity_identifier = "Development_user_123456"
@@ -458,6 +463,19 @@ def test_export_edge_identities(
                     }
                 ],
             },
+            # An override with missing feature_state_value and multivariate_feature_state_values
+            # that should be handled correctly
+            {
+                "django_id": None,
+                "enabled": False,
+                "feature": {
+                    "id": feature_to_override_with_missing_attrs.id,
+                    "name": feature_to_override_with_missing_attrs.name,
+                    "type": "STANDARD",
+                },
+                "featurestate_uuid": "53ab4fc5-0027-47a4-85b3-825abe2287eb",
+                "feature_segment": None,
+            },
         ],
         "identity_traits": [
             {"trait_key": "int_trait", "trait_value": 123},
@@ -517,7 +535,7 @@ def test_export_edge_identities(
     assert bool_trait.trait_value is True
 
     all_feature_states = identity.get_all_feature_states()
-    assert len(all_feature_states) == 6
+    assert len(all_feature_states) == 7
 
     actual_mv_override = all_feature_states[0]
     assert str(actual_mv_override.uuid) == mv_override_fs_uuid
@@ -549,6 +567,13 @@ def test_export_edge_identities(
     )
     assert override_without_mv_option.identity == identity
 
+    override_with_missing_attributes = all_feature_states[6]
+    assert override_with_missing_attributes.feature_state_value.value is None
+    assert (
+        override_with_missing_attributes.multivariate_feature_state_values.exists()
+        is False
+    )
+
 
 @mock_s3
 def test_organisation_exporter_export_to_s3(organisation):