kiln-ai 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

kiln_ai/utils/test_async_job_runner.py

@@ -3,30 +3,51 @@ from unittest.mock import AsyncMock, patch
 
 import pytest
 
-from kiln_ai.utils.async_job_runner import AsyncJobRunner, Progress
+from kiln_ai.utils.async_job_runner import (
+    AsyncJobRunner,
+    AsyncJobRunnerObserver,
+    Progress,
+)
+
+
+@pytest.fixture
+def mock_async_run_job_fn_success():
+    return AsyncMock(return_value=True)
+
+
+@pytest.fixture
+def mock_async_run_job_fn_failure():
+    return AsyncMock(return_value=False)
 
 
 @pytest.mark.parametrize("concurrency", [0, -1, -25])
-def test_invalid_concurrency_raises(concurrency):
+def test_invalid_concurrency_raises(concurrency, mock_async_run_job_fn_success):
     with pytest.raises(ValueError):
-        AsyncJobRunner(concurrency=concurrency)
+        AsyncJobRunner(
+            concurrency=concurrency,
+            jobs=[],
+            run_job_fn=mock_async_run_job_fn_success,
+        )
 
 
 # Test with and without concurrency
 @pytest.mark.parametrize("concurrency", [1, 25])
 @pytest.mark.asyncio
-async def test_async_job_runner_status_updates(concurrency):
+async def test_async_job_runner_status_updates(
+    concurrency, mock_async_run_job_fn_success
+):
     job_count = 50
     jobs = [{"id": i} for i in range(job_count)]
 
-    runner = AsyncJobRunner(concurrency=concurrency)
-
-    # fake run_job that succeeds
-    mock_run_job_success = AsyncMock(return_value=True)
+    runner = AsyncJobRunner(
+        concurrency=concurrency,
+        jobs=jobs,
+        run_job_fn=mock_async_run_job_fn_success,
+    )
 
     # Expect the status updates in order, and 1 for each job
     expected_completed_count = 0
-    async for progress in runner.run(jobs, mock_run_job_success):
+    async for progress in runner.run():
         assert progress.complete == expected_completed_count
         expected_completed_count += 1
         assert progress.errors == 0
@@ -36,26 +57,29 @@ async def test_async_job_runner_status_updates(concurrency):
     assert expected_completed_count == job_count + 1
 
     # Verify run_job was called for each job
-    assert mock_run_job_success.call_count == job_count
+    assert mock_async_run_job_fn_success.call_count == job_count
 
     # Verify run_job was called with the correct arguments
     for i in range(job_count):
-        mock_run_job_success.assert_any_await(jobs[i])
+        mock_async_run_job_fn_success.assert_any_await(jobs[i])
 
 
 # Test with and without concurrency
 @pytest.mark.parametrize("concurrency", [1, 25])
 @pytest.mark.asyncio
-async def test_async_job_runner_status_updates_empty_job_list(concurrency):
+async def test_async_job_runner_status_updates_empty_job_list(
+    concurrency, mock_async_run_job_fn_success
+):
     empty_job_list = []
 
-    runner = AsyncJobRunner(concurrency=concurrency)
-
-    # fake run_job that succeeds
-    mock_run_job_success = AsyncMock(return_value=True)
+    runner = AsyncJobRunner(
+        concurrency=concurrency,
+        jobs=empty_job_list,
+        run_job_fn=mock_async_run_job_fn_success,
+    )
 
     updates: List[Progress] = []
-    async for progress in runner.run(empty_job_list, mock_run_job_success):
+    async for progress in runner.run():
         updates.append(progress)
 
     # Verify last status update was complete
@@ -66,23 +90,26 @@ async def test_async_job_runner_status_updates_empty_job_list(concurrency):
     assert updates[0].total == 0
 
     # Verify run_job was called for each job
-    assert mock_run_job_success.call_count == 0
+    assert mock_async_run_job_fn_success.call_count == 0
 
 
 @pytest.mark.parametrize("concurrency", [1, 25])
 @pytest.mark.asyncio
-async def test_async_job_runner_all_failures(concurrency):
+async def test_async_job_runner_all_failures(
+    concurrency, mock_async_run_job_fn_failure
+):
     job_count = 50
     jobs = [{"id": i} for i in range(job_count)]
 
