From 87cc6c2f20e86ffdfe2d99055d2c300e63b7b17b Mon Sep 17 00:00:00 2001
From: Kim Gustyr <kim.gustyr@flagsmith.com>
Date: Fri, 7 Jun 2024 19:29:08 +0100
Subject: [PATCH] fix: V2 environments are not rebuilt on feature version
 publish

---
 api/environments/tasks.py                         | 12 +++---------
 .../environments/test_unit_environments_tasks.py  | 15 ++++++++++-----
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/api/environments/tasks.py b/api/environments/tasks.py
index 286492755b1a..bfc0450f250e 100644
--- a/api/environments/tasks.py
+++ b/api/environments/tasks.py
@@ -1,8 +1,5 @@
 from audit.models import AuditLog
-from environments.dynamodb import (
-    DynamoEnvironmentWrapper,
-    DynamoIdentityWrapper,
-)
+from environments.dynamodb import DynamoIdentityWrapper
 from environments.models import (
     Environment,
     environment_v2_wrapper,
@@ -17,11 +14,8 @@
 
 
 @register_task_handler(priority=TaskPriority.HIGH)
-def rebuild_environment_document(environment_id: int):
-    wrapper = DynamoEnvironmentWrapper()
-    if wrapper.is_enabled:
-        environment = Environment.objects.get(id=environment_id)
-        wrapper.write_environment(environment)
+def rebuild_environment_document(environment_id: int) -> None:
+    Environment.write_environments_to_dynamodb(environment_id=environment_id)
 
 
 @register_task_handler(priority=TaskPriority.HIGHEST)
diff --git a/api/tests/unit/environments/test_unit_environments_tasks.py b/api/tests/unit/environments/test_unit_environments_tasks.py
index 0669c528176f..d83408383340 100644
--- a/api/tests/unit/environments/test_unit_environments_tasks.py
+++ b/api/tests/unit/environments/test_unit_environments_tasks.py
@@ -1,6 +1,7 @@
 from pytest_mock import MockerFixture
 
 from audit.models import AuditLog
+from environments.models import Environment
 from environments.tasks import (
     delete_environment_from_dynamo,
     process_environment_update,
@@ -8,18 +9,22 @@
 )
 
 
-def test_rebuild_environment_document(environment, mocker):
+def test_rebuild_environment_document(
+    environment: Environment,
+    mocker: MockerFixture,
+) -> None:
     # Given
-    mock_dynamo_wrapper = mocker.MagicMock(is_enabled=True)
-    mocker.patch(
-        "environments.tasks.DynamoEnvironmentWrapper", return_value=mock_dynamo_wrapper
+    mock_write_environments_to_dynamodb = mocker.patch(
+        "environments.tasks.Environment.write_environments_to_dynamodb",
     )
 
     # When
     rebuild_environment_document(environment_id=environment.id)
 
     # Then
-    mock_dynamo_wrapper.write_environment.assert_called_once_with(environment)
+    mock_write_environments_to_dynamodb.assert_called_once_with(
+        environment_id=environment.id
+    )
 
 
 def test_process_environment_update_with_environment_audit_log(environment, mocker):