@@ -42,29 +42,34 @@ def _make_mock_sgp_span() -> MagicMock:
4242
4343
4444class TestSGPSyncTracingProcessorMemoryLeak :
45- @patch (f"{ MODULE } .EnvironmentVariables" )
46- @patch (f"{ MODULE } .flush_queue" )
47- @patch (f"{ MODULE } .create_span" )
48- @patch (f"{ MODULE } .SGPClient" )
49- @patch (f"{ MODULE } .tracing" )
50- def _make_processor (self , mock_tracing , mock_sgp_client , mock_create_span , mock_flush , mock_env ):
45+ @staticmethod
46+ def _make_processor ():
47+ mock_env = MagicMock ()
5148 mock_env .refresh .return_value = MagicMock (ACP_TYPE = None , AGENT_NAME = None , AGENT_ID = None )
52- mock_create_span . side_effect = lambda ** kwargs : _make_mock_sgp_span ()
49+ mock_create_span = MagicMock ( side_effect = lambda ** kwargs : _make_mock_sgp_span () )
5350
54- from agentex .lib .core .tracing .processors .sgp_tracing_processor import (
55- SGPSyncTracingProcessor ,
56- )
51+ with patch (f"{ MODULE } .EnvironmentVariables" , mock_env ), \
52+ patch (f"{ MODULE } .SGPClient" ), \
53+ patch (f"{ MODULE } .tracing" ), \
54+ patch (f"{ MODULE } .flush_queue" ), \
55+ patch (f"{ MODULE } .create_span" , mock_create_span ):
56+ from agentex .lib .core .tracing .processors .sgp_tracing_processor import (
57+ SGPSyncTracingProcessor ,
58+ )
59+
60+ processor = SGPSyncTracingProcessor (_make_config ())
5761
58- return SGPSyncTracingProcessor ( _make_config ()) , mock_create_span
62+ return processor , mock_create_span
5963
6064 def test_spans_not_leaked_after_completed_lifecycle (self ):
6165 processor , _ = self ._make_processor ()
6266
63- for _ in range (100 ):
64- span = _make_span ()
65- processor .on_span_start (span )
66- span .end_time = datetime .now (UTC )
67- processor .on_span_end (span )
67+ with patch (f"{ MODULE } .create_span" , side_effect = lambda ** kw : _make_mock_sgp_span ()):
68+ for _ in range (100 ):
69+ span = _make_span ()
70+ processor .on_span_start (span )
71+ span .end_time = datetime .now (UTC )
72+ processor .on_span_end (span )
6873
6974 assert len (processor ._spans ) == 0 , (
7075 f"Expected 0 spans after 100 complete lifecycles, got { len (processor ._spans )} — memory leak!"
@@ -73,13 +78,14 @@ def test_spans_not_leaked_after_completed_lifecycle(self):
7378 def test_spans_present_during_active_lifecycle (self ):
7479 processor , _ = self ._make_processor ()
7580
76- span = _make_span ()
77- processor .on_span_start (span )
78- assert len (processor ._spans ) == 1 , "Span should be tracked while active"
81+ with patch (f"{ MODULE } .create_span" , side_effect = lambda ** kw : _make_mock_sgp_span ()):
82+ span = _make_span ()
83+ processor .on_span_start (span )
84+ assert len (processor ._spans ) == 1 , "Span should be tracked while active"
7985
80- span .end_time = datetime .now (UTC )
81- processor .on_span_end (span )
82- assert len (processor ._spans ) == 0 , "Span should be removed after end"
86+ span .end_time = datetime .now (UTC )
87+ processor .on_span_end (span )
88+ assert len (processor ._spans ) == 0 , "Span should be removed after end"
8389
8490 def test_span_end_for_unknown_span_is_noop (self ):
8591 processor , _ = self ._make_processor ()
0 commit comments