-    runner = AsyncJobRunner(concurrency=concurrency)
-
-    # fake run_job that fails
-    mock_run_job_failure = AsyncMock(return_value=False)
+    runner = AsyncJobRunner(
+        concurrency=concurrency,
+        jobs=jobs,
+        run_job_fn=mock_async_run_job_fn_failure,
+    )
 
     # Expect the status updates in order, and 1 for each job
     expected_error_count = 0
-    async for progress in runner.run(jobs, mock_run_job_failure):
+    async for progress in runner.run():
         assert progress.complete == 0
         assert progress.errors == expected_error_count
         expected_error_count += 1
@@ -92,11 +119,11 @@ async def test_async_job_runner_all_failures(concurrency):
     assert expected_error_count == job_count + 1
 
     # Verify run_job was called for each job
-    assert mock_run_job_failure.call_count == job_count
+    assert mock_async_run_job_fn_failure.call_count == job_count
 
     # Verify run_job was called with the correct arguments
     for i in range(job_count):
-        mock_run_job_failure.assert_any_await(jobs[i])
+        mock_async_run_job_fn_failure.assert_any_await(jobs[i])
 
 
 @pytest.mark.parametrize("concurrency", [1, 25])
@@ -108,16 +135,20 @@ async def test_async_job_runner_partial_failures(concurrency):
     # we want to fail on some jobs and succeed on others
     jobs_to_fail = set([0, 2, 4, 6, 8, 20, 25])
 
-    runner = AsyncJobRunner(concurrency=concurrency)
-
     # fake run_job that fails
     mock_run_job_partial_success = AsyncMock(
         # return True for jobs that should succeed
         side_effect=lambda job: job["id"] not in jobs_to_fail
     )
 
+    runner = AsyncJobRunner(
+        concurrency=concurrency,
+        jobs=jobs,
+        run_job_fn=mock_run_job_partial_success,
+    )
+
     # Expect the status updates in order, and 1 for each job
-    async for progress in runner.run(jobs, mock_run_job_partial_success):
+    async for progress in runner.run():
         assert progress.total == job_count
 
     # Verify last status update was complete
@@ -140,8 +171,6 @@ async def test_async_job_runner_partial_raises(concurrency):
     job_count = 50
     jobs = [{"id": i} for i in range(job_count)]
 
-    runner = AsyncJobRunner(concurrency=concurrency)
-
     ids_to_fail = set([10, 25])
 
     def failure_fn(job):
@@ -152,6 +181,12 @@ async def test_async_job_runner_partial_raises(concurrency):
     # fake run_job that fails
     mock_run_job_partial_success = AsyncMock(side_effect=failure_fn)
 
+    runner = AsyncJobRunner(
+        concurrency=concurrency,
+        jobs=jobs,
+        run_job_fn=mock_run_job_partial_success,
+    )
+
     # generate all the values we expect to see in progress updates
     complete_values_expected = set([i for i in range(job_count - len(ids_to_fail) + 1)])
     errors_values_expected = set([i for i in range(len(ids_to_fail) + 1)])
@@ -164,7 +199,7 @@ async def test_async_job_runner_partial_raises(concurrency):
     errors_values_actual = set()
 
     # Expect the status updates in order, and 1 for each job
-    async for progress in runner.run(jobs, mock_run_job_partial_success):
+    async for progress in runner.run():
         updates.append(progress)
         complete_values_actual.add(progress.complete)
         errors_values_actual.add(progress.errors)
@@ -184,9 +219,13 @@ async def test_async_job_runner_partial_raises(concurrency):
 
 @pytest.mark.parametrize("concurrency", [1, 25])
 @pytest.mark.asyncio
-async def test_async_job_runner_cancelled(concurrency):
-    runner = AsyncJobRunner(concurrency=concurrency)
+async def test_async_job_runner_cancelled(concurrency, mock_async_run_job_fn_success):
     jobs = [{"id": i} for i in range(10)]
+    runner = AsyncJobRunner(
+        concurrency=concurrency,
+        jobs=jobs,
+        run_job_fn=mock_async_run_job_fn_success,
+    )
 
     with patch.object(
         runner,
@@ -195,5 +234,82 @@ async def test_async_job_runner_cancelled(concurrency):
     ):
         # if an exception is raised in the task, we should see it bubble up
         with pytest.raises(Exception, match="run_worker raised an exception"):
-            async for _ in runner.run(jobs, AsyncMock(return_value=True)):
+            async for _ in runner.run():
                 pass
