From 7fa6d2219d4f81862431fc16d5087bcc86ff1282 Mon Sep 17 00:00:00 2001
From: Colin Chartier <colin.chartier@sentry.io>
Date: Mon, 22 Jul 2024 15:43:42 -0400
Subject: [PATCH] Add tests for @ai_track decorator

---
 tests/test_ai_monitoring.py | 59 +++++++++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)
 create mode 100644 tests/test_ai_monitoring.py

diff --git a/tests/test_ai_monitoring.py b/tests/test_ai_monitoring.py
new file mode 100644
index 0000000000..4329cc92af
--- /dev/null
+++ b/tests/test_ai_monitoring.py
@@ -0,0 +1,59 @@
+import sentry_sdk
+from sentry_sdk.ai.monitoring import ai_track
+
+
+def test_ai_track(sentry_init, capture_events):
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events()
+
+    @ai_track("my tool")
+    def tool(**kwargs):
+        pass
+
+    @ai_track("some test pipeline")
+    def pipeline():
+        tool()
+
+    with sentry_sdk.start_transaction():
+        pipeline()
+
+    transaction = events[0]
+    assert transaction["type"] == "transaction"
+    assert len(transaction["spans"]) == 2
+    spans = transaction["spans"]
+
+    ai_pipeline_span = spans[0] if spans[0]["op"] == "ai.pipeline" else spans[1]
+    ai_run_span = spans[0] if spans[0]["op"] == "ai.run" else spans[1]
+
+    assert ai_pipeline_span["description"] == "some test pipeline"
+    assert ai_run_span["description"] == "my tool"
+
+
+def test_ai_track_with_tags(sentry_init, capture_events):
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events()
+
+    @ai_track("my tool")
+    def tool(**kwargs):
+        pass
+
+    @ai_track("some test pipeline")
+    def pipeline():
+        tool()
+
+    with sentry_sdk.start_transaction():
+        pipeline(sentry_tags={"user": "colin"}, sentry_data={"some_data": "value"})
+
+    transaction = events[0]
+    assert transaction["type"] == "transaction"
+    assert len(transaction["spans"]) == 2
+    spans = transaction["spans"]
+
+    ai_pipeline_span = spans[0] if spans[0]["op"] == "ai.pipeline" else spans[1]
+    ai_run_span = spans[0] if spans[0]["op"] == "ai.run" else spans[1]
+
+    assert ai_pipeline_span["description"] == "some test pipeline"
+    print(ai_pipeline_span)
+    assert ai_pipeline_span["tags"]["user"] == "colin"
+    assert ai_pipeline_span["data"]["some_data"] == "value"
+    assert ai_run_span["description"] == "my tool"