From 844f13a7934c2a483d9471abad26c4466b4e035e Mon Sep 17 00:00:00 2001 From: NiveditJain Date: Tue, 16 Sep 2025 16:22:08 +0530 Subject: [PATCH 1/4] feat: add manual retry state functionality - Introduced a new API endpoint for manual state retries, allowing users to trigger retries for specific states. - Implemented request and response models for manual retry operations. - Added error handling for duplicate retry states and invalid API keys. - Enhanced logging for better traceability of retry operations. This feature improves the system's ability to manage state retries effectively. --- .../app/controller/manul_retry_state.py | 49 +++++++++++++++++++ state-manager/app/models/manual_retry.py | 11 +++++ state-manager/app/routes.py | 22 +++++++++ 3 files changed, 82 insertions(+) create mode 100644 state-manager/app/controller/manul_retry_state.py create mode 100644 state-manager/app/models/manual_retry.py diff --git a/state-manager/app/controller/manul_retry_state.py b/state-manager/app/controller/manul_retry_state.py new file mode 100644 index 00000000..fc97bec9 --- /dev/null +++ b/state-manager/app/controller/manul_retry_state.py @@ -0,0 +1,49 @@ +from pymongo.errors import DuplicateKeyError +from app.models.manual_retry import ManualRetryRequestModel, ManualRetryResponseModel +from beanie import PydanticObjectId +from app.singletons.logs_manager import LogsManager +from app.models.state_status_enum import StateStatusEnum +from fastapi import HTTPException, status +from app.models.db.state import State + + +logger = LogsManager().get_logger() + +async def manual_retry_state(namespace_name: str, state_id: PydanticObjectId, body: ManualRetryRequestModel, x_exosphere_request_id: str): + try: + logger.info(f"Manual retry state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) + + state = await State.find_one(State.id == state_id) + if not state: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="State not found") + + try: + retry_state = State( + node_name=state.node_name, + namespace_name=state.namespace_name, + identifier=state.identifier, + graph_name=state.graph_name, + run_id=state.run_id, + status=StateStatusEnum.CREATED, + inputs=state.inputs, + outputs={}, + error=None, + parents=state.parents, + does_unites=state.does_unites, + fanout_id=body.fanout_id # this will ensure that multiple unwanted retries are not formed because of index in database + ) + retry_state = await retry_state.insert() + logger.info(f"Retry state {retry_state.id} created for state {state_id}", x_exosphere_request_id=x_exosphere_request_id) + + state.status = StateStatusEnum.RETRY_CREATED + await state.save() + + return ManualRetryResponseModel(id=str(retry_state.id), status=retry_state.status) + except DuplicateKeyError: + logger.info(f"Duplicate retry state detected for state {state_id}. A retry state with the same unique key already exists.", x_exosphere_request_id=x_exosphere_request_id) + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Duplicate retry state detected") + + + except Exception as e: + logger.error(f"Error manual retry state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) + raise e diff --git a/state-manager/app/models/manual_retry.py b/state-manager/app/models/manual_retry.py new file mode 100644 index 00000000..0aec686b --- /dev/null +++ b/state-manager/app/models/manual_retry.py @@ -0,0 +1,11 @@ +from pydantic import BaseModel, Field +from .state_status_enum import StateStatusEnum + + +class ManualRetryRequestModel(BaseModel): + fanout_id: str = Field(..., description="Fanout ID of the state") + + +class ManualRetryResponseModel(BaseModel): + id: str = Field(..., description="ID of the state") + status: StateStatusEnum = Field(..., description="Status of the state") diff --git a/state-manager/app/routes.py b/state-manager/app/routes.py index 53f73ff5..d202e65b 100644 --- a/state-manager/app/routes.py +++ b/state-manager/app/routes.py @@ -50,6 +50,10 @@ from .models.signal_models import ReEnqueueAfterRequestModel from .controller.re_queue_after_signal import re_queue_after_signal +# manual_retry +from .models.manual_retry import ManualRetryRequestModel, ManualRetryResponseModel +from .controller.manul_retry_state import manual_retry_state + logger = LogsManager().get_logger() @@ -176,6 +180,24 @@ async def re_enqueue_after_state_route(namespace_name: str, state_id: str, body: return await re_queue_after_signal(namespace_name, PydanticObjectId(state_id), body, x_exosphere_request_id) +@router.post( + "state/{state_id}/manual-retry", + response_model=ManualRetryResponseModel, + status_code=status.HTTP_200_OK, + response_description="State manual retry successfully", + tags=["state"] +) +async def manual_retry_state_route(namespace_name: str, state_id: str, body: ManualRetryRequestModel, request: Request, api_key: str = Depends(check_api_key)): + x_exosphere_request_id = getattr(request.state, "x_exosphere_request_id", str(uuid4())) + + if api_key: + logger.info(f"API key is valid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) + else: + logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key") + + return await manual_retry_state(namespace_name, PydanticObjectId(state_id), body, x_exosphere_request_id) + @router.put( "/graph/{graph_name}", From 3d68362622f1f84c8bc165a996eed137158bce72 Mon Sep 17 00:00:00 2001 From: NiveditJain Date: Tue, 16 Sep 2025 16:56:45 +0530 Subject: [PATCH 2/4] fix: correct import path for manual retry state and update route definition - Fixed the import path for the manual retry state controller in routes.py. - Updated the route definition for manual retry to include a leading slash for consistency. - Added a new controller file for manual retry state functionality, implementing the logic for handling manual retries. - Introduced unit tests for the manual retry request and response models, ensuring validation and functionality. - Enhanced unit tests for the manual retry state route, covering various scenarios including valid and invalid API keys. These changes improve the structure and reliability of the manual retry feature. --- ...l_retry_state.py => manual_retry_state.py} | 0 state-manager/app/routes.py | 4 +- .../tests/unit/models/test_manual_retry.py | 241 ++++++++++++++++++ state-manager/tests/unit/test_routes.py | 88 ++++++- 4 files changed, 329 insertions(+), 4 deletions(-) rename state-manager/app/controller/{manul_retry_state.py => manual_retry_state.py} (100%) create mode 100644 state-manager/tests/unit/models/test_manual_retry.py diff --git a/state-manager/app/controller/manul_retry_state.py b/state-manager/app/controller/manual_retry_state.py similarity index 100% rename from state-manager/app/controller/manul_retry_state.py rename to state-manager/app/controller/manual_retry_state.py diff --git a/state-manager/app/routes.py b/state-manager/app/routes.py index d202e65b..e552081f 100644 --- a/state-manager/app/routes.py +++ b/state-manager/app/routes.py @@ -52,7 +52,7 @@ # manual_retry from .models.manual_retry import ManualRetryRequestModel, ManualRetryResponseModel -from .controller.manul_retry_state import manual_retry_state +from .controller.manual_retry_state import manual_retry_state logger = LogsManager().get_logger() @@ -181,7 +181,7 @@ async def re_enqueue_after_state_route(namespace_name: str, state_id: str, body: return await re_queue_after_signal(namespace_name, PydanticObjectId(state_id), body, x_exosphere_request_id) @router.post( - "state/{state_id}/manual-retry", + "/state/{state_id}/manual-retry", response_model=ManualRetryResponseModel, status_code=status.HTTP_200_OK, response_description="State manual retry successfully", diff --git a/state-manager/tests/unit/models/test_manual_retry.py b/state-manager/tests/unit/models/test_manual_retry.py new file mode 100644 index 00000000..5869702c --- /dev/null +++ b/state-manager/tests/unit/models/test_manual_retry.py @@ -0,0 +1,241 @@ +import pytest +from pydantic import ValidationError + +from app.models.manual_retry import ManualRetryRequestModel, ManualRetryResponseModel +from app.models.state_status_enum import StateStatusEnum + + +class TestManualRetryRequestModel: + """Test cases for ManualRetryRequestModel""" + + def test_manual_retry_request_model_valid_data(self): + """Test ManualRetryRequestModel with valid fanout_id""" + # Arrange & Act + fanout_id = "test-fanout-id-123" + model = ManualRetryRequestModel(fanout_id=fanout_id) + + # Assert + assert model.fanout_id == fanout_id + + def test_manual_retry_request_model_empty_fanout_id(self): + """Test ManualRetryRequestModel with empty fanout_id""" + # Arrange & Act + fanout_id = "" + model = ManualRetryRequestModel(fanout_id=fanout_id) + + # Assert + assert model.fanout_id == fanout_id + + def test_manual_retry_request_model_uuid_fanout_id(self): + """Test ManualRetryRequestModel with UUID fanout_id""" + # Arrange & Act + fanout_id = "550e8400-e29b-41d4-a716-446655440000" + model = ManualRetryRequestModel(fanout_id=fanout_id) + + # Assert + assert model.fanout_id == fanout_id + + def test_manual_retry_request_model_long_fanout_id(self): + """Test ManualRetryRequestModel with long fanout_id""" + # Arrange & Act + fanout_id = "a" * 1000 # Very long string + model = ManualRetryRequestModel(fanout_id=fanout_id) + + # Assert + assert model.fanout_id == fanout_id + + def test_manual_retry_request_model_special_characters_fanout_id(self): + """Test ManualRetryRequestModel with special characters in fanout_id""" + # Arrange & Act + fanout_id = "test-fanout@#$%^&*()_+-={}[]|\\:;\"'<>?,./" + model = ManualRetryRequestModel(fanout_id=fanout_id) + + # Assert + assert model.fanout_id == fanout_id + + def test_manual_retry_request_model_missing_fanout_id(self): + """Test ManualRetryRequestModel with missing fanout_id field""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryRequestModel() # type: ignore + + assert "fanout_id" in str(exc_info.value) + assert "Field required" in str(exc_info.value) + + def test_manual_retry_request_model_none_fanout_id(self): + """Test ManualRetryRequestModel with None fanout_id""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryRequestModel(fanout_id=None) # type: ignore + + assert "fanout_id" in str(exc_info.value) + + def test_manual_retry_request_model_numeric_fanout_id(self): + """Test ManualRetryRequestModel with numeric fanout_id (should fail validation)""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryRequestModel(fanout_id=12345) # type: ignore + + assert "string_type" in str(exc_info.value) + + def test_manual_retry_request_model_dict_representation(self): + """Test ManualRetryRequestModel dict representation""" + # Arrange & Act + fanout_id = "test-fanout-id" + model = ManualRetryRequestModel(fanout_id=fanout_id) + + # Assert + expected_dict = {"fanout_id": fanout_id} + assert model.model_dump() == expected_dict + + def test_manual_retry_request_model_json_serialization(self): + """Test ManualRetryRequestModel JSON serialization""" + # Arrange & Act + fanout_id = "test-fanout-id" + model = ManualRetryRequestModel(fanout_id=fanout_id) + + # Assert + json_str = model.model_dump_json() + assert f'"fanout_id":"{fanout_id}"' in json_str + + +class TestManualRetryResponseModel: + """Test cases for ManualRetryResponseModel""" + + def test_manual_retry_response_model_valid_data(self): + """Test ManualRetryResponseModel with valid data""" + # Arrange & Act + state_id = "507f1f77bcf86cd799439011" + status = StateStatusEnum.CREATED + model = ManualRetryResponseModel(id=state_id, status=status) + + # Assert + assert model.id == state_id + assert model.status == status + + def test_manual_retry_response_model_all_status_types(self): + """Test ManualRetryResponseModel with all possible status values""" + # Arrange & Act & Assert + state_id = "507f1f77bcf86cd799439011" + + for status in StateStatusEnum: + model = ManualRetryResponseModel(id=state_id, status=status) + assert model.id == state_id + assert model.status == status + + def test_manual_retry_response_model_created_status(self): + """Test ManualRetryResponseModel with CREATED status""" + # Arrange & Act + state_id = "507f1f77bcf86cd799439011" + status = StateStatusEnum.CREATED + model = ManualRetryResponseModel(id=state_id, status=status) + + # Assert + assert model.id == state_id + assert model.status == StateStatusEnum.CREATED + + def test_manual_retry_response_model_retry_created_status(self): + """Test ManualRetryResponseModel with RETRY_CREATED status""" + # Arrange & Act + state_id = "507f1f77bcf86cd799439011" + status = StateStatusEnum.RETRY_CREATED + model = ManualRetryResponseModel(id=state_id, status=status) + + # Assert + assert model.id == state_id + assert model.status == StateStatusEnum.RETRY_CREATED + + def test_manual_retry_response_model_missing_id(self): + """Test ManualRetryResponseModel with missing id field""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryResponseModel(status=StateStatusEnum.CREATED) # type: ignore + + assert "id" in str(exc_info.value) + assert "Field required" in str(exc_info.value) + + def test_manual_retry_response_model_missing_status(self): + """Test ManualRetryResponseModel with missing status field""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryResponseModel(id="507f1f77bcf86cd799439011") # type: ignore + + assert "status" in str(exc_info.value) + assert "Field required" in str(exc_info.value) + + def test_manual_retry_response_model_none_id(self): + """Test ManualRetryResponseModel with None id""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryResponseModel(id=None, status=StateStatusEnum.CREATED) # type: ignore + + assert "id" in str(exc_info.value) + + def test_manual_retry_response_model_none_status(self): + """Test ManualRetryResponseModel with None status""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryResponseModel(id="507f1f77bcf86cd799439011", status=None) # type: ignore + + assert "status" in str(exc_info.value) + + def test_manual_retry_response_model_invalid_status(self): + """Test ManualRetryResponseModel with invalid status""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryResponseModel(id="507f1f77bcf86cd799439011", status="INVALID_STATUS") # type: ignore + + assert "status" in str(exc_info.value) + + def test_manual_retry_response_model_numeric_id(self): + """Test ManualRetryResponseModel with numeric id (should fail validation)""" + # Arrange & Act & Assert + with pytest.raises(ValidationError) as exc_info: + ManualRetryResponseModel(id=12345, status=StateStatusEnum.CREATED) # type: ignore + + assert "string_type" in str(exc_info.value) + + def test_manual_retry_response_model_dict_representation(self): + """Test ManualRetryResponseModel dict representation""" + # Arrange & Act + state_id = "507f1f77bcf86cd799439011" + status = StateStatusEnum.CREATED + model = ManualRetryResponseModel(id=state_id, status=status) + + # Assert + expected_dict = {"id": state_id, "status": status} + assert model.model_dump() == expected_dict + + def test_manual_retry_response_model_json_serialization(self): + """Test ManualRetryResponseModel JSON serialization""" + # Arrange & Act + state_id = "507f1f77bcf86cd799439011" + status = StateStatusEnum.CREATED + model = ManualRetryResponseModel(id=state_id, status=status) + + # Assert + json_str = model.model_dump_json() + assert f'"id":"{state_id}"' in json_str + assert f'"status":"{status.value}"' in json_str + + def test_manual_retry_response_model_empty_id(self): + """Test ManualRetryResponseModel with empty string id""" + # Arrange & Act + state_id = "" + status = StateStatusEnum.CREATED + model = ManualRetryResponseModel(id=state_id, status=status) + + # Assert + assert model.id == state_id + assert model.status == status + + def test_manual_retry_response_model_long_id(self): + """Test ManualRetryResponseModel with very long id""" + # Arrange & Act + state_id = "a" * 1000 # Very long string + status = StateStatusEnum.CREATED + model = ManualRetryResponseModel(id=state_id, status=status) + + # Assert + assert model.id == state_id + assert model.status == status \ No newline at end of file diff --git a/state-manager/tests/unit/test_routes.py b/state-manager/tests/unit/test_routes.py index af50f066..c3ca8539 100644 --- a/state-manager/tests/unit/test_routes.py +++ b/state-manager/tests/unit/test_routes.py @@ -8,6 +8,7 @@ from app.models.secrets_response import SecretsResponseModel from app.models.list_models import ListRegisteredNodesResponse, ListGraphTemplatesResponse from app.models.run_models import RunsResponse, RunListItem, RunStatusEnum +from app.models.manual_retry import ManualRetryRequestModel, ManualRetryResponseModel import pytest @@ -32,6 +33,7 @@ def test_router_has_correct_routes(self): assert any('/v0/namespace/{namespace_name}/state/{state_id}/errored' in path for path in paths) assert any('/v0/namespace/{namespace_name}/state/{state_id}/prune' in path for path in paths) assert any('/v0/namespace/{namespace_name}/state/{state_id}/re-enqueue-after' in path for path in paths) + assert any('/v0/namespace/{namespace_name}/state/{state_id}/manual-retry' in path for path in paths) # Graph template routes (there are two /graph/{graph_name} routes - GET and PUT) assert any('/v0/namespace/{namespace_name}/graph/{graph_name}' in path for path in paths) @@ -273,6 +275,26 @@ def test_list_graph_templates_response_validation(self): assert model.namespace == "test" assert model.count == 0 + def test_manual_retry_request_model_validation(self): + """Test ManualRetryRequestModel validation""" + # Test with valid data + valid_data = {"fanout_id": "test-fanout-id-123"} + model = ManualRetryRequestModel(**valid_data) + assert model.fanout_id == "test-fanout-id-123" + + def test_manual_retry_response_model_validation(self): + """Test ManualRetryResponseModel validation""" + from app.models.state_status_enum import StateStatusEnum + + # Test with valid data + valid_data = { + "id": "507f1f77bcf86cd799439011", + "status": StateStatusEnum.CREATED + } + model = ManualRetryResponseModel(**valid_data) + assert model.id == "507f1f77bcf86cd799439011" + assert model.status == StateStatusEnum.CREATED + @@ -295,7 +317,8 @@ def test_route_handlers_exist(self): list_graph_templates_route, get_runs_route, get_graph_structure_route, - get_node_run_details_route + get_node_run_details_route, + manual_retry_state_route ) @@ -313,6 +336,7 @@ def test_route_handlers_exist(self): assert callable(get_runs_route) assert callable(get_graph_structure_route) assert callable(get_node_run_details_route) + assert callable(manual_retry_state_route) @@ -1033,4 +1057,64 @@ async def test_get_node_run_details_route_with_invalid_api_key(self, mock_get_no assert exc_info.value.status_code == 401 assert exc_info.value.detail == "Invalid API key" - mock_get_node_run_details.assert_not_called() \ No newline at end of file + mock_get_node_run_details.assert_not_called() + + @patch('app.routes.manual_retry_state') + async def test_manual_retry_state_route_with_valid_api_key(self, mock_manual_retry_state, mock_request): + """Test manual_retry_state_route with valid API key""" + from app.routes import manual_retry_state_route + + # Arrange + mock_manual_retry_state.return_value = MagicMock() + body = ManualRetryRequestModel(fanout_id="test-fanout-id") + + # Act + result = await manual_retry_state_route("test_namespace", "507f1f77bcf86cd799439011", body, mock_request, "valid_key") + + # Assert + mock_manual_retry_state.assert_called_once() + call_args = mock_manual_retry_state.call_args + assert call_args[0][0] == "test_namespace" # namespace_name + assert str(call_args[0][1]) == "507f1f77bcf86cd799439011" # state_id as PydanticObjectId + assert call_args[0][2] == body # body + assert call_args[0][3] == "test-request-id" # x_exosphere_request_id + assert result == mock_manual_retry_state.return_value + + @patch('app.routes.manual_retry_state') + async def test_manual_retry_state_route_with_invalid_api_key(self, mock_manual_retry_state, mock_request): + """Test manual_retry_state_route with invalid API key""" + from app.routes import manual_retry_state_route + from fastapi import HTTPException + + # Arrange + body = ManualRetryRequestModel(fanout_id="test-fanout-id") + + # Act & Assert + with pytest.raises(HTTPException) as exc_info: + await manual_retry_state_route("test_namespace", "507f1f77bcf86cd799439011", body, mock_request, None) # type: ignore + + assert exc_info.value.status_code == 401 + assert exc_info.value.detail == "Invalid API key" + mock_manual_retry_state.assert_not_called() + + @patch('app.routes.manual_retry_state') + async def test_manual_retry_state_route_without_request_id(self, mock_manual_retry_state, mock_request_no_id): + """Test manual_retry_state_route without x_exosphere_request_id""" + from app.routes import manual_retry_state_route + + # Arrange + mock_manual_retry_state.return_value = MagicMock() + body = ManualRetryRequestModel(fanout_id="test-fanout-id") + + # Act + result = await manual_retry_state_route("test_namespace", "507f1f77bcf86cd799439011", body, mock_request_no_id, "valid_key") + + # Assert + mock_manual_retry_state.assert_called_once() + call_args = mock_manual_retry_state.call_args + assert call_args[0][0] == "test_namespace" # namespace_name + assert str(call_args[0][1]) == "507f1f77bcf86cd799439011" # state_id as PydanticObjectId + assert call_args[0][2] == body # body + # Should generate a UUID when no request ID is present + assert len(call_args[0][3]) > 0 # x_exosphere_request_id should be generated + assert result == mock_manual_retry_state.return_value \ No newline at end of file From edde4f03c7406d495c01f7d6dfa27c560849883a Mon Sep 17 00:00:00 2001 From: NiveditJain Date: Tue, 16 Sep 2025 17:02:57 +0530 Subject: [PATCH 3/4] fix: refine manual retry state error handling and query logic - Updated the query for fetching the state to include the namespace name, ensuring accurate state retrieval. - Changed the HTTP status code for duplicate retry state errors from 400 to 409 to better reflect the conflict nature of the error. - Simplified exception handling by removing the unnecessary re-raise of the caught exception. These changes enhance the reliability and clarity of the manual retry state functionality. --- state-manager/app/controller/manual_retry_state.py | 8 ++++---- state-manager/tests/unit/test_routes.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/state-manager/app/controller/manual_retry_state.py b/state-manager/app/controller/manual_retry_state.py index fc97bec9..17c926e5 100644 --- a/state-manager/app/controller/manual_retry_state.py +++ b/state-manager/app/controller/manual_retry_state.py @@ -13,7 +13,7 @@ async def manual_retry_state(namespace_name: str, state_id: PydanticObjectId, bo try: logger.info(f"Manual retry state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) - state = await State.find_one(State.id == state_id) + state = await State.find_one(State.id == state_id, State.namespace_name == namespace_name) if not state: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="State not found") @@ -41,9 +41,9 @@ async def manual_retry_state(namespace_name: str, state_id: PydanticObjectId, bo return ManualRetryResponseModel(id=str(retry_state.id), status=retry_state.status) except DuplicateKeyError: logger.info(f"Duplicate retry state detected for state {state_id}. A retry state with the same unique key already exists.", x_exosphere_request_id=x_exosphere_request_id) - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Duplicate retry state detected") + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="Duplicate retry state detected") - except Exception as e: + except Exception as _: logger.error(f"Error manual retry state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) - raise e + raise diff --git a/state-manager/tests/unit/test_routes.py b/state-manager/tests/unit/test_routes.py index c3ca8539..97568887 100644 --- a/state-manager/tests/unit/test_routes.py +++ b/state-manager/tests/unit/test_routes.py @@ -91,7 +91,7 @@ def test_trigger_graph_request_model_validation(self): "store": {"s1": "v1"}, "inputs": {"input1": "value1"} } - model = TriggerGraphRequestModel(**valid_data) + model = TriggerGraphRequestModel(**valid_data) # type: ignore assert model.store == {"s1": "v1"} assert model.inputs == {"input1": "value1"} From 19b3a65fb9363d8e5d6ad89c42b674c0792373df Mon Sep 17 00:00:00 2001 From: NiveditJain Date: Tue, 16 Sep 2025 17:17:57 +0530 Subject: [PATCH 4/4] feat: add unit tests for manual retry state functionality - Introduced a new test file for the manual retry state, covering various scenarios including successful state creation, error handling for not found states, and duplicate key errors. - Enhanced tests to verify logging, database error handling, and preservation of original state fields during retries. - Updated the README to include the new test file and detailed coverage of the manual retry state functionality. These changes improve the test coverage and reliability of the manual retry state feature. --- state-manager/.coverage | Bin 69632 -> 0 bytes state-manager/tests/README.md | 17 +- .../controller/test_manual_retry_state.py | 518 ++++++++++++++++++ 3 files changed, 534 insertions(+), 1 deletion(-) delete mode 100644 state-manager/.coverage create mode 100644 state-manager/tests/unit/controller/test_manual_retry_state.py diff --git a/state-manager/.coverage b/state-manager/.coverage deleted file mode 100644 index e0975eb1650d316acdcc42e9e285cde87bd1f60b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 69632 zcmeHQ32@uinI=e(-~~L%I%r9jL|K+4**Yy-wk1VACCjakIO!z2jT(lcAc-+WG5~2? zX&aeMJI*ANc4v2|o!#cz>5(?xCY?A<+DzB-PCL!jy4&nJX_K|RZMHdUXX`Xg9Ut2N zJpd#qQ7M2l({1)~r1t<4@5BFp|NoBv0S~x(&yHwH42~tF$#5#zWh=5VjBR5uXtUWo z@Lvr7+J}P;R@;G`(QPlZ>9Mul^4C1N#^!h3X+ziZyHTtEE`GW1XZ}gg?Y?`suqOgP z&<|yRGC&!i4E+CQAhXBM1u80-&wM%+9*K!5DI5_cCFQqVHL&Sp1Hq4N+Pq^RsC*Y} z@&@74+8W#xloI=cNl^-pMPp(x8Xt{D!l`I{JeZmgwWHIrIEp(u_zt={d|(?XkGDe6 zQ7DxX#{nWKMJL13RB*32)j|$X&z2Goq?9iJM2wEd@e#o$wRm$-630YIj7LOSk)kO& z+U#wo>V(Aq*rsI6rV-9#J((-RtrrsV8~C(miJUoYS~`1ghV2G^2&C&9HS77HfE zNR`2*gNkv$oAq_3T_Ynb;ffYas0V_ie^zh9!h#IJkyv6R-`9x8$Rbn_XbZH%#?tW! zt|~ItA(Gf9O5t&_P5o7A7XI9h1g_KRaRmly7jm7f3xqneo8C-?iwjg$F`w}eH&U&o zW?Y^NX_hR|jA}nkieWh|X${GHHIiD1Tq~*75%m{Pr+D42 zLluPmNoe%(zavj%{7p{ekBU>R`@%73O_JfL1XUK=V5J8L;sZy+aBz4ekxmi+#k<5k zao>}ZR{_ZFYuv+8&E#;X-d84MISSn>?zZF}?b~4EL|DeHklYv?NhD%oIIi>!P-43b z=ZwlhBE$S%X+sV}p28U8`x=Q%TIYhB`Hw-R)9M#&OiqU5qiTPJyI}Btxe@8P>#B zckOvjGUvyfzwCwv7e+gJKf&DT=)k+m8F1sKSa3vPK|4wv>v)@akdi?cZ16`vlmW^B zWq>k38K4YM1}FoR0m=YnfHFWCpbUKY7;rFlhQsl{9sQRLorQ#cCy2`9*BrZvi>VHnLPlouUrGz4OXxv5f%i>DOf6* zkGN$uAo{EzhC4gfj|~3d?)tVK8uXY$`M^g_9Eiv3nIj)EkAsOR%Iw zEC$S{lQJyARd7NnaWWYLy8z&pPE(43fn)M;N{WtyIYX0p7eKak07!`isWc!>y`Aj< zRBIH7oL(N5#PO({0&(IXjvN9PNy3UitlW+^fUADLaKs}3ZKxH{%8a7vxX%|5kao8K zk`*gTVhE&3!^+b!tgsBrm&4cqTUG*~aracLsFYTWCQ&XbtTICT_VI z;#9AdMOcxR%5&F&2EaiU^dxdAcp%Q^WdP$b3Zsgnu%zq&#*TWxFkEwVzK%u7XB?v~ zHG`HXk(`p!krXUm&H3rhIzThrNEDPUi;9H{nh|aeTqG^Q!vZv~0JQ#EKx0gjKRN>a zBrH^h@i+37p+<$D{eYm=xhr45aQ!Km$mzeG%XUjC!hygFBgQ?~pafHN^ncX%Z^G-3!5L z?kKE60|iiz$L=~K-ne$is6<* zU$;~wt!nMy0myK#A%oT6@h~t}8@^iG+=?HX5{9-A#2WMinqjxlp~Vv8a!3t10B5rg za18f{IvkygZu6R^NkuVw1_Paq0!|g<_xMrYD9F#X4 zf#k~Tm>6^cgkh_w5Mo>zw(7>_15Us(-7jQuv_W}4Yk?7+IztX9Z@9nJ`@Sj)p%z!| z&2~UC>}EP7oiyEy!zLT8&Pq85!i(10A!X3_aX^mH_tEE2uYZI4CGHB>S=Uw0C!Fo< zY4-E%pyS()G5bO0b=yH^gekH;B3uxf`NO^kxM%p+xI4Tv!XrY=ciw%MOK{!iyUNEo z7Tcfl>~~JH3_s#&wfFjO<7@oSxvzJZG532j-c9&)%BRb**SJ?I@&DTE*vxw4g|gwj zk@&xR51VPQguPJwzh*0&@mWIDS$FkjHp3YPDH#8EUCCzZ4MXUhY@zsnl{v+f$WaU1 zN$m#L?;K(?C03-;fJpZSKpneSc-V%~3ZOPF+R%4?{NKKl%~Zc{IIXE^+sf%t#99A$v1?0|N>*D`fQ;m#DRl)eb##E#H9`S$m6^?1c8qJIU zmzZ*DRiSLqo^{A7Q@rS`l!em9z?G(uRDlY{{}on{3daA7O?eGy2h$Mx_&;E!yl!Mt zZmRtJ_`l4Ii_DAvORcn@qx0kcl8>;NfMHKnEn7JLUvw>-=`;&9CrB?8{}-ETyHpXO zde4C-KS0}6G{9zfvtrJR|BOOijtzUyc7grs`>~xYqW$O+_HN@;W9sQ&v$S3dH{|OZ~#U_}^)T z5uG}$sWx2i`*Y)ehbcG9AwdK!SW$uu8kCE;Zuk38K4YM1}FoR z0m=YnfHFWC_)s$7fLjK5y#Aj-7j5uIKa>H=0A+wOKpCJ6PzERilmW^BWq>k38K4Y& z2pMoVcn`k*-{(JKL+8*GVb-Ty=eV|H#6(j!Ircx{@9uL{ML=_Tsfra@pE* z+s}=5Pwjg7^rm0iJ>BqCY4_2K)3yFJuvxk0=#CfOsQ=b(*J{`kRx@Q?khgTbd-KW3 z=imMDt0zAh`RY3_?5`D8!Phmbj=p{V(*IoAe@W&#VZW$zA6L`?8%sOheg5vLnddgV zJ#(Dl+VM~AFUwpT-e`MS7FuDWrZpaY{da-c-@o)5PYdjqwwyZgLYiBNk6C$q(z^mS z%2%Adaq@Y_+l;rHpJng+UXQm4Z#OC1jd;89S$5-Fue+ARwy=E6-vD`K!}WIC6|da& zx@Q^el`gydsml(3J^rEo`u>^z>1oEh6!yxOo|=8@mD%z-*lw$P)^qUXOB)`ZUAyL$|5FZm%PF0JyYtO!O!^LWE4c@NlpS{KJYN&?o9o3g--@fUc@#!xP@4WlL zPri9NV0-l8e>*=N@A#LWABQ#Ow=k|HaG0=!;i@1ns*1aWO4z8WeD3hcqu-zX=?R9f zfc+&Ek0U)KRnu_LnZYN4|Xd@`aNp-#J~&6~i}0#pjsL zBG?!$I(zx@CpSF$rBjn{UUTWxGcUk#mu5~+cRe;cGyBfzOYz4K47mUMqc5K6{`rv^ zhC{e48iT+XHiE*<5AAy~G&}q0=i$wJ_!s`+fZ&IX8h`&2M}NHM8^;d0Tt3(rd<^Hs z2Y8>|zO{Ebs`u{dWL z!2x-VAmBG2wf{!@^**oc`t>ARw~l0M*OIKen`CR&kZkp8l67^F zY}G1~b#{`hql0Aa?Ide!BUx)J$y!=SwsIxOR;(adb2G`Bnn>2zNV4V2N!HLnvSrIi zR$ouDrAtXxS4Xm7kYu&BB&(?*S#>qZmMkGzRTar9D@j&SL9)e*Nfro@th}6LWo0BQ zEhSk=3CR{MB3W@U$%=|dh7idFfn+>SGQXc>J|D@vUXpn{By*FA0EXkpp36ltr;}tX zOEQOpWOh3X69DA;zpdy);A&I_$^d16GC&!i3{VCr1C#;E0A+wOKpCJ6PzDxc0AK&7 z@&AJK&>tuRlmW^BWq>k38K4YM1}FoR0m=YnfHLr*V1TUu7ybj^|Nnb<0>E$3E9h6~ zB&-4WcXS**g?&ZA7oyKr!$N&MYczT3qmu0*!KYU40nfjMtDSs`OdrVatW^cd{_B6$71_ap8d{g zmf=S{t@d93ZG4UYIrsJMGUk47#=9wJ*e=Ij<8L7iN8|A7(zVyInf1mCMMpDSof~U z)%ISS8;jNBm@+J#`m*0hLXttVKCACmPbXzjN`>;@BMX45>`&PS3=5;a&NC%N$HzrU zc_lZ1beeiQYVOsS!x|6~3xMD$Dz>}LN+XHmQ8@+T;8*#|A#jl-c!UUW?WW#Gt134S zj(7y1wOM^tXg(q!wOX~JB!)nmv=|DHVTI-Bcsv{fz!vAUc`GWVbMIAOY3j|cIx8ln zbR6Dus|mBhs&+9jBbZOICRA@?PtdM;H=C(5>=in}2w=X;<4V+I6%W7Mp1ev|4vkEO zKx=3x0I$&{*oF+tr@EOeM&MQ8dG5O0)GKjydJ?%5JP@bBb)RjkS#cDW6xje{nW>sH zN0fZV0j-|9!?w$sL~=?>M^b5U{H$zCO}#_Apln%G5U%e1s2WL2k{H*#0?>l)GTTP8 zs*+dvCM4Ws)tYK#R9}MfbE*v>)p%NM#|=x83d?)tVQ@+O%IdgyK&Mi*nO9aTIJ!=6 z35c=Huo$XB;sG&|PC*lpHF}jPUUWv!LTT^7u7rWh-44+-|c`+%*N0quK0SzEU$w$~sz$^~#WI`!%G8x03Prw0c z(Y0)*(=1eeyBFFQjzveqczCM510HaTO|@N`sw2a5pfxdz{I}T#&59{X5;(K2ZXv6~ zHw{UvS_1zXQ*T%t8LSo)Q!rK=zFOPdD#AcF43!PI;`!HXC#;~w65|jMgyUftqXLfK z)Qj15Qt4dOXNn@#ZiE)q_8~Ezo&;gMW@<{#{v%;Bp!Enbi@l5bqWH~hhBHiQ?({k) zIEebbU=~7*E5laZ*xY4mu%n8S!_fxiPE#$P3ZqkpHPwde`xR9sWjqL$b(nH99g|=T6U7599uUd!zf-Tr^L& zC&1U;r^#&oh24(HZ55{BajDFM?z zcxGdrxeG5aZicA|?Xm)?BLLEAYLKr635BP8OnFHb2~REH0eo7{#(`G!dVe$anluF_ zo~`gxcp}Tn{DwAn`GH@+`3smIGF9bNiz<^t**Wdpr0G&qA!#0(vIBKznWZ>={`~C# zo8b*te^nH1ekm&qGSxb2C?DM9gB5W`3rt*^>O57o^cM~A{r_ui(>C-ZJm0@Zcvm=l~$SNR9{FubSlg#VDgAL9RS`u6x--lN`w-i=n2$N&B?Dhk1?pIumA)?FgJs`>cXHGWHss zwwP3||9`xjwPn^DE~LAn{>QKXU%3Xb8?0b2bp3zJYC!Z^LDbFeZ0Z6K&L~L1QOfRB z08wugLO02@(DnZVhj4sx&FVq*`js} zUH{+J1b}rG1Y6+x|E@-hcne(rA6#yRIMr(lT>l?v032jN&;0BEo0kEM;aOx=9Ayll z?0|>vs0R$gHAm;`Sd{#WRTyol8MHi!3XE!Y)&ZK~Mxvl>SyU`s(2Q_%;O1Tb@2>?k zv*R3HS3Xn&WGg5z< zney{zZ~EXKSHmtc@B05>37p+<$D{fD{9*L=MSv49>ZvMD;p_j`6a!SJNvQmGui*9n zokf6axb4zZUGVzy~;afmXG4@Bn1E*Z5--MuUDpGwc>R zv^i5@n|**|xIfh4=qkuIuNjK6jmG?WpFX&)+i*?E*-|q%747U@x;+laraZuiH z1d=PSV`9(+5QeRyLgcR7XYb7&Z~~6$ejyvmYlHHB)&e6sb%q>J-f(}b_x-uo|2NwK z$*`O0kaW^?!_EA%(dw)#hvR>zaLfj8{5yoU!d?Bxw9uU9lmW^BWq>k38K4YM1}FoR z0m=YnfHFWC_`k!zdRXf}Wi;+pr$5yd(c5|<>^3}krOro?`H)clO>mof0Lbjz$Q+3E VW*`dQN4FZDaF8&{MD{MJ{{@07ua5u# diff --git a/state-manager/tests/README.md b/state-manager/tests/README.md index aafe1dff..c2a89da1 100644 --- a/state-manager/tests/README.md +++ b/state-manager/tests/README.md @@ -14,6 +14,7 @@ tests/ │ ├── test_errored_state.py │ ├── test_get_graph_template.py │ ├── test_get_secrets.py +│ ├── test_manual_retry_state.py │ ├── test_register_nodes.py │ └── test_upsert_graph_template.py └── README.md @@ -80,7 +81,21 @@ The unit tests cover all controller functions in the state-manager: - ✅ Complex schema handling - ✅ Database error handling -### 8. `upsert_graph_template.py` +### 8. `manual_retry_state.py` +- ✅ Successful manual retry state creation +- ✅ State not found scenarios +- ✅ Duplicate retry state detection (DuplicateKeyError) +- ✅ Different fanout_id handling +- ✅ Complex inputs and multiple parents preservation +- ✅ Database errors during state lookup +- ✅ Database errors during state save +- ✅ Database errors during retry state insert +- ✅ Empty inputs and parents handling +- ✅ Namespace mismatch scenarios +- ✅ Field preservation and reset logic +- ✅ Logging verification + +### 9. `upsert_graph_template.py` - ✅ Existing template updates - ✅ New template creation - ✅ Empty nodes handling diff --git a/state-manager/tests/unit/controller/test_manual_retry_state.py b/state-manager/tests/unit/controller/test_manual_retry_state.py new file mode 100644 index 00000000..072372e0 --- /dev/null +++ b/state-manager/tests/unit/controller/test_manual_retry_state.py @@ -0,0 +1,518 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock, patch +from fastapi import HTTPException, status +from beanie import PydanticObjectId +from pymongo.errors import DuplicateKeyError + +from app.controller.manual_retry_state import manual_retry_state +from app.models.manual_retry import ManualRetryRequestModel, ManualRetryResponseModel +from app.models.state_status_enum import StateStatusEnum + + +class TestManualRetryState: + """Test cases for manual_retry_state function""" + + @pytest.fixture + def mock_request_id(self): + return "test-request-id" + + @pytest.fixture + def mock_namespace(self): + return "test_namespace" + + @pytest.fixture + def mock_state_id(self): + return PydanticObjectId() + + @pytest.fixture + def mock_manual_retry_request(self): + return ManualRetryRequestModel( + fanout_id="test-fanout-id-123" + ) + + @pytest.fixture + def mock_original_state(self): + state = MagicMock() + state.id = PydanticObjectId() + state.node_name = "test_node" + state.namespace_name = "test_namespace" + state.identifier = "test_identifier" + state.graph_name = "test_graph" + state.run_id = "test_run_id" + state.status = StateStatusEnum.EXECUTED + state.inputs = {"key": "value"} + state.outputs = {"result": "success"} + state.error = "Original error" + state.parents = {"parent1": PydanticObjectId()} + state.does_unites = False + state.save = AsyncMock() + return state + + @pytest.fixture + def mock_retry_state(self): + retry_state = MagicMock() + retry_state.id = PydanticObjectId() + retry_state.status = StateStatusEnum.CREATED + retry_state.insert = AsyncMock(return_value=retry_state) + return retry_state + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_success( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_original_state, + mock_retry_state, + mock_request_id + ): + """Test successful manual retry state creation""" + # Arrange + mock_state_class.find_one = AsyncMock(return_value=mock_original_state) + mock_state_class.return_value = mock_retry_state + + # Act + result = await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + # Assert + assert isinstance(result, ManualRetryResponseModel) + assert result.id == str(mock_retry_state.id) + assert result.status == StateStatusEnum.CREATED + + # Verify State.find_one was called with correct parameters + mock_state_class.find_one.assert_called_once() + call_args = mock_state_class.find_one.call_args[0] + # Check that both conditions were passed + assert len(call_args) == 2 + + # Verify original state was updated to RETRY_CREATED + assert mock_original_state.status == StateStatusEnum.RETRY_CREATED + mock_original_state.save.assert_called_once() + + # Verify retry state was created with correct attributes + mock_state_class.assert_called_once() + retry_state_args = mock_state_class.call_args[1] + assert retry_state_args['node_name'] == mock_original_state.node_name + assert retry_state_args['namespace_name'] == mock_original_state.namespace_name + assert retry_state_args['identifier'] == mock_original_state.identifier + assert retry_state_args['graph_name'] == mock_original_state.graph_name + assert retry_state_args['run_id'] == mock_original_state.run_id + assert retry_state_args['status'] == StateStatusEnum.CREATED + assert retry_state_args['inputs'] == mock_original_state.inputs + assert retry_state_args['outputs'] == {} + assert retry_state_args['error'] is None + assert retry_state_args['parents'] == mock_original_state.parents + assert retry_state_args['does_unites'] == mock_original_state.does_unites + assert retry_state_args['fanout_id'] == mock_manual_retry_request.fanout_id + + # Verify retry state was inserted + mock_retry_state.insert.assert_called_once() + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_not_found( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ): + """Test when original state is not found""" + # Arrange + mock_state_class.find_one = AsyncMock(return_value=None) + + # Act & Assert + with pytest.raises(HTTPException) as exc_info: + await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND + assert exc_info.value.detail == "State not found" + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_duplicate_key_error( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_original_state, + mock_retry_state, + mock_request_id + ): + """Test when duplicate retry state is detected""" + # Arrange + mock_state_class.find_one = AsyncMock(return_value=mock_original_state) + mock_retry_state.insert = AsyncMock(side_effect=DuplicateKeyError("Duplicate key")) + mock_state_class.return_value = mock_retry_state + + # Act & Assert + with pytest.raises(HTTPException) as exc_info: + await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + assert exc_info.value.status_code == status.HTTP_409_CONFLICT + assert exc_info.value.detail == "Duplicate retry state detected" + + # Verify original state was not updated since duplicate was detected + mock_original_state.save.assert_not_called() + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_with_different_fanout_id( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_original_state, + mock_retry_state, + mock_request_id + ): + """Test manual retry with different fanout_id""" + # Arrange + different_fanout_request = ManualRetryRequestModel( + fanout_id="different-fanout-id-456" + ) + mock_state_class.find_one = AsyncMock(return_value=mock_original_state) + mock_state_class.return_value = mock_retry_state + + # Act + result = await manual_retry_state( + mock_namespace, + mock_state_id, + different_fanout_request, + mock_request_id + ) + + # Assert + assert isinstance(result, ManualRetryResponseModel) + assert result.id == str(mock_retry_state.id) + assert result.status == StateStatusEnum.CREATED + + # Verify retry state was created with the different fanout_id + retry_state_args = mock_state_class.call_args[1] + assert retry_state_args['fanout_id'] == "different-fanout-id-456" + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_with_complex_inputs_and_parents( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_retry_state, + mock_request_id + ): + """Test manual retry with complex inputs and multiple parents""" + # Arrange + complex_state = MagicMock() + complex_state.id = PydanticObjectId() + complex_state.node_name = "complex_node" + complex_state.namespace_name = "test_namespace" + complex_state.identifier = "complex_identifier" + complex_state.graph_name = "complex_graph" + complex_state.run_id = "complex_run_id" + complex_state.status = StateStatusEnum.ERRORED + complex_state.inputs = { + "nested_data": {"key1": "value1", "key2": [1, 2, 3]}, + "simple_value": "test", + "number": 42 + } + complex_state.outputs = {"previous_result": "some_output"} + complex_state.error = "Complex error message" + complex_state.parents = { + "parent1": PydanticObjectId(), + "parent2": PydanticObjectId(), + "parent3": PydanticObjectId() + } + complex_state.does_unites = True + complex_state.save = AsyncMock() + + mock_state_class.find_one = AsyncMock(return_value=complex_state) + mock_state_class.return_value = mock_retry_state + + # Act + result = await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + # Assert + assert isinstance(result, ManualRetryResponseModel) + + # Verify retry state preserves complex data structures + retry_state_args = mock_state_class.call_args[1] + assert retry_state_args['inputs'] == complex_state.inputs + assert retry_state_args['parents'] == complex_state.parents + assert retry_state_args['does_unites'] == complex_state.does_unites + assert retry_state_args['outputs'] == {} # Should be reset + assert retry_state_args['error'] is None # Should be reset + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_database_error_on_find( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ): + """Test handling of database error during state lookup""" + # Arrange + mock_state_class.find_one = AsyncMock(side_effect=Exception("Database connection error")) + + # Act & Assert + with pytest.raises(Exception) as exc_info: + await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + assert str(exc_info.value) == "Database connection error" + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_database_error_on_save( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_original_state, + mock_retry_state, + mock_request_id + ): + """Test handling of database error during original state save""" + # Arrange + mock_state_class.find_one = AsyncMock(return_value=mock_original_state) + mock_state_class.return_value = mock_retry_state + mock_original_state.save = AsyncMock(side_effect=Exception("Save operation failed")) + + # Act & Assert + with pytest.raises(Exception) as exc_info: + await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + assert str(exc_info.value) == "Save operation failed" + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_database_error_on_insert( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_original_state, + mock_retry_state, + mock_request_id + ): + """Test handling of database error during retry state insert (non-duplicate)""" + # Arrange + mock_state_class.find_one = AsyncMock(return_value=mock_original_state) + mock_retry_state.insert = AsyncMock(side_effect=Exception("Insert operation failed")) + mock_state_class.return_value = mock_retry_state + + # Act & Assert + with pytest.raises(Exception) as exc_info: + await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + assert str(exc_info.value) == "Insert operation failed" + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_empty_inputs_and_parents( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_retry_state, + mock_request_id + ): + """Test manual retry with empty inputs and parents""" + # Arrange + empty_state = MagicMock() + empty_state.id = PydanticObjectId() + empty_state.node_name = "empty_node" + empty_state.namespace_name = "test_namespace" + empty_state.identifier = "empty_identifier" + empty_state.graph_name = "empty_graph" + empty_state.run_id = "empty_run_id" + empty_state.status = StateStatusEnum.EXECUTED + empty_state.inputs = {} + empty_state.outputs = {} + empty_state.error = None + empty_state.parents = {} + empty_state.does_unites = False + empty_state.save = AsyncMock() + + mock_state_class.find_one = AsyncMock(return_value=empty_state) + mock_state_class.return_value = mock_retry_state + + # Act + result = await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + # Assert + assert isinstance(result, ManualRetryResponseModel) + + # Verify retry state handles empty collections correctly + retry_state_args = mock_state_class.call_args[1] + assert retry_state_args['inputs'] == {} + assert retry_state_args['parents'] == {} + assert retry_state_args['outputs'] == {} + assert retry_state_args['error'] is None + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_namespace_mismatch( + self, + mock_state_class, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ): + """Test manual retry with namespace that doesn't match any state""" + # Arrange + different_namespace = "different_namespace" + mock_state_class.find_one = AsyncMock(return_value=None) + + # Act & Assert + with pytest.raises(HTTPException) as exc_info: + await manual_retry_state( + different_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND + assert exc_info.value.detail == "State not found" + + # Verify find_one was called with the different namespace + mock_state_class.find_one.assert_called_once() + + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_preserves_all_original_fields( + self, + mock_state_class, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_retry_state, + mock_request_id + ): + """Test that all relevant fields from original state are preserved in retry state""" + # Arrange + original_state = MagicMock() + original_state.id = PydanticObjectId() + original_state.node_name = "preserve_test_node" + original_state.namespace_name = "preserve_test_namespace" + original_state.identifier = "preserve_test_identifier" + original_state.graph_name = "preserve_test_graph" + original_state.run_id = "preserve_test_run_id" + original_state.status = StateStatusEnum.EXECUTED + original_state.inputs = {"preserve": "input_data"} + original_state.outputs = {"should_be": "reset"} + original_state.error = "should_be_reset" + original_state.parents = {"preserve_parent": PydanticObjectId()} + original_state.does_unites = True + original_state.save = AsyncMock() + + mock_state_class.find_one = AsyncMock(return_value=original_state) + mock_state_class.return_value = mock_retry_state + + # Act + await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + # Assert - verify all fields are correctly set + retry_state_args = mock_state_class.call_args[1] + + # Fields that should be preserved + assert retry_state_args['node_name'] == original_state.node_name + assert retry_state_args['namespace_name'] == original_state.namespace_name + assert retry_state_args['identifier'] == original_state.identifier + assert retry_state_args['graph_name'] == original_state.graph_name + assert retry_state_args['run_id'] == original_state.run_id + assert retry_state_args['inputs'] == original_state.inputs + assert retry_state_args['parents'] == original_state.parents + assert retry_state_args['does_unites'] == original_state.does_unites + assert retry_state_args['fanout_id'] == mock_manual_retry_request.fanout_id + + # Fields that should be reset/set to specific values + assert retry_state_args['status'] == StateStatusEnum.CREATED + assert retry_state_args['outputs'] == {} + assert retry_state_args['error'] is None + + @patch('app.controller.manual_retry_state.logger') + @patch('app.controller.manual_retry_state.State') + async def test_manual_retry_state_logging_calls( + self, + mock_state_class, + mock_logger, + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_original_state, + mock_retry_state, + mock_request_id + ): + """Test that appropriate logging calls are made""" + # Arrange + mock_state_class.find_one = AsyncMock(return_value=mock_original_state) + mock_state_class.return_value = mock_retry_state + + # Act + await manual_retry_state( + mock_namespace, + mock_state_id, + mock_manual_retry_request, + mock_request_id + ) + + # Assert - verify logging calls were made + assert mock_logger.info.call_count >= 2 # At least initial log and success log + + # Check that the initial log contains expected information + first_call_args = mock_logger.info.call_args_list[0] + assert str(mock_state_id) in first_call_args[0][0] + assert mock_namespace in first_call_args[0][0] + assert first_call_args[1]['x_exosphere_request_id'] == mock_request_id + + # Check that the success log contains retry state id + second_call_args = mock_logger.info.call_args_list[1] + assert str(mock_retry_state.id) in second_call_args[0][0] + assert str(mock_state_id) in second_call_args[0][0] + assert second_call_args[1]['x_exosphere_request_id'] == mock_request_id