+
+
+@pytest.mark.parametrize("concurrency", [1, 25])
+@pytest.mark.asyncio
+async def test_async_job_runner_observers(concurrency):
+    class MockAsyncJobRunnerObserver(AsyncJobRunnerObserver[dict[str, int]]):
+        def __init__(self):
+            self.on_error_calls = []
+            self.on_success_calls = []
+
+        async def on_error(self, job: dict[str, int], error: Exception):
+            self.on_error_calls.append((job, error))
+
+        async def on_success(self, job: dict[str, int]):
+            self.on_success_calls.append(job)
+
+    mock_observer_a = MockAsyncJobRunnerObserver()
+    mock_observer_b = MockAsyncJobRunnerObserver()
+
+    jobs = [{"id": i} for i in range(10)]
+
+    async def run_job_fn(job: dict[str, int]) -> bool:
+        # we simulate the job 5 and 6 crashing, which should trigger the observers on_error handlers
+        if job["id"] == 5 or job["id"] == 6:
+            raise ValueError(f"job failed unexpectedly {job['id']}")
+        return True
+
+    runner = AsyncJobRunner(
+        concurrency=concurrency,
+        jobs=jobs,
+        run_job_fn=run_job_fn,
+        observers=[mock_observer_a, mock_observer_b],
+    )
+
+    async for _ in runner.run():
+        pass
+
+    assert len(mock_observer_a.on_error_calls) == 2
+    assert len(mock_observer_b.on_error_calls) == 2
+
+    # not necessarily in order, but we should have seen both 5 and 6
+    assert len(mock_observer_a.on_success_calls) == 8
+    assert len(mock_observer_b.on_success_calls) == 8
+
+    # check that 5 and 6 are in the error calls
+    for job_idx in [5, 6]:
+        # check that 5 and 6 are in the error calls for both observers
+        assert any(call[0] == jobs[job_idx] for call in mock_observer_a.on_error_calls)
+        assert any(call[0] == jobs[job_idx] for call in mock_observer_b.on_error_calls)
+
+        # check that the error is the correct exception
+        assert (
+            str(mock_observer_a.on_error_calls[0][1]) == "job failed unexpectedly 5"
+            or str(mock_observer_a.on_error_calls[1][1]) == "job failed unexpectedly 6"
+        )
+        assert (
+            str(mock_observer_b.on_error_calls[0][1]) == "job failed unexpectedly 5"
+            or str(mock_observer_b.on_error_calls[1][1]) == "job failed unexpectedly 6"
+        )
+
+        # check that 5 and 6 are not in the success calls for both observers
+        assert not any(
+            call == jobs[job_idx] for call in mock_observer_a.on_success_calls
+        )
+        assert not any(
+            call == jobs[job_idx] for call in mock_observer_b.on_success_calls
+        )
+
+    # check that the other jobs are in the success calls for both observers
+    for job_idx in range(10):
+        if job_idx not in [5, 6]:
+            assert any(
+                call == jobs[job_idx] for call in mock_observer_a.on_success_calls
+            )
+            assert any(
+                call == jobs[job_idx] for call in mock_observer_b.on_success_calls
+            )

kiln_ai/utils/test_config.py

@@ -6,7 +6,7 @@ from unittest.mock import patch
 import pytest
 import yaml
 
-from kiln_ai.utils.config import Config, ConfigProperty, _get_user_id
+from kiln_ai.utils.config import MCP_SECRETS_KEY, Config, ConfigProperty, _get_user_id
 
 
 @pytest.fixture
@@ -322,3 +322,140 @@ def test_update_settings_thread_safety(config_with_yaml):
 
     assert not exceptions
     assert config.int_property in range(5)
