# import pytest
# from langchain_openai import ChatOpenAI
# from langchain.prompts import ChatPromptTemplate
# from opentelemetry.semconv_ai import SpanAttributes
# from lmnr.traceloop_sdk import Traceloop
# from lmnr.traceloop_sdk.decorators import task, workflow


# def test_association_properties(exporter):
#     @workflow(name="test_workflow")
#     def test_workflow():
#         return test_task()

#     @task(name="test_task")
#     def test_task():
#         return

#     Traceloop.set_association_properties({"user_id": 1, "user_name": "John Doe"})
#     test_workflow()

#     spans = exporter.get_finished_spans()
#     assert [span.name for span in spans] == [
#         "test_task.task",
#         "test_workflow.workflow",
#     ]

#     some_task_span = spans[0]
#     some_workflow_span = spans[1]
#     assert (
#         some_workflow_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
#         ]
#         == 1
#     )
#     assert (
#         some_workflow_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_name"
#         ]
#         == "John Doe"
#     )
#     assert (
#         some_task_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
#         ]
#         == 1
#     )
#     assert (
#         some_task_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_name"
#         ]
#         == "John Doe"
#     )


# def test_association_properties_within_workflow(exporter):
#     @workflow(name="test_workflow_within")
#     def test_workflow():
#         Traceloop.set_association_properties({"session_id": 15})
#         return

#     test_workflow()

#     spans = exporter.get_finished_spans()
#     assert [span.name for span in spans] == [
#         "test_workflow_within.workflow",
#     ]

#     some_workflow_span = spans[0]
#     assert (
#         some_workflow_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id"
#         ]
#         == 15
#     )


# @pytest.mark.vcr
# def test_langchain_association_properties(exporter):
#     prompt = ChatPromptTemplate.from_messages(
#         [("system", "You are helpful assistant"), ("user", "{input}")]
#     )
#     model = ChatOpenAI(model="gpt-3.5-turbo")

#     chain = prompt | model
#     chain.invoke(
#         {"input": "tell me a short joke"},
#         {"metadata": {"user_id": "1234", "session_id": 456}},
#     )

#     spans = exporter.get_finished_spans()

#     assert [
#         "ChatPromptTemplate.task",
#         "ChatOpenAI.chat",
#         "RunnableSequence.workflow",
#     ] == [span.name for span in spans], [span.name for span in spans]

#     workflow_span = next(
#         span for span in spans if span.name == "RunnableSequence.workflow"
#     )
#     prompt_span = next(span for span in spans if span.name == "ChatPromptTemplate.task")
#     chat_span = next(span for span in spans if span.name == "ChatOpenAI.chat")

#     assert (
#         workflow_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
#         ]
#         == "1234"
#     )
#     assert (
#         workflow_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id"
#         ]
#         == 456
#     )
#     assert (
#         chat_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
#         ]
#         == "1234"
#     )
#     assert (
#         chat_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id"
#         ]
#         == 456
#     )
#     assert (
#         prompt_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
#         ]
#         == "1234"
#     )
#     assert (
#         prompt_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id"
#         ]
#         == 456
#     )


# @pytest.mark.vcr
# def test_langchain_and_external_association_properties(exporter):
#     @workflow(name="test_workflow_external")
#     def test_workflow_external():
#         Traceloop.set_association_properties({"workspace_id": "789"})

#         prompt = ChatPromptTemplate.from_messages(
#             [("system", "You are helpful assistant"), ("user", "{input}")]
#         )
#         model = ChatOpenAI(model="gpt-3.5-turbo")

#         chain = prompt | model
#         chain.invoke(
#             {"input": "tell me a short joke"},
#             {"metadata": {"user_id": "1234", "session_id": 456}},
#         )

#     test_workflow_external()

#     spans = exporter.get_finished_spans()

#     assert [
#         "ChatPromptTemplate.task",
#         "ChatOpenAI.chat",
#         "RunnableSequence.workflow",
#         "test_workflow_external.workflow",
#     ] == [span.name for span in spans], [span.name for span in spans]

#     workflow_span = next(
#         span for span in spans if span.name == "RunnableSequence.workflow"
#     )
#     prompt_span = next(span for span in spans if span.name == "ChatPromptTemplate.task")
#     chat_span = next(span for span in spans if span.name == "ChatOpenAI.chat")

#     assert (
#         workflow_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
#         ]
#         == "1234"
#     )
#     assert (
#         workflow_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id"
#         ]
#         == 456
#     )
#     assert (
#         workflow_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.workspace_id"
#         ]
#         == "789"
#     )
#     assert (
#         chat_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
#         ]
#         == "1234"
#     )
#     assert (
#         chat_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id"
#         ]
#         == 456
#     )
#     assert (
#         chat_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.workspace_id"
#         ]
#         == "789"
#     )
#     assert (
#         prompt_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
#         ]
#         == "1234"
#     )
#     assert (
#         prompt_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id"
#         ]
#         == 456
#     )
#     assert (
#         prompt_span.attributes[
#             f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.workspace_id"
#         ]
#         == "789"
#     )
