Update tests

traceloop · Jun 7, 2024 · b5f7825 · b5f7825
1 parent d01b9b5
commit b5f7825
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/packages/opentelemetry-instrumentation-haystack/tests/conftest.py b/packages/opentelemetry-instrumentation-haystack/tests/conftest.py
@@ -28,6 +28,7 @@ def exporter():
 @pytest.fixture(autouse=True)
 def environment():
     os.environ["OPENAI_API_KEY"] = "test_api_key"
+    os.environ["TRACELOOP_TRACE_CONTENT"] = "true"
 
 
 @pytest.fixture(autouse=True)

diff --git a/packages/opentelemetry-instrumentation-haystack/tests/test_simple_pipeline.py b/packages/opentelemetry-instrumentation-haystack/tests/test_simple_pipeline.py
@@ -5,6 +5,7 @@
 from haystack.components.builders import DynamicChatPromptBuilder
 from haystack.dataclasses import ChatMessage
 from haystack.utils import Secret
+from opentelemetry.semconv.ai import SpanAttributes
 
 
 @pytest.mark.vcr
@@ -34,3 +35,6 @@ def test_haystack(exporter):
         "haystack.openai.chat",
         "haystack_pipeline.workflow",
     } == {span.name for span in spans}
+    span_workflow = next(span for span in spans if span.name == "haystack_pipeline.workflow")
+    assert SpanAttributes.TRACELOOP_ENTITY_INPUT in span_workflow.attributes
+    assert SpanAttributes.TRACELOOP_ENTITY_OUTPUT in span_workflow.attributes