Update unit tests and backend emulator

comet-ml · Jan 21, 2025 · 53b8e18 · 53b8e18
1 parent a70df2b
commit 53b8e18
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 1 deletion.
diff --git a/sdks/python/tests/testlib/backend_emulator_message_processor.py b/sdks/python/tests/testlib/backend_emulator_message_processor.py
@@ -129,6 +129,7 @@ def _dispatch_message(self, message: messages.BaseMessage) -> None:
                 model=message.model,
                 provider=message.provider,
                 error_info=message.error_info,
+                total_cost=message.total_cost,
             )
 
             self._span_to_parent_span[span.id] = message.parent_span_id
@@ -156,6 +157,7 @@ def _dispatch_message(self, message: messages.BaseMessage) -> None:
                 "error_info": message.error_info,
                 "tags": message.tags,
                 "input": message.input,
+                "total_cost": message.total_cost,
             }
             cleaned_update_payload = dict_utils.remove_none_from_dict(update_payload)
             span.__dict__.update(cleaned_update_payload)

diff --git a/sdks/python/tests/testlib/models.py b/sdks/python/tests/testlib/models.py
@@ -29,6 +29,7 @@ class SpanModel:
     model: Optional[str] = None
     provider: Optional[str] = None
     error_info: Optional[ErrorInfoDict] = None
+    total_cost: Optional[float] = None
 
 
 @dataclasses.dataclass

diff --git a/sdks/python/tests/unit/decorator/test_tracker_outputs.py b/sdks/python/tests/unit/decorator/test_tracker_outputs.py
@@ -1058,7 +1058,9 @@ def test_track__span_and_trace_updated_via_opik_context(fake_backend):
     @tracker.track
     def f(x):
         opik_context.update_current_span(
-            name="span-name", metadata={"span-metadata-key": "span-metadata-value"}
+            name="span-name",
+            metadata={"span-metadata-key": "span-metadata-value"},
+            total_cost=0.42,
         )
         opik_context.update_current_trace(
             name="trace-name",
@@ -1087,6 +1089,7 @@ def f(x):
                 output={"output": "f-output"},
                 start_time=ANY_BUT_NONE,
                 end_time=ANY_BUT_NONE,
+                total_cost=0.42,
                 spans=[],
             )
         ],

diff --git a/sdks/python/tests/unit/message_processing/batching/test_batch_manager.py b/sdks/python/tests/unit/message_processing/batching/test_batch_manager.py
@@ -100,6 +100,7 @@ def test_batch_manager__start_and_stop_were_called__accumulated_data_is_flushed(
         model=NOT_USED,
         provider=NOT_USED,
         error_info=NOT_USED,
+        total_cost=NOT_USED,
     )
 
     example_span_batcher = batchers.CreateSpanMessageBatcher(