kiln_ai.test_run_context

  1"""Tests for agent run context management using contextvars."""
  2
  3import asyncio
  4
  5import pytest
  6
  7from kiln_ai.run_context import (
  8    clear_agent_run_id,
  9    generate_agent_run_id,
 10    get_agent_run_id,
 11    set_agent_run_id,
 12)
 13
 14
 15class TestAdapterRunContext:
 16    """Unit tests for agent run context management."""
 17
 18    def test_default_is_none(self):
 19        """Test that the default run ID is None."""
 20        # Clear any existing context first
 21        clear_agent_run_id()
 22        assert get_agent_run_id() is None
 23
 24    def test_set_and_get(self):
 25        """Test basic set and get cycle."""
 26        clear_agent_run_id()
 27        assert get_agent_run_id() is None
 28
 29        test_run_id = "test_run_123"
 30        set_agent_run_id(test_run_id)
 31        assert get_agent_run_id() == test_run_id
 32
 33    def test_clear(self):
 34        """Test that clear resets to None."""
 35        set_agent_run_id("some_run")
 36        assert get_agent_run_id() == "some_run"
 37
 38        clear_agent_run_id()
 39        assert get_agent_run_id() is None
 40
 41    def test_generate_agent_run_id_format(self):
 42        """Test that generated run IDs have the correct format."""
 43        run_id = generate_agent_run_id()
 44        assert run_id.startswith("run_")
 45        assert len(run_id) == 20  # "run_" + 16 hex characters
 46
 47    def test_generate_agent_run_id_unique(self):
 48        """Test that generated run IDs are unique."""
 49        ids = [generate_agent_run_id() for _ in range(100)]
 50        assert len(set(ids)) == 100  # All should be unique
 51
 52    @pytest.mark.asyncio
 53    async def test_asyncio_gather_propagation(self):
 54        """Test that contextvar propagates through asyncio.gather."""
 55        clear_agent_run_id()
 56        parent_run_id = "parent_run"
 57        set_agent_run_id(parent_run_id)
 58
 59        child_results = await asyncio.gather(
 60            self._read_context_var(),
 61            self._read_context_var(),
 62            self._read_context_var(),
 63        )
 64
 65        # All children should see the parent's run ID
 66        for result in child_results:
 67            assert result == parent_run_id
 68
 69    @pytest.mark.asyncio
 70    async def test_asyncio_gather_isolation(self):
 71        """Test that child writes don't affect parent or siblings."""
 72        clear_agent_run_id()
 73        parent_run_id = "parent_run"
 74        set_agent_run_id(parent_run_id)
 75
 76        async def child_modify_context():
 77            # Child sets a different value
 78            set_agent_run_id("child_modified")
 79            return get_agent_run_id()
 80
 81        child_results = await asyncio.gather(
 82            child_modify_context(),
 83            child_modify_context(),
 84        )
 85
 86        # Children should see their own modifications
 87        for result in child_results:
 88            assert result == "child_modified"
 89
 90        # Parent should still have the original value
 91        assert get_agent_run_id() == parent_run_id
 92
 93    @pytest.mark.asyncio
 94    async def test_nested_async_calls_propagate(self):
 95        """Test that contextvar propagates through nested async calls."""
 96        clear_agent_run_id()
 97        root_run_id = "root_run"
 98        set_agent_run_id(root_run_id)
 99
100        async def level_2():
101            return get_agent_run_id()
102
103        async def level_1():
104            return await level_2()
105
106        result = await level_1()
107        assert result == root_run_id
108
109    async def _read_context_var(self) -> str | None:
110        """Helper to read the context var in an async context."""
111        return get_agent_run_id()