+
+
+def test_mcp_secrets_property():
+    """Test mcp_secrets configuration property"""
+    config = Config.shared()
+
+    # Initially should be None/empty
+    assert config.mcp_secrets is None
+
+    # Set some secrets
+    secrets = {
+        "server1::Authorization": "Bearer token123",
+        "server1::X-API-Key": "api-key-456",
+        "server2::Token": "secret-token",
+    }
+    config.mcp_secrets = secrets
+
+    # Verify they are stored correctly
+    assert config.mcp_secrets == secrets
+    assert config.mcp_secrets["server1::Authorization"] == "Bearer token123"
+    assert config.mcp_secrets["server1::X-API-Key"] == "api-key-456"
+    assert config.mcp_secrets["server2::Token"] == "secret-token"
+
+
+def test_mcp_secrets_sensitive_hiding():
+    """Test that mcp_secrets are hidden when hide_sensitive=True"""
+    config = Config.shared()
+
+    # Set some secrets
+    secrets = {
+        "server1::Authorization": "Bearer secret123",
+        "server2::X-API-Key": "secret-key",
+    }
+    config.mcp_secrets = secrets
+
+    # Test without hiding sensitive data
+    visible_settings = config.settings(hide_sensitive=False)
+    assert MCP_SECRETS_KEY in visible_settings
+    assert visible_settings[MCP_SECRETS_KEY] == secrets
+
+    # Test with hiding sensitive data
+    hidden_settings = config.settings(hide_sensitive=True)
+    assert MCP_SECRETS_KEY in hidden_settings
+    assert hidden_settings[MCP_SECRETS_KEY] == "[hidden]"
+
+
+def test_mcp_secrets_persistence(mock_yaml_file):
+    """Test that mcp_secrets are persisted to YAML correctly"""
+    with patch(
+        "kiln_ai.utils.config.Config.settings_path",
+        return_value=mock_yaml_file,
+    ):
+        config = Config()
+
+        # Set some secrets
+        secrets = {
+            "server1::Authorization": "Bearer persist123",
+            "server2::Token": "persist-token",
+        }
+        config.mcp_secrets = secrets
+
+        # Check that the value was saved to the YAML file
+        with open(mock_yaml_file, "r") as f:
+            saved_settings = yaml.safe_load(f)
+        assert saved_settings[MCP_SECRETS_KEY] == secrets
+
+        # Create a new config instance to test loading from YAML
+        new_config = Config()
+
+        # Check that the value is loaded from YAML
+        assert new_config.mcp_secrets == secrets
+
+
+def test_mcp_secrets_get_value():
+    """Test that mcp_secrets can be retrieved using get_value method"""
+    config = Config.shared()
+
+    # Initially should be None
+    assert config.get_value(MCP_SECRETS_KEY) is None
+
+    # Set some secrets
+    secrets = {"server::key": "value"}
+    config.mcp_secrets = secrets
+
+    # Should be retrievable via get_value
+    assert config.get_value(MCP_SECRETS_KEY) == secrets
+
+
+def test_mcp_secrets_update_settings():
+    """Test updating mcp_secrets using update_settings method"""
+    config = Config.shared()
+
+    # Set initial secrets
+    initial_secrets = {"server1::key1": "value1"}
+    config.update_settings({MCP_SECRETS_KEY: initial_secrets})
+    assert config.mcp_secrets == initial_secrets
+
+    # Update with new secrets (should replace, not merge)
+    new_secrets = {
+        "server1::key1": "updated_value1",
+        "server2::key2": "value2",
+    }
+    config.update_settings({MCP_SECRETS_KEY: new_secrets})
+    assert config.mcp_secrets == new_secrets
+    assert config.mcp_secrets["server1::key1"] == "updated_value1"
+    assert config.mcp_secrets["server2::key2"] == "value2"
+
+
+def test_mcp_secrets_empty_dict():
+    """Test mcp_secrets with empty dict"""
+    config = Config.shared()
+
+    # Set empty dict
+    config.mcp_secrets = {}
+    assert config.mcp_secrets == {}
+
+    # Should still be dict type, not None
+    assert isinstance(config.mcp_secrets, dict)
+
+
+def test_mcp_secrets_type_validation():
+    """Test that mcp_secrets enforces dict[str, str] type"""
+    config = Config.shared()
+
+    # Valid dict[str, str]
+    valid_secrets = {"server::key": "value"}
+    config.mcp_secrets = valid_secrets
+    assert config.mcp_secrets == valid_secrets
+
+    # The config system applies type conversion when retrieving values
+    mixed_types = {"server::key": 123}  # int value
+    config.mcp_secrets = mixed_types
+    # The type conversion happens when the value is retrieved, not when set
+    # So the underlying storage may preserve the original type
+    assert config.mcp_secrets == mixed_types or config.mcp_secrets == {
+        "server::key": "123"
+    }

kiln_ai/utils/test_env.py

@@ -0,0 +1,142 @@
+import os
+
+import pytest
+
+from kiln_ai.utils.env import temporary_env
+
+
+class TestTemporaryEnv:
+    def test_set_new_env_var(self):
+        """Test setting a new environment variable that doesn't exist."""
+        var_name = "TEST_NEW_VAR"
+        test_value = "test_value"
+
+        # Ensure the variable doesn't exist initially
+        assert var_name not in os.environ
+
+        with temporary_env(var_name, test_value):
+            assert os.environ[var_name] == test_value
+
+        # Verify it's removed after context
+        assert var_name not in os.environ
+
+    def test_modify_existing_env_var(self):
+        """Test modifying an existing environment variable."""
+        var_name = "TEST_EXISTING_VAR"
+        original_value = "original_value"
+        new_value = "new_value"
+
+        # Set up initial state
+        os.environ[var_name] = original_value
+
+        with temporary_env(var_name, new_value):
+            assert os.environ[var_name] == new_value
+
+        # Verify original value is restored
+        assert os.environ[var_name] == original_value
+
+    def test_restore_nonexistent_var(self):
+        """Test that a variable that didn't exist is properly removed."""
+        var_name = "TEST_NONEXISTENT_VAR"
+        test_value = "test_value"
+
+        # Ensure the variable doesn't exist initially
+        if var_name in os.environ:
+            del os.environ[var_name]
+
+        with temporary_env(var_name, test_value):
+            assert os.environ[var_name] == test_value
+
+        # Verify it's removed after context
+        assert var_name not in os.environ
+
+    def test_exception_handling(self):
+        """Test that environment is restored even when an exception occurs."""
+        var_name = "TEST_EXCEPTION_VAR"
+        original_value = "original_value"
+        new_value = "new_value"
+
+        # Set up initial state
+        os.environ[var_name] = original_value
+
+        with pytest.raises(ValueError):
+            with temporary_env(var_name, new_value):
+                assert os.environ[var_name] == new_value
+                raise ValueError("Test exception")
+
+        # Verify original value is restored even after exception
+        assert os.environ[var_name] == original_value
+
+    def test_exception_handling_new_var(self):
+        """Test that new variable is removed even when an exception occurs."""
+        var_name = "TEST_EXCEPTION_NEW_VAR"
+        test_value = "test_value"
+
+        # Ensure the variable doesn't exist initially
+        if var_name in os.environ:
+            del os.environ[var_name]
+
+        with pytest.raises(RuntimeError):
+            with temporary_env(var_name, test_value):
+                assert os.environ[var_name] == test_value
+                raise RuntimeError("Test exception")
+
+        # Verify variable is removed even after exception
+        assert var_name not in os.environ
+
+    def test_nested_context_managers(self):
+        """Test using multiple temporary_env context managers."""
+        var1 = "TEST_NESTED_VAR1"
+        var2 = "TEST_NESTED_VAR2"
+        value1 = "value1"
+        value2 = "value2"
+
+        # Set up initial state
+        os.environ[var1] = "original1"
+        if var2 in os.environ:
+            del os.environ[var2]
+
+        with temporary_env(var1, value1):
+            assert os.environ[var1] == value1
+
+            with temporary_env(var2, value2):
+                assert os.environ[var1] == value1
+                assert os.environ[var2] == value2
+
+            # Inner context should be cleaned up
+            assert var2 not in os.environ
+            assert os.environ[var1] == value1
+
+        # Both contexts should be cleaned up
+        assert os.environ[var1] == "original1"
+        assert var2 not in os.environ
+
+    def test_empty_string_value(self):
+        """Test setting an empty string value."""
+        var_name = "TEST_EMPTY_VAR"
+        test_value = ""
+
+        with temporary_env(var_name, test_value):
+            assert os.environ[var_name] == test_value
+
+        assert var_name not in os.environ
+
+    def test_none_value_handling(self):
+        """Test that None values are handled properly."""
+        var_name = "TEST_NONE_VAR"
+        test_value = "test_value"
+
+        with temporary_env(var_name, test_value):
+            assert os.environ[var_name] == test_value
+
+        assert var_name not in os.environ
+
+    def test_unicode_value(self):
+        """Test setting unicode values."""
+        var_name = "TEST_UNICODE_VAR"
+        test_value = "测试值 🚀"
+
+        with temporary_env(var_name, test_value):
+            assert os.environ[var_name] == test_value
+
+        assert var_name not in os.environ