feat: enable feature versioning for new environments (#5108)

Author: matthewelwell

api/environments/models.py

@@ -15,6 +15,7 @@
     AFTER_DELETE,
     AFTER_SAVE,
     AFTER_UPDATE,
+    BEFORE_CREATE,
     LifecycleModel,
     hook,
 )
@@ -49,6 +50,7 @@
 )
 from features.models import Feature, FeatureSegment, FeatureState
 from features.multivariate.models import MultivariateFeatureStateValue
+from integrations.flagsmith.client import get_client
 from metadata.models import Metadata
 from projects.models import Project
 from segments.models import Segment
@@ -195,6 +197,24 @@ def delete_environment_document_from_cache(self) -> None:
         ):
             environment_document_cache.delete(self.api_key)
 
+    @hook(BEFORE_CREATE)  # type: ignore[misc]
+    def enable_v2_versioning(self) -> None:
+        if self.use_v2_feature_versioning:
+            # if the environment has already been created with versioning enabled,
+            # we don't want to disable it based on the flag state.
+            return
+
+        flagsmith_client = get_client("local", local_eval=True)
+        organisation = self.project.organisation
+        enable_v2_versioning = flagsmith_client.get_identity_flags(
+            organisation.flagsmith_identifier,
+            traits={
+                "organisation_id": organisation.id,
+                "subscription.plan": organisation.subscription.plan,
+            },
+        ).is_feature_enabled("enable_feature_versioning_for_new_environments")
+        self.use_v2_feature_versioning = enable_v2_versioning
+
     def __str__(self):  # type: ignore[no-untyped-def]
         return "Project %s - Environment %s" % (self.project.name, self.name)
 

api/tests/unit/environments/conftest.py

@@ -1,3 +1,4 @@
+import typing
 from unittest.mock import Mock
 
 import pytest
@@ -12,3 +13,31 @@ def mock_dynamo_env_wrapper(mocker: MockerFixture) -> Mock:
 @pytest.fixture()
 def mock_dynamo_env_v2_wrapper(mocker: MockerFixture) -> Mock:
     return mocker.patch("environments.models.environment_v2_wrapper")
+
+
+@pytest.fixture()
+def enable_v2_versioning_for_new_environments(
+    mocker: MockerFixture,
+) -> typing.Callable[[], None]:
+    """
+    This fixture returns a callable that allows us to enable the flag such that new environments
+    are created with feature versioning v2 enabled. It returns a callable to ensure that we can
+    control when the flag is enabled in the test code, to ensure that other fixtures are unaffected.
+
+    Relevant issue for improving this: https://github.com/Flagsmith/flagsmith-python-client/issues/135
+    """
+
+    def _enable_v2_versioning_for_new_environments() -> None:
+        mock_flagsmith_client = mocker.MagicMock()
+        mocker.patch(
+            "environments.models.get_client", return_value=mock_flagsmith_client
+        )
+
+        mocked_identity_flags = mock_flagsmith_client.get_identity_flags.return_value
+
+        def is_feature_enabled(feature_name: str) -> bool:
+            return feature_name == "enable_feature_versioning_for_new_environments"
+
+        mocked_identity_flags.is_feature_enabled.side_effect = is_feature_enabled
+
+    return _enable_v2_versioning_for_new_environments

api/tests/unit/environments/test_unit_environments_models.py

@@ -1208,3 +1208,59 @@ def test_environment_metric_query_helpers_match_expected_counts(
     assert change_request_count_result == change_request_count
     assert scheduled_count_result == scheduled_change_count
     assert identity_override_count == 0
+
+
+def test_environment_create_with_use_v2_feature_versioning_true(
+    project: Project,
+    feature: Feature,
+    enable_v2_versioning_for_new_environments: typing.Callable[[], None],
+) -> None:
+    # Given
+    enable_v2_versioning_for_new_environments()
+
+    # When
+    new_environment = Environment.objects.create(
+        name="new-environment",
+        project=project,
+    )
+
+    # Then
+    assert EnvironmentFeatureVersion.objects.filter(
+        environment=new_environment, feature=feature
+    ).exists()
+
+
+def test_environment_clone_from_versioned_environment_with_use_v2_feature_versioning_true(
+    project: Project,
+    environment_v2_versioning: Environment,
+    feature: Feature,
+    enable_v2_versioning_for_new_environments: typing.Callable[[], None],
+) -> None:
+    # Given
+    enable_v2_versioning_for_new_environments()
+
+    # When
+    new_environment = environment_v2_versioning.clone(name="new-environment")
+
+    # Then
+    assert EnvironmentFeatureVersion.objects.filter(
+        environment=new_environment, feature=feature
+    ).exists()
+
+
+def test_environment_clone_from_non_versioned_environment_with_use_v2_feature_versioning_true(
+    project: Project,
+    environment: Environment,
+    feature: Feature,
+    enable_v2_versioning_for_new_environments: typing.Callable[[], None],
+) -> None:
+    # Given
+    enable_v2_versioning_for_new_environments()
+
+    # When
+    new_environment = environment.clone(name="new-environment")
+
+    # Then
+    assert not EnvironmentFeatureVersion.objects.filter(
+        environment=new_environment, feature=feature
+    ).exists()