From 53b8e1849c511e59d544e9456841467adc00c95d Mon Sep 17 00:00:00 2001
From: Aliaksandr Kuzmik <alexander.kuzmik99@gmail.com>
Date: Tue, 21 Jan 2025 16:57:41 +0100
Subject: [PATCH] Update unit tests and backend emulator

---
 .../tests/testlib/backend_emulator_message_processor.py      | 2 ++
 sdks/python/tests/testlib/models.py                          | 1 +
 sdks/python/tests/unit/decorator/test_tracker_outputs.py     | 5 ++++-
 .../unit/message_processing/batching/test_batch_manager.py   | 1 +
 4 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/sdks/python/tests/testlib/backend_emulator_message_processor.py b/sdks/python/tests/testlib/backend_emulator_message_processor.py
index 9240cd999e..3cd6f4ed2e 100644
--- 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
index bfd6d19587..e214d0305e 100644
--- 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
index b8c07eceff..2cedc30640 100644
--- 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
index 94b9ccfb4a..933177756d 100644
--- 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(