kiln-ai 0.20.1__py3-none-any.whl → 0.22.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_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