"""Tests for session context manager."""
from __future__ import annotations

import asyncio

import staso as st
from staso.context import session_id_var


class TestSession:
    def test_session_sets_id(self):
        with st.conversation("my-session"):
            assert session_id_var.get("") == "my-session"
        assert session_id_var.get("") == ""

    def test_session_propagates_to_spans(self, capturing):
        @st.agent(name="a")
        def run():
            return "ok"

        with st.conversation("sess-123"):
            run()

        assert capturing[0]["session_id"] == "sess-123"

    def test_no_session_empty_id(self, capturing):
        @st.agent(name="a")
        def run():
            return "ok"

        run()
        assert capturing[0]["session_id"] == ""

    def test_nested_sessions(self):
        with st.conversation("outer"):
            assert session_id_var.get("") == "outer"
            with st.conversation("inner"):
                assert session_id_var.get("") == "inner"
            assert session_id_var.get("") == "outer"
        assert session_id_var.get("") == ""

    def test_session_restored_on_exception(self):
        try:
            with st.conversation("test"):
                raise ValueError("oops")
        except ValueError:
            pass
        assert session_id_var.get("") == ""

    def test_session_propagates_to_nested_tools(self, capturing):
        @st.tool(name="inner")
        def inner():
            return "ok"

        @st.agent(name="outer")
        def outer():
            return inner()

        with st.conversation("propagated"):
            outer()

        for span in capturing:
            assert span["session_id"] == "propagated"

    def test_session_works_with_async(self, capturing):
        @st.agent(name="async-agent")
        async def run():
            return "ok"

        with st.conversation("async-sess"):
            asyncio.run(run())

        assert capturing[0]["session_id"] == "async-sess"
