Refactor mocking of v4.models in test settings and plan_to_mplan_converter to use real classes for improved type handling and consistency

Author: Kingshuk-Microsoft

src/tests/backend/v4/config/test_settings.py

@@ -42,16 +42,27 @@
 sys.modules['azure.keyvault.secrets'] = Mock()
 sys.modules['azure.keyvault.secrets.aio'] = Mock()
 
-# Mock v4.models for relative imports used in settings.py
-mock_v4_models_messages = Mock()
-mock_mplan = Mock()
-mock_websocket_message_type = Mock()
-mock_websocket_message_type.SYSTEM_MESSAGE = 'system_message'
-mock_v4_models_messages.MPlan = mock_mplan
-mock_v4_models_messages.WebsocketMessageType = mock_websocket_message_type
-sys.modules['v4'] = Mock()
-sys.modules['v4.models'] = Mock()
+# Import the real v4.models classes first to avoid type annotation issues
+from backend.v4.models.messages import MPlan, WebsocketMessageType
+from backend.v4.models.models import MPlan as MPlanModel, MStep
+
+# Mock v4.models for relative imports used in settings.py, using REAL classes
+from types import ModuleType
+mock_v4 = ModuleType('v4')
+mock_v4_models = ModuleType('v4.models')
+mock_v4_models_messages = ModuleType('v4.models.messages')
+mock_v4_models_models = ModuleType('v4.models.models')
+
+# Assign real classes to mock modules
+mock_v4_models_messages.MPlan = MPlan
+mock_v4_models_messages.WebsocketMessageType = WebsocketMessageType
+mock_v4_models_models.MPlan = MPlanModel
+mock_v4_models_models.MStep = MStep
+
+sys.modules['v4'] = mock_v4
+sys.modules['v4.models'] = mock_v4_models
 sys.modules['v4.models.messages'] = mock_v4_models_messages
+sys.modules['v4.models.models'] = mock_v4_models_models
 
 # Mock common.config.app_config 
 sys.modules['common'] = Mock()

src/tests/backend/v4/orchestration/helper/test_plan_to_mplan_converter.py

@@ -21,15 +21,23 @@
 # Import the models first (from backend path)
 from backend.v4.models.models import MPlan, MStep, PlanStatus
 
-# Mock v4.models.models with the real classes so relative imports work
-from types import ModuleType
-mock_v4_models_models = ModuleType('models')
-mock_v4_models_models.MPlan = MPlan
-mock_v4_models_models.MStep = MStep
-mock_v4_models_models.PlanStatus = PlanStatus
-sys.modules['v4'] = ModuleType('v4')
-sys.modules['v4.models'] = ModuleType('models')
-sys.modules['v4.models.models'] = mock_v4_models_models
+# Check if v4.models.models is already properly set up (running in full test suite)
+_existing_v4_models = sys.modules.get('v4.models.models')
+_need_mock = _existing_v4_models is None or not hasattr(_existing_v4_models, 'MPlan')
+
+if _need_mock:
+    # Mock v4.models.models with the real classes so relative imports work
+    from types import ModuleType
+    mock_v4_models_models = ModuleType('models')
+    mock_v4_models_models.MPlan = MPlan
+    mock_v4_models_models.MStep = MStep
+    mock_v4_models_models.PlanStatus = PlanStatus
+    
+    if 'v4' not in sys.modules:
+        sys.modules['v4'] = ModuleType('v4')
+    if 'v4.models' not in sys.modules:
+        sys.modules['v4.models'] = ModuleType('models')
+    sys.modules['v4.models.models'] = mock_v4_models_models
 
 # Now import the converter
 from backend.v4.orchestration.helper.plan_to_mplan_converter import PlanToMPlanConverter