fix: V2 environments are not rebuilt on feature version publish

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)

api/tests/unit/environments/test_unit_environments_tasks.py

@@ -1,25 +1,30 @@
 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,
     rebuild_environment_document,
 )
 
 
-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):