"""Simplified unit tests for DraftVarLoader focusing on core functionality.""" import json from unittest.mock import Mock, patch import pytest from sqlalchemy import Engine from core.workflow.file_reference import build_file_reference from graphon.file import File, FileTransferMethod, FileType from graphon.variables.segments import ObjectSegment, StringSegment from graphon.variables.types import SegmentType from models.model import UploadFile from models.workflow import WorkflowDraftVariable, WorkflowDraftVariableFile from services.workflow_draft_variable_service import DraftVarLoader class TestDraftVarLoaderSimple: """Simplified unit tests for DraftVarLoader core methods.""" @pytest.fixture def mock_engine(self) -> Engine: return Mock(spec=Engine) @pytest.fixture def draft_var_loader(self, mock_engine): """Create DraftVarLoader instance for testing.""" return DraftVarLoader( engine=mock_engine, app_id="test-app-id", tenant_id="test-tenant-id", user_id="test-user-id", fallback_variables=[], ) def test_load_offloaded_variable_object_type_unit(self, draft_var_loader): """Test _load_offloaded_variable with object type - isolated unit test.""" # Create mock objects upload_file = Mock(spec=UploadFile) upload_file.key = "storage/key/test.json" variable_file = Mock(spec=WorkflowDraftVariableFile) variable_file.value_type = SegmentType.OBJECT variable_file.upload_file = upload_file draft_var = Mock(spec=WorkflowDraftVariable) draft_var.id = "draft-var-id" draft_var.node_id = "test-node-id" draft_var.name = "test_object" draft_var.description = "test description" draft_var.get_selector.return_value = ["test-node-id", "test_object"] draft_var.variable_file = variable_file test_object = {"key1": "value1", "key2": 42} test_json_content = json.dumps(test_object, ensure_ascii=False, separators=(",", ":")) with patch("services.workflow_draft_variable_service.storage") as mock_storage: mock_storage.load.return_value = test_json_content.encode() mock_segment = ObjectSegment(value=test_object) draft_var.build_segment_from_serialized_value.return_value = mock_segment # Execute the method selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) # Verify results assert selector_tuple == ("test-node-id", "test_object") assert variable.id == "draft-var-id" assert variable.name == "test_object" assert variable.description == "test description" assert variable.value == test_object # Verify method calls mock_storage.load.assert_called_once_with("storage/key/test.json") draft_var.build_segment_from_serialized_value.assert_called_once_with(SegmentType.OBJECT, test_object) def test_load_offloaded_variable_missing_variable_file_unit(self, draft_var_loader): """Test that assertion error is raised when variable_file is None.""" draft_var = Mock(spec=WorkflowDraftVariable) draft_var.variable_file = None with pytest.raises(AssertionError): draft_var_loader._load_offloaded_variable(draft_var) def test_load_offloaded_variable_missing_upload_file_unit(self, draft_var_loader): """Test that assertion error is raised when upload_file is None.""" variable_file = Mock(spec=WorkflowDraftVariableFile) variable_file.upload_file = None draft_var = Mock(spec=WorkflowDraftVariable) draft_var.variable_file = variable_file with pytest.raises(AssertionError): draft_var_loader._load_offloaded_variable(draft_var) def test_load_variables_empty_selectors_unit(self, draft_var_loader): """Test load_variables returns empty list for empty selectors.""" result = draft_var_loader.load_variables([]) assert result == [] def test_selector_to_tuple_unit(self, draft_var_loader): """Test _selector_to_tuple method.""" selector = ["node_id", "var_name", "extra_field"] result = draft_var_loader._selector_to_tuple(selector) assert result == ("node_id", "var_name") def test_load_offloaded_variable_array_type_unit(self, draft_var_loader): """Test _load_offloaded_variable with array type - isolated unit test.""" # Create mock objects upload_file = Mock(spec=UploadFile) upload_file.key = "storage/key/test_array.json" variable_file = Mock(spec=WorkflowDraftVariableFile) variable_file.value_type = SegmentType.ARRAY_ANY variable_file.upload_file = upload_file draft_var = Mock(spec=WorkflowDraftVariable) draft_var.id = "draft-var-id" draft_var.node_id = "test-node-id" draft_var.name = "test_array" draft_var.description = "test array description" draft_var.get_selector.return_value = ["test-node-id", "test_array"] draft_var.variable_file = variable_file test_array = ["item1", "item2", "item3"] test_json_content = json.dumps(test_array) with patch("services.workflow_draft_variable_service.storage") as mock_storage: mock_storage.load.return_value = test_json_content.encode() from graphon.variables.segments import ArrayAnySegment mock_segment = ArrayAnySegment(value=test_array) draft_var.build_segment_from_serialized_value.return_value = mock_segment # Execute the method selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) # Verify results assert selector_tuple == ("test-node-id", "test_array") assert variable.id == "draft-var-id" assert variable.name == "test_array" assert variable.description == "test array description" # Verify method calls mock_storage.load.assert_called_once_with("storage/key/test_array.json") draft_var.build_segment_from_serialized_value.assert_called_once_with(SegmentType.ARRAY_ANY, test_array) def test_load_offloaded_variable_file_type_rebuilds_storage_backed_payload(self, draft_var_loader): upload_file = Mock(spec=UploadFile) upload_file.key = "storage/key/test_file.json" variable_file = Mock(spec=WorkflowDraftVariableFile) variable_file.value_type = SegmentType.FILE variable_file.upload_file = upload_file draft_var = WorkflowDraftVariable( id="draft-var-id", app_id="app-1", node_id="test-node-id", name="test_file", description="test file description", ) draft_var._set_selector(["test-node-id", "test_file"]) draft_var.variable_file = variable_file persisted_file = File( file_id="file-1", file_type=FileType.DOCUMENT, transfer_method=FileTransferMethod.LOCAL_FILE, reference=build_file_reference(record_id="upload-1", storage_key="legacy-storage-key"), filename="test.txt", extension=".txt", mime_type="text/plain", size=12, ) rebuilt_file = File( file_id="file-1", file_type=FileType.DOCUMENT, transfer_method=FileTransferMethod.LOCAL_FILE, reference=build_file_reference(record_id="upload-1"), filename="test.txt", extension=".txt", mime_type="text/plain", size=12, storage_key="canonical-storage-key", ) raw_file = { **persisted_file.model_dump(mode="json"), "tenant_id": "legacy-tenant", } with ( patch("services.workflow_draft_variable_service.storage") as mock_storage, patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"), patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file, ): mock_storage.load.return_value = json.dumps(raw_file).encode() selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) assert selector_tuple == ("test-node-id", "test_file") assert variable.id == "draft-var-id" assert variable.name == "test_file" assert variable.description == "test file description" assert variable.value == rebuilt_file rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1") def test_load_variables_with_offloaded_variables_unit(self, draft_var_loader): """Test load_variables method with mix of regular and offloaded variables.""" selectors = [["node1", "regular_var"], ["node2", "offloaded_var"]] # Mock regular variable regular_draft_var = Mock(spec=WorkflowDraftVariable) regular_draft_var.is_truncated.return_value = False regular_draft_var.node_id = "node1" regular_draft_var.name = "regular_var" regular_draft_var.get_value.return_value = StringSegment(value="regular_value") regular_draft_var.get_selector.return_value = ["node1", "regular_var"] regular_draft_var.id = "regular-var-id" regular_draft_var.description = "regular description" # Mock offloaded variable upload_file = Mock(spec=UploadFile) upload_file.key = "storage/key/offloaded.txt" variable_file = Mock(spec=WorkflowDraftVariableFile) variable_file.value_type = SegmentType.STRING variable_file.upload_file = upload_file offloaded_draft_var = Mock(spec=WorkflowDraftVariable) offloaded_draft_var.is_truncated.return_value = True offloaded_draft_var.node_id = "node2" offloaded_draft_var.name = "offloaded_var" offloaded_draft_var.get_selector.return_value = ["node2", "offloaded_var"] offloaded_draft_var.variable_file = variable_file offloaded_draft_var.id = "offloaded-var-id" offloaded_draft_var.description = "offloaded description" draft_vars = [regular_draft_var, offloaded_draft_var] with patch("services.workflow_draft_variable_service.Session") as mock_session_cls: mock_session = Mock() mock_session_cls.return_value.__enter__.return_value = mock_session mock_service = Mock() mock_service.get_draft_variables_by_selectors.return_value = draft_vars with patch( "services.workflow_draft_variable_service.WorkflowDraftVariableService", return_value=mock_service ): with patch("services.workflow_draft_variable_service.StorageKeyLoader"): with patch("factories.variable_factory.segment_to_variable") as mock_segment_to_variable: # Mock regular variable creation regular_variable = Mock() regular_variable.selector = ["node1", "regular_var"] # Mock offloaded variable creation offloaded_variable = Mock() offloaded_variable.selector = ["node2", "offloaded_var"] mock_segment_to_variable.return_value = regular_variable with patch("services.workflow_draft_variable_service.storage") as mock_storage: mock_storage.load.return_value = b"offloaded_content" with patch.object(draft_var_loader, "_load_offloaded_variable") as mock_load_offloaded: mock_load_offloaded.return_value = (("node2", "offloaded_var"), offloaded_variable) with patch("concurrent.futures.ThreadPoolExecutor") as mock_executor_cls: mock_executor = Mock() mock_executor_cls.return_value.__enter__.return_value = mock_executor mock_executor.map.return_value = [(("node2", "offloaded_var"), offloaded_variable)] # Execute the method result = draft_var_loader.load_variables(selectors) # Verify results assert len(result) == 2 # Verify service method was called mock_service.get_draft_variables_by_selectors.assert_called_once_with( draft_var_loader._app_id, selectors, user_id=draft_var_loader._user_id, ) # Verify offloaded variable loading was called mock_load_offloaded.assert_called_once_with(offloaded_draft_var) def test_load_variables_all_offloaded_variables_unit(self, draft_var_loader): """Test load_variables method with only offloaded variables.""" selectors = [["node1", "offloaded_var1"], ["node2", "offloaded_var2"]] # Mock first offloaded variable offloaded_var1 = Mock(spec=WorkflowDraftVariable) offloaded_var1.is_truncated.return_value = True offloaded_var1.node_id = "node1" offloaded_var1.name = "offloaded_var1" # Mock second offloaded variable offloaded_var2 = Mock(spec=WorkflowDraftVariable) offloaded_var2.is_truncated.return_value = True offloaded_var2.node_id = "node2" offloaded_var2.name = "offloaded_var2" draft_vars = [offloaded_var1, offloaded_var2] with patch("services.workflow_draft_variable_service.Session") as mock_session_cls: mock_session = Mock() mock_session_cls.return_value.__enter__.return_value = mock_session mock_service = Mock() mock_service.get_draft_variables_by_selectors.return_value = draft_vars with patch( "services.workflow_draft_variable_service.WorkflowDraftVariableService", return_value=mock_service ): with patch("services.workflow_draft_variable_service.StorageKeyLoader"): with patch("services.workflow_draft_variable_service.ThreadPoolExecutor") as mock_executor_cls: mock_executor = Mock() mock_executor_cls.return_value.__enter__.return_value = mock_executor mock_executor.map.return_value = [ (("node1", "offloaded_var1"), Mock()), (("node2", "offloaded_var2"), Mock()), ] # Execute the method result = draft_var_loader.load_variables(selectors) # Verify results - since we have only offloaded variables, should have 2 results assert len(result) == 2 # Verify ThreadPoolExecutor was used mock_executor_cls.assert_called_once_with(max_workers=10) mock_executor.map.assert_called_once()