@@ -276,17 +276,21 @@ def test_add_trace_creates_child_span(self, tracer, mock_logger):
276276 def test_add_trace_no_op_when_not_ready (self , trace_id ):
277277 with patch .dict ("os.environ" , {}, clear = True ):
278278 t = BraintrustTracer (
279- trace_name = "Test" , trace_type = "chain" ,
280- project_name = "P" , trace_id = trace_id ,
279+ trace_name = "Test" ,
280+ trace_type = "chain" ,
281+ project_name = "P" ,
282+ trace_id = trace_id ,
281283 )
282284 t .add_trace (trace_id = "c1" , trace_name = "X" , trace_type = "llm" , inputs = {})
283285 assert len (t .spans ) == 0
284286
285287 def test_end_trace_logs_and_closes_span (self , tracer , mock_logger ):
286288 _logger , _root_span , child_span = mock_logger
287289 tracer .add_trace (
288- trace_id = "comp-1" , trace_name = "OpenAI (comp-1)" ,
289- trace_type = "llm" , inputs = {"prompt" : "hello" },
290+ trace_id = "comp-1" ,
291+ trace_name = "OpenAI (comp-1)" ,
292+ trace_type = "llm" ,
293+ inputs = {"prompt" : "hello" },
290294 )
291295 tracer .end_trace (
292296 trace_id = "comp-1" ,
@@ -302,24 +306,30 @@ def test_end_trace_logs_and_closes_span(self, tracer, mock_logger):
302306 def test_end_trace_with_error (self , tracer , mock_logger ):
303307 _logger , _root_span , child_span = mock_logger
304308 tracer .add_trace (
305- trace_id = "comp-1" , trace_name = "Fail (comp-1)" ,
306- trace_type = "llm" , inputs = {},
309+ trace_id = "comp-1" ,
310+ trace_name = "Fail (comp-1)" ,
311+ trace_type = "llm" ,
312+ inputs = {},
307313 )
308314 error = ValueError ("something went wrong" )
309315 tracer .end_trace (
310- trace_id = "comp-1" , trace_name = "Fail (comp-1)" ,
316+ trace_id = "comp-1" ,
317+ trace_name = "Fail (comp-1)" ,
311318 error = error ,
312319 )
313320 assert child_span .log .call_args [1 ]["error" ] == "something went wrong"
314321
315322 def test_end_trace_with_logs (self , tracer , mock_logger ):
316323 _logger , _root_span , child_span = mock_logger
317324 tracer .add_trace (
318- trace_id = "comp-1" , trace_name = "X (comp-1)" ,
319- trace_type = "llm" , inputs = {},
325+ trace_id = "comp-1" ,
326+ trace_name = "X (comp-1)" ,
327+ trace_type = "llm" ,
328+ inputs = {},
320329 )
321330 tracer .end_trace (
322- trace_id = "comp-1" , trace_name = "X (comp-1)" ,
331+ trace_id = "comp-1" ,
332+ trace_name = "X (comp-1)" ,
323333 outputs = {"out" : "value" },
324334 logs = [{"step" : 1 }, {"step" : 2 }],
325335 )
@@ -358,8 +368,10 @@ def test_end_with_error(self, tracer, mock_logger):
358368 def test_end_no_op_when_not_ready (self , trace_id ):
359369 with patch .dict ("os.environ" , {}, clear = True ):
360370 t = BraintrustTracer (
361- trace_name = "Test" , trace_type = "chain" ,
362- project_name = "P" , trace_id = trace_id ,
371+ trace_name = "Test" ,
372+ trace_type = "chain" ,
373+ project_name = "P" ,
374+ trace_id = trace_id ,
363375 )
364376 # Should not raise
365377 t .end (inputs = {}, outputs = {})
@@ -376,8 +388,10 @@ class TestGetLangchainCallback:
376388 def test_returns_none_when_not_ready (self , trace_id ):
377389 with patch .dict ("os.environ" , {}, clear = True ):
378390 t = BraintrustTracer (
379- trace_name = "Test" , trace_type = "chain" ,
380- project_name = "P" , trace_id = trace_id ,
391+ trace_name = "Test" ,
392+ trace_type = "chain" ,
393+ project_name = "P" ,
394+ trace_id = trace_id ,
381395 )
382396 assert t .get_langchain_callback () is None
383397
@@ -400,8 +414,10 @@ def test_returns_handler_with_most_recent_component_span(self, tracer, mock_logg
400414
401415 # Add a component span
402416 tracer .add_trace (
403- trace_id = "comp-1" , trace_name = "OpenAI (comp-1)" ,
404- trace_type = "llm" , inputs = {},
417+ trace_id = "comp-1" ,
418+ trace_name = "OpenAI (comp-1)" ,
419+ trace_type = "llm" ,
420+ inputs = {},
405421 )
406422
407423 mock_handler = MagicMock ()
0 commit comments