intuned-runtime 1.3.1__py3-none-any.whl → 1.3.2__py3-none-any.whl

intuned_cli/controller/__test__/test_api.py

@@ -2,13 +2,16 @@ from dataclasses import dataclass
 from typing import Any
 from typing import Generator
 from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
 from unittest.mock import Mock
 from unittest.mock import patch
 
 import pytest
 
+from intuned_cli.controller.api import get_cli_run_options as get_cli_run_options
 from runtime.types.run_types import PayloadToAppend
 from runtime.types.run_types import StorageState
+from runtime.types.run_types import TracingDisabled
 
 
 def get_mock_console():
@@ -18,23 +21,6 @@ def get_mock_console():
     return mock_console
 
 
-class AsyncContextManagerMock:
-    def __init__(self):
-        self.aenter_called = False
-        self.aexit_called = False
-        self.exception_info = None
-
-    async def __aenter__(self):
-        self.aenter_called = True
-        return self
-
-    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
-        self.aexit_called = True
-        self.exception_info = (exc_type, exc_val, exc_tb)
-        # Return False/None to propagate exceptions
-        return False
-
-
 @dataclass
 class SharedMocks:
     console: Mock
@@ -69,6 +55,8 @@ def shared_mocks() -> Generator[SharedMocks, Any, None]:
 class AttemptApiMocks:
     extendable_timeout: AsyncMock
     run_api: AsyncMock
+    cli_trace: AsyncMock
+    get_cli_run_options: AsyncMock
 
 
 @pytest.fixture
@@ -76,13 +64,17 @@ def attempt_api_mocks() -> Generator[AttemptApiMocks, Any, None]:
     """Mock dependencies for attempt_api tests."""
     _mock_timeout_patch = patch("intuned_cli.controller.api.extendable_timeout")
     _mock_run_api_patch = patch("intuned_cli.controller.api.run_api", new_callable=AsyncMock)
+    _mock_cli_trace_patch = patch("intuned_cli.controller.api.cli_trace")
+    _mock_get_cli_run_options_patch = patch("intuned_cli.controller.api.get_cli_run_options", new_callable=AsyncMock)
 
     with (
         _mock_timeout_patch as mock_timeout,
         _mock_run_api_patch as mock_run_api,
+        _mock_cli_trace_patch as mock_cli_trace,
+        _mock_get_cli_run_options_patch as mock_get_cli_run_options,
     ):
         # Setup default return values
-        mock_timeout.return_value = AsyncContextManagerMock()
+        mock_timeout.return_value = MagicMock()
 
         # Mock run_api to return success by default
         mock_result = Mock()
@@ -90,9 +82,23 @@ def attempt_api_mocks() -> Generator[AttemptApiMocks, Any, None]:
         mock_result.payload_to_append = []
         mock_run_api.return_value = mock_result
 
+        mock_cli_trace_return_value = MagicMock()
+        mock_cli_trace_return_value.__enter__.return_value = TracingDisabled()
+        mock_cli_trace.return_value = mock_cli_trace_return_value
+
+        # _mock_get_cli_run_options_patch keeps the original implementation but always passes keep_browser_open=False
+        async def get_cli_run_options_side_effect(*args: Any, **kwargs: Any):
+            # Force keep_browser_open to False
+            kwargs["keep_browser_open"] = False
+            return await get_cli_run_options(*args, **kwargs)
+
+        mock_get_cli_run_options.side_effect = get_cli_run_options_side_effect
+
         yield AttemptApiMocks(
             extendable_timeout=mock_timeout,
             run_api=mock_run_api,
+            cli_trace=mock_cli_trace,
+            get_cli_run_options=mock_get_cli_run_options,
         )
 
 
@@ -146,10 +152,77 @@ class TestAttemptApi:
             parameters={},
             headless=False,
             timeout=6000,
+            trace_id=None,
+            keep_browser_open=False,
         )
 
         attempt_api_mocks.extendable_timeout.assert_called_once_with(6000)
 
+    @pytest.mark.asyncio
+    async def test_attempt_api_uses_trace_context_manager(self, attempt_api_mocks: AttemptApiMocks):
+        """Test that attempt_api uses the trace context manager."""
+        # Import inside test to avoid circular import issues
+        from intuned_cli.controller.api import attempt_api
+
+        await attempt_api(
+            api_name="testApi",
+            parameters={},
+            headless=False,
+            timeout=6000,
+            trace_id=None,
+            keep_browser_open=False,
+        )
+
+        attempt_api_mocks.cli_trace.assert_called_once_with(None)
+
+        attempt_api_mocks.cli_trace.reset_mock()
+        await attempt_api(
+            api_name="testApi",
+            parameters={},
+            headless=False,
+            timeout=6000,
+            trace_id="trace-id",
+            keep_browser_open=False,
+        )
+
+        attempt_api_mocks.cli_trace.assert_called_once_with("trace-id")
+
+    @pytest.mark.asyncio
+    async def test_attempt_api_uses_get_cli_run_options(self, attempt_api_mocks: AttemptApiMocks):
+        """Test that attempt_api calls get_cli_run_options with the correct parameters."""
+        # Import inside test to avoid circular import issues
+        from intuned_cli.controller.api import attempt_api
+
+        await attempt_api(
+            api_name="testApi",
+            parameters={},
+            headless=False,
+            timeout=6000,
+            trace_id=None,
+            keep_browser_open=False,
+        )
+
+        attempt_api_mocks.get_cli_run_options.assert_called_once_with(
+            headless=False,
+            proxy=None,
+            keep_browser_open=False,
+        )
+
+        attempt_api_mocks.get_cli_run_options.reset_mock()
+        await attempt_api(
+            api_name="testApi",
+            parameters={},
+            headless=True,
+            proxy="proxy",
+            timeout=6000,
+            trace_id=None,
+            keep_browser_open=True,
+        )
+        call_args = attempt_api_mocks.get_cli_run_options.call_args[1]
+        assert call_args["headless"] is True
+        assert call_args["proxy"] is not None
+        assert call_args["keep_browser_open"] is True
+
     @pytest.mark.asyncio
     async def test_attempt_api_calls_run_api_with_correct_parameters_and_parses_proxy(
         self, attempt_api_mocks: AttemptApiMocks
@@ -175,6 +248,8 @@ class TestAttemptApi:
                 auth=auth,
                 proxy="proxy",
                 timeout=999999999,
+                trace_id=None,
+                keep_browser_open=False,
             )
 
             # mock_parse_proxy.assert_called_once_with("proxy")
@@ -208,6 +283,8 @@ class TestAttemptApi:
             parameters="inputData",
             headless=False,
             timeout=999999999,
+            trace_id=None,
+            keep_browser_open=False,
         )
 
         assert result is expected_result
@@ -227,12 +304,50 @@ class TestAttemptApi:
                 parameters={},
                 headless=False,
                 timeout=999999999,
+                trace_id=None,
+                keep_browser_open=False,
             )
 
 
 class TestExecuteRunApiCli:
     """Test suite for API controller functions."""
 
+    @pytest.mark.asyncio
+    async def test_execute_run_api_cli_calls_attempt_api_with_trace_id_correctly(
+        self, execute_cli_mocks: ExecuteCLIMocks
+    ):
+        """Test that execute_run_api_cli calls attemptApi with trace_id correctly."""
+        from intuned_cli.controller.api import execute_run_api_cli
+
+        await execute_run_api_cli(
+            api_name="testApi",
+            input_data={},
+            retries=3,
+            headless=False,
+            timeout=30,
+            trace=False,
+            keep_browser_open=False,
+        )
+
+        execute_cli_mocks.attempt_api.assert_called()
+        call_args = execute_cli_mocks.attempt_api.call_args
+        assert call_args[1]["trace_id"] is None
+
+        execute_cli_mocks.attempt_api.reset_mock()
+
+        await execute_run_api_cli(
+            api_name="testApi",
+            input_data={},
+            retries=3,
+            headless=False,
+            timeout=30,
+            trace=True,
+            keep_browser_open=False,
+        )
+        execute_cli_mocks.attempt_api.assert_called()
+        call_args = execute_cli_mocks.attempt_api.call_args
+        assert isinstance(call_args[1]["trace_id"], str)
+
     @pytest.mark.asyncio
     async def test_execute_run_api_cli_calls_attempt_api_once_if_success(self, execute_cli_mocks: ExecuteCLIMocks):
         """Test that execute_run_api_cli calls attemptApi once if successful."""
@@ -244,6 +359,8 @@ class TestExecuteRunApiCli:
             retries=3,
             headless=False,
             timeout=30,
+            trace=False,
+            keep_browser_open=False,
         )
 
         execute_cli_mocks.attempt_api.assert_called_once()
@@ -264,6 +381,8 @@ class TestExecuteRunApiCli:
                 retries=10,
                 headless=False,
                 timeout=30,
+                trace=False,
+                keep_browser_open=False,
             )
 
         assert execute_cli_mocks.attempt_api.call_count == 10
@@ -284,6 +403,8 @@ class TestExecuteRunApiCli:
                 retries=3,
                 headless=False,
                 timeout=30,
+                trace=False,
+                keep_browser_open=False,
             )
 
         execute_cli_mocks.attempt_api.assert_called_once()
@@ -305,6 +426,8 @@ class TestExecuteRunApiCli:
             retries=10,
             headless=False,
             timeout=30,
+            trace=False,
+            keep_browser_open=False,
         )
 
         assert execute_cli_mocks.attempt_api.call_count == 2
@@ -340,6 +463,8 @@ class TestExecuteRunApiCli:
                 retries=10,
                 headless=False,
                 timeout=30,
+                trace=False,
+                keep_browser_open=False,
             )
 
             assert mock_validate_auth.call_count == 2
@@ -363,6 +488,8 @@ class TestExecuteRunApiCli:
             retries=1,
             headless=False,
             timeout=30,
+            trace=False,
+            keep_browser_open=False,
         )
 
         execute_cli_mocks.execute_run_validate_auth_session_cli.assert_not_called()
@@ -393,6 +520,8 @@ class TestExecuteRunApiCli:
                 retries=10,
                 headless=False,
                 timeout=30,
+                trace=False,
+                keep_browser_open=False,
             )
 
         # Verify auth validation was called
@@ -425,6 +554,8 @@ class TestExecuteRunApiCli:
             retries=1,
             headless=False,
             timeout=30,
+            trace=False,
+            keep_browser_open=False,
         )
 
         execute_cli_mocks.write_results_to_file.assert_called_once()
@@ -446,6 +577,8 @@ class TestExecuteRunApiCli:
             retries=1,
             headless=False,
             timeout=30,
+            trace=False,
+            keep_browser_open=False,
         )
 
         shared_mocks.assert_api_file_exists.assert_called_once_with("api", "testApi")
@@ -465,6 +598,8 @@ class TestExecuteAttemptApiCli:
             input_data={},
             headless=False,
             timeout=30,
+            trace=False,
+            keep_browser_open=False,
         )
 
         execute_cli_mocks.attempt_api.assert_called_once()
@@ -479,6 +614,8 @@ class TestExecuteAttemptApiCli:
                 input_data={},
                 headless=False,
                 timeout=30,
+                trace=False,
+                keep_browser_open=False,
             )
 
         execute_cli_mocks.attempt_api.assert_called_once()
@@ -502,6 +639,8 @@ class TestExecuteAttemptApiCli:
             output_file="output.json",
             headless=False,
             timeout=30,
+            trace=False,
+            keep_browser_open=False,
         )
 
         # Verify file was written
@@ -523,6 +662,8 @@ class TestExecuteAttemptApiCli:
             input_data={},
             headless=False,
             timeout=30,
+            trace=False,
+            keep_browser_open=False,
         )
 
         shared_mocks.assert_api_file_exists.assert_called_once_with("api", "testApi")