From cc16608f9e69299dcde48d886249926c84ea63a8 Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Tue, 1 Oct 2024 09:55:11 +0200 Subject: [PATCH] Add unit test --- integrations/langfuse/tests/test_tracing.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/integrations/langfuse/tests/test_tracing.py b/integrations/langfuse/tests/test_tracing.py index 936064e0a..657b6eae1 100644 --- a/integrations/langfuse/tests/test_tracing.py +++ b/integrations/langfuse/tests/test_tracing.py @@ -43,7 +43,12 @@ def test_tracing_integration(llm_class, env_var, expected_trace): ChatMessage.from_user("Tell me about {{location}}"), ] - response = pipe.run(data={"prompt_builder": {"template_variables": {"location": "Berlin"}, "template": messages}}) + response = pipe.run( + data={ + "prompt_builder": {"template_variables": {"location": "Berlin"}, "template": messages}, + "tracer": {"invocation_context": {"user_id": "user_42"}}, + } + ) assert "Berlin" in response["llm"]["replies"][0].content assert response["tracer"]["trace_url"] @@ -65,5 +70,7 @@ def test_tracing_integration(llm_class, env_var, expected_trace): assert expected_trace in str(response.content) # check if the trace contains the expected generation span assert "GENERATION" in str(response.content) + # check if the trace contains the expected user_id + assert "user_42" in str(response.content) except requests.exceptions.RequestException as e: pytest.fail(f"Failed to retrieve data from Langfuse API: {e}")