From b5f78255fee4bca5c5aa3b339c9181bc6603492d Mon Sep 17 00:00:00 2001
From: Tibor Reiss <tibor.reiss@gmail.com>
Date: Fri, 7 Jun 2024 19:43:24 +0200
Subject: [PATCH] Update tests

---
 .../opentelemetry-instrumentation-haystack/tests/conftest.py  | 1 +
 .../tests/test_simple_pipeline.py                             | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/packages/opentelemetry-instrumentation-haystack/tests/conftest.py b/packages/opentelemetry-instrumentation-haystack/tests/conftest.py
index 6e6a152bb..e550a8088 100644
--- 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
index 5971db015..c1c546b2e 100644
--- 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