From 5a75eda3c7fb2e76ad73c83b863fbdf6986b9c84 Mon Sep 17 00:00:00 2001
From: Teo <teocns@gmail.com>
Date: Tue, 26 Nov 2024 17:13:20 -0600
Subject: [PATCH] tests(session): migrate span processor tests to exporter
 tests

---
 tests/session/test_exporter.py | 47 ++++++++++++++++++++++++++++++----
 tests/test_session.py          | 21 ---------------
 2 files changed, 42 insertions(+), 26 deletions(-)

diff --git a/tests/session/test_exporter.py b/tests/session/test_exporter.py
index ec978773..68d48ca7 100644
--- a/tests/session/test_exporter.py
+++ b/tests/session/test_exporter.py
@@ -1,14 +1,15 @@
-import pytest
 from unittest.mock import Mock, patch
 from uuid import uuid4
-from opentelemetry.sdk.trace.export import SpanExportResult
-from opentelemetry.trace import SpanKind
+
+import pytest
 from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExportResult
+from opentelemetry.trace import SpanKind
 
-from agentops.session.exporter import SessionExporter
-from agentops.event import Event, ErrorEvent
 from agentops.config import Configuration
+from agentops.event import ErrorEvent, Event
 from agentops.session import Session
+from agentops.session.exporter import SessionExporter
 
 
 @pytest.fixture
@@ -121,3 +122,39 @@ def test_export_error_handling(exporter, mock_session):
 
     result = exporter.export([span])
     assert result == SpanExportResult.FAILURE
+
+
+def test_span_processor_config(mock_session):
+    """Test that BatchSpanProcessor is configured correctly for testing"""
+    # Create a mock BatchSpanProcessor instead of using real Session
+    mock_processor = Mock(spec=BatchSpanProcessor)
+    mock_processor._max_export_batch_size = 1
+    mock_processor._schedule_delay_millis = 0
+
+    with patch("agentops.session.Session._span_processor", mock_processor):
+        session = Session(mock_session.config)
+        # Verify BatchSpanProcessor configuration
+        assert session._span_processor._max_export_batch_size == 1
+        assert session._span_processor._schedule_delay_millis == 0
+
+
+def test_event_batching(mock_session):
+    """Test that events are properly batched for export"""
+    mock_processor = Mock(spec=BatchSpanProcessor)
+    mock_processor.on_end = Mock()
+
+    with patch("agentops.session.Session._span_processor", mock_processor), patch(
+        "agentops.session.Session._tracer_provider"
+    ) as mock_provider:
+        session = Session(mock_session.config)
+        session._span_processor = mock_processor  # Explicitly set the processor
+
+        # Create mock events
+        mock_events = [Mock(spec=Event) for _ in range(3)]
+
+        # Record events
+        for event in mock_events:
+            session.record(event)
+
+        # Verify batching
+        assert mock_processor.on_end.call_count == 3
diff --git a/tests/test_session.py b/tests/test_session.py
index e97cb649..e694959f 100644
--- a/tests/test_session.py
+++ b/tests/test_session.py
@@ -276,27 +276,6 @@ def test_get_analytics(self, mock_req):
         session.end_session(end_state="Success")
         agentops.end_all_sessions()
 
-    def test_span_processor_config(self):
-        session = agentops.start_session()
-
-        # Verify BatchSpanProcessor is configured for immediate export in tests
-        processor = session._span_processor
-        assert processor._max_export_batch_size == 1
-        assert processor._schedule_delay_millis == 0
-
-    def test_event_batching(self):
-        with patch("agentops.session.exporter.BatchSpanProcessor") as mock_processor:
-            session = agentops.start_session()
-
-            # Record multiple events
-            events = [ActionEvent(f"event_{i}") for i in range(3)]
-            for event in events:
-                session.record(event)
-
-            # Verify events were batched correctly
-            mock_processor.return_value.on_end.assert_called()
-            assert len(mock_processor.return_value.on_end.call_args[0][0]) == 3
-
 
 class TestMultiSessions:
     def setup_method(self):