import time
import unittest.mock
from typing import Any
from uuid import UUID

from langchainplus_sdk import LangChainPlusClient

from langchain.callbacks.tracers.langchain import LangChainTracer
from langchain.callbacks.tracers.schemas import Run
from langchain.schema.output import LLMResult


def test_example_id_assignment_threadsafe() -> None:
    """Test that example assigned at callback start/end is honored."""
    example_ids = {}

    def mock_create_run(self: Any, **kwargs: Any) -> Any:
        example_ids[kwargs.get("id")] = kwargs.get("reference_example_id")
        return unittest.mock.MagicMock()

    with unittest.mock.patch.object(
        LangChainPlusClient, "create_run", new=mock_create_run
    ):
        client = LangChainPlusClient()
        tracer = LangChainTracer(client=client)
        old_persist_run_single = tracer._persist_run_single

        def new_persist_run_single(run: Run) -> None:
            time.sleep(0.01)
            old_persist_run_single(run)

        with unittest.mock.patch.object(
            tracer, "_persist_run_single", new=new_persist_run_single
        ):
            run_id_1 = UUID("9d878ab3-e5ca-4218-aef6-44cbdc90160a")
            run_id_2 = UUID("f1f9fa53-8b2f-4742-bdbc-38215f7bd1e1")
            example_id_1 = UUID("57e42c57-8c79-4d9f-8765-bf6cd3a98055")
            tracer.example_id = example_id_1
            tracer.on_llm_start({"name": "example_1"}, ["foo"], run_id=run_id_1)
            tracer.on_llm_end(LLMResult(generations=[], llm_output={}), run_id=run_id_1)
            example_id_2 = UUID("4f31216e-7c26-4027-a5fd-0bbf9ace17dc")
            tracer.example_id = example_id_2
            tracer.on_llm_start({"name": "example_2"}, ["foo"], run_id=run_id_2)
            tracer.on_llm_end(LLMResult(generations=[], llm_output={}), run_id=run_id_2)
            tracer.example_id = None
            expected_example_ids = {
                run_id_1: example_id_1,
                run_id_2: example_id_2,
            }
            tracer.wait_for_futures()
            assert example_ids == expected_example_ids

langchain.callbacks.tracers.langchain.LangChainTracer langchainplus_sdk.LangChainPlusClient langchain.schema.output.LLMResult
[((741, 762), 'langchainplus_sdk.LangChainPlusClient', 'LangChainPlusClient', ([], {}), '()\n', (760, 762), False, 'from langchainplus_sdk import LangChainPlusClient\n'), ((780, 810), 'langchain.callbacks.tracers.langchain.LangChainTracer', 'LangChainTracer', ([], {'client': 'client'}), '(client=client)\n', (795, 810), False, 'from langchain.callbacks.tracers.langchain import LangChainTracer\n'), ((938, 954), 'time.sleep', 'time.sleep', (['(0.01)'], {}), '(0.01)\n', (948, 954), False, 'import time\n'), ((1141, 1185), 'uuid.UUID', 'UUID', (['"""9d878ab3-e5ca-4218-aef6-44cbdc90160a"""'], {}), "('9d878ab3-e5ca-4218-aef6-44cbdc90160a')\n", (1145, 1185), False, 'from uuid import UUID\n'), ((1209, 1253), 'uuid.UUID', 'UUID', (['"""f1f9fa53-8b2f-4742-bdbc-38215f7bd1e1"""'], {}), "('f1f9fa53-8b2f-4742-bdbc-38215f7bd1e1')\n", (1213, 1253), False, 'from uuid import UUID\n'), ((1281, 1325), 'uuid.UUID', 'UUID', (['"""57e42c57-8c79-4d9f-8765-bf6cd3a98055"""'], {}), "('57e42c57-8c79-4d9f-8765-bf6cd3a98055')\n", (1285, 1325), False, 'from uuid import UUID\n'), ((1568, 1612), 'uuid.UUID', 'UUID', (['"""4f31216e-7c26-4027-a5fd-0bbf9ace17dc"""'], {}), "('4f31216e-7c26-4027-a5fd-0bbf9ace17dc')\n", (1572, 1612), False, 'from uuid import UUID\n'), ((1482, 1522), 'langchain.schema.output.LLMResult', 'LLMResult', ([], {'generations': '[]', 'llm_output': '{}'}), '(generations=[], llm_output={})\n', (1491, 1522), False, 'from langchain.schema.output import LLMResult\n'), ((1769, 1809), 'langchain.schema.output.LLMResult', 'LLMResult', ([], {'generations': '[]', 'llm_output': '{}'}), '(generations=[], llm_output={})\n', (1778, 1809), False, 'from langchain.schema.output import LLMResult\n')]