tests: Unit Tests for SpanIOProcessor and groq utils

tests/fixtures/__init__.py

@@ -0,0 +1 @@
+from .base_provider_utils import BaseProviderUtils

tests/fixtures/base_provider_utils.py

@@ -0,0 +1,82 @@
+import itertools
+from typing import Callable
+from unittest.mock import MagicMock
+
+from opentelemetry.semconv_ai import SpanAttributes
+
+
+class BaseProviderUtils:
+    """Base class to be inherited by test cases that are going to test `utils.py` from `netra/instrumentation/openai/`, `netra/instrumentation/groq/` etc."""
+
+    ALIASES = {
+        "prompt_tokens": ["prompt_tokens", "input_tokens"],
+        "completion_tokens": ["completion_tokens", "output_tokens"],
+        "prompt_tokens_details": ["prompt_tokens_details", "input_tokens_details"],
+    }
+
+    P_TOKEN = 100
+    C_TOKEN = 50
+    P_DETAIL = 10
+    T_TOKEN = 160
+
+    _set_usage_attributes_method: Callable = None
+
+    def _build_input_data(self):
+        keys_groups = [
+            self.ALIASES["prompt_tokens"],
+            self.ALIASES["completion_tokens"],
+            self.ALIASES["prompt_tokens_details"],
+        ]
+
+        for p_token, c_token, p_detail in itertools.product(*keys_groups):
+            data = dict()
+
+            data[p_token] = self.P_TOKEN
+            data[c_token] = self.C_TOKEN
+            data[p_detail] = {"cached_tokens": self.P_DETAIL}
+            data["total_tokens"] = self.T_TOKEN
+
+            yield data
+
+    def _build_no_details_data(self):
+        keys_group = [self.ALIASES["prompt_tokens"], self.ALIASES["completion_tokens"]]
+
+        for p_token, c_token in itertools.product(*keys_group):
+            data = dict()
+
+            data[p_token] = self.P_TOKEN
+            data[c_token] = self.C_TOKEN
+            data["total_tokens"] = self.T_TOKEN
+
+            yield data
+
+    def test_set_usage_attributes(self):
+        """Tests _set_usage_attributes"""
+        for i, dummy_data in enumerate(self._build_input_data()):
+            with self.subTest(scenario=f"Combination {i}", payload=dummy_data):
+                mock_span = MagicMock()
+                self._set_usage_attributes_method(mock_span, dummy_data)
+
+                mock_span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, self.P_TOKEN)
+                mock_span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, self.C_TOKEN)
+                mock_span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, self.T_TOKEN)
+                mock_span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS, self.P_DETAIL)
+
+    def test_set_usage_attributes_no_prompt_tokens_details(self):
+        """Tests _set_usage_attributes without prompt token details"""
+        for i, dummy_data in enumerate(self._build_no_details_data()):
+            with self.subTest(scenario=f"Combination {i}", payload=dummy_data):
+                mock_span = MagicMock()
+                self._set_usage_attributes_method(mock_span, dummy_data)
+
+                mock_span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, self.P_TOKEN)
+                mock_span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, self.C_TOKEN)
+                mock_span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, self.T_TOKEN)
+
+    def test_empty_dict(self):
+        """Tests _set_usage_attributes with empty dictionary"""
+        mock_span = MagicMock()
+        self._set_usage_attributes_method(mock_span, dict())
+
+        called_keys = [call[0][0] for call in mock_span.set_attribute.call_args_list]
+        assert len(called_keys) == 0

tests/test_groq_provider_utils.py

@@ -0,0 +1,9 @@
+import unittest
+
+from netra.instrumentation.groq.utils import _set_usage_attributes
+
+from .fixtures.base_provider_utils import BaseProviderUtils
+
+
+class TestGroqProviderUtils(unittest.TestCase, BaseProviderUtils):
+    _set_usage_attributes_method = staticmethod(_set_usage_attributes)

tests/test_span_io_processor.py

@@ -0,0 +1,39 @@
+import unittest
+from unittest.mock import MagicMock
+
+from netra.processors.span_io_processor import SpanIOProcessor
+
+
+class TestSpanIOProcessor(unittest.TestCase):
+    """Test cases for `SpanIOProcessor` as per `4.3` in PRD"""
+
+    _USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
+    _USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
+    _USAGE_PROMPT_TOKENS = "gen_ai.usage.prompt_tokens"
+    _USAGE_COMPLETION_TOKENS = "gen_ai.usage.completion_tokens"
+
+    def _get_mocks(self, original_value: str, desired_value: str):
+        mock_span = MagicMock()
+        mock_original_set_attribute = MagicMock()
+        mock_span.set_attribute = mock_original_set_attribute
+        # creating fake span
+
+        SpanIOProcessor._wrap_set_attribute(mock_span)
+        # adding set_attribute method to span
+
+        mock_span.set_attribute(original_value, 100)
+        # setting value
+
+        mock_original_set_attribute.assert_called_with(desired_value, 100)
+
+    def test_alias_to_prompt_tokens(self):
+        self._get_mocks(self._USAGE_INPUT_TOKENS, self._USAGE_PROMPT_TOKENS)
+
+    def test_alias_to_completion_tokens(self):
+        self._get_mocks(self._USAGE_OUTPUT_TOKENS, self._USAGE_COMPLETION_TOKENS)
+
+    def test_pass_through_prompt_tokens(self):
+        self._get_mocks(self._USAGE_PROMPT_TOKENS, self._USAGE_PROMPT_TOKENS)
+
+    def test_pass_through_completion_tokens(self):
+        self._get_mocks(self._USAGE_COMPLETION_TOKENS, self._USAGE_COMPLETION_TOKENS)