Add tests for @ai_track decorator (#3325)

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"