litellm/tests/unit/integrations/cloudzero/test_transform.py

from datetime import datetime
from unittest.mock import MagicMock, patch

import polars as pl
import pytest


from litellm.integrations.cloudzero.transform import CBFTransformer
from litellm.types.integrations.cloudzero import CBFRecord


class TestCBFTransformer:
    """Test suite for CBFTransformer class."""

    def test_init(self):
        """Test CBFTransformer initialization."""
        transformer = CBFTransformer()
        assert hasattr(transformer, "czrn_generator")
        assert transformer.czrn_generator is not None

    def test_transform_empty_dataframe(self):
        """Test transform method with empty DataFrame."""
        transformer = CBFTransformer()
        empty_df = pl.DataFrame()

        result = transformer.transform(empty_df)

        assert result.is_empty()
        assert isinstance(result, pl.DataFrame)

    def test_transform_with_zero_successful_requests(self):
        """Test transform method filters out records with zero successful_requests."""
        transformer = CBFTransformer()
        data = pl.DataFrame(
            {
                "date": ["2025-01-19"],
                "successful_requests": [0],
                "spend": [10.0],
                "entity_id": ["test_entity"],
                "model": ["gpt-4"],
            }
        )

        result = transformer.transform(data)

        assert result.is_empty()

    def test_transform_with_valid_data(self):
        """Test transform method with valid data."""
        transformer = CBFTransformer()
        with patch.object(transformer, "_create_cbf_record") as mock_create:
            mock_create.return_value = CBFRecord({"test": "data"})

            data = pl.DataFrame(
                {
                    "date": ["2025-01-19"],
                    "successful_requests": [5],
                    "spend": [10.0],
                    "entity_id": ["test_entity"],
                    "model": ["gpt-4"],
                }
            )

            result = transformer.transform(data)

            assert len(result) == 1
            mock_create.assert_called_once()

    def test_transform_handles_czrn_generation_failures(self):
        """Test transform method handles CZRN generation failures gracefully."""
        transformer = CBFTransformer()
        with patch.object(transformer, "_create_cbf_record") as mock_create:
            mock_create.side_effect = Exception("CZRN generation failed")

            data = pl.DataFrame(
                {
                    "date": ["2025-01-19"],
                    "successful_requests": [5],
                    "spend": [10.0],
                    "entity_id": ["test_entity"],
                    "model": ["gpt-4"],
                }
            )

            result = transformer.transform(data)

            assert result.is_empty()

    def test_transform_keeps_tags_first_seen_after_row_100(self):
        transformer = CBFTransformer()
        teamless_rows = 101
        team_rows = 2
        total_rows = teamless_rows + team_rows
        data = pl.DataFrame(
            {
                "date": ["2025-01-19"] * total_rows,
                "successful_requests": [1] * total_rows,
                "spend": [0.5] * total_rows,
                "prompt_tokens": [10] * total_rows,
                "completion_tokens": [5] * total_rows,
                "model": ["gpt-4"] * total_rows,
                "custom_llm_provider": ["openai"] * total_rows,
                "api_key": ["sk-late-team"] * total_rows,
                "team_id": pl.Series([None] * teamless_rows + ["team-late"] * team_rows, dtype=pl.String),
                "team_alias": pl.Series([None] * teamless_rows + ["Late Team"] * team_rows, dtype=pl.String),
            }
        )

        result = transformer.transform(data)

        assert len(result) == total_rows
        assert "resource/tag:team_alias" in result.columns
        assert result["resource/tag:team_alias"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows
        assert result["resource/tag:entity_id"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows

    def test_create_cbf_record(self):
        """Test _create_cbf_record method with valid row data."""
        transformer = CBFTransformer()
        with (
            patch.object(
                transformer.czrn_generator, "create_from_litellm_data"
            ) as mock_czrn,
            patch.object(
                transformer.czrn_generator, "extract_components"
            ) as mock_extract,
        ):

            mock_czrn.return_value = "test-czrn"
            mock_extract.return_value = (
                "service",
                "provider",
                "region",
                "account",
                "resource",
                "local_id",
            )

            row = {
                "date": "2025-01-19",
                "spend": 10.5,
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "entity_id": "test_entity",
                "model": "gpt-4",
                "entity_type": "user",
                "model_group": "openai",
                "custom_llm_provider": "openai",
                "api_key": "sk-test123",
                "api_requests": 5,
                "successful_requests": 5,
                "failed_requests": 0,
            }

            result = transformer._create_cbf_record(row)

            assert isinstance(result, CBFRecord)
            assert result["cost/cost"] == 10.5
            assert result["usage/amount"] == 150  # 100 + 50
            assert result["usage/units"] == "tokens"
            assert result["resource/id"] == "test-czrn"

    def test_create_cbf_record_adds_user_email_tag(self):
        """Test that user_email field is emitted as a resource tag when present."""
        transformer = CBFTransformer()
        with (
            patch.object(
                transformer.czrn_generator, "create_from_litellm_data"
            ) as mock_czrn,
            patch.object(
                transformer.czrn_generator, "extract_components"
            ) as mock_extract,
        ):

            mock_czrn.return_value = "test-czrn"
            mock_extract.return_value = (
                "service",
                "provider",
                "region",
                "account",
                "resource",
                "local_id",
            )

            row = {
                "date": "2025-01-19",
                "spend": 1.0,
                "prompt_tokens": 10,
                "completion_tokens": 5,
                "model": "gpt-4",
                "api_key": "sk-useremail",
                "team_id": "team-123",
                "team_alias": "Dev Team",
                "user_email": "user@example.com",
            }

            result = transformer._create_cbf_record(row)

            assert result["resource/tag:user_email"] == "user@example.com"

    def test_create_cbf_record_omits_empty_user_email(self):
        """Test that empty user_email values are not added as resource tags."""
        transformer = CBFTransformer()
        with (
            patch.object(
                transformer.czrn_generator, "create_from_litellm_data"
            ) as mock_czrn,
            patch.object(
                transformer.czrn_generator, "extract_components"
            ) as mock_extract,
        ):

            mock_czrn.return_value = "test-czrn"
            mock_extract.return_value = (
                "service",
                "provider",
                "region",
                "account",
                "resource",
                "local_id",
            )

            row = {
                "date": "2025-01-19",
                "spend": 1.0,
                "prompt_tokens": 10,
                "completion_tokens": 5,
                "model": "gpt-4",
                "api_key": "sk-useremail",
                "team_id": "team-123",
                "team_alias": "Dev Team",
                "user_email": None,
            }

            result = transformer._create_cbf_record(row)

            assert "resource/tag:user_email" not in result

    def test_create_cbf_record_minimal_data(self):
        """Test _create_cbf_record method with minimal row data."""
        transformer = CBFTransformer()
        with (
            patch.object(
                transformer.czrn_generator, "create_from_litellm_data"
            ) as mock_czrn,
            patch.object(
                transformer.czrn_generator, "extract_components"
            ) as mock_extract,
        ):

            mock_czrn.return_value = "test-czrn"
            mock_extract.return_value = (
                "service",
                "provider",
                "region",
                "account",
                "resource",
                "local_id",
            )

            row = {"date": "2025-01-19", "spend": 0.0}

            result = transformer._create_cbf_record(row)

            assert isinstance(result, CBFRecord)
            assert result["cost/cost"] == 0.0
            assert result["usage/amount"] == 0  # no tokens
            assert result["usage/units"] == "tokens"

    def test_parse_date_with_valid_string(self):
        """Test _parse_date method with valid date string."""
        transformer = CBFTransformer()

        result = transformer._parse_date("2025-01-19")

        assert isinstance(result, datetime)
        assert result.year == 2025
        assert result.month == 1
        assert result.day == 19

    def test_parse_date_with_datetime_object(self):
        """Test _parse_date method with datetime object."""
        transformer = CBFTransformer()
        dt = datetime(2025, 1, 19)

        result = transformer._parse_date(dt)

        assert result == dt

    def test_parse_date_with_none(self):
        """Test _parse_date method with None."""
        transformer = CBFTransformer()

        result = transformer._parse_date(None)

        assert result is None

    def test_parse_date_with_invalid_string(self):
        """Test _parse_date method with invalid date string."""
        transformer = CBFTransformer()

        result = transformer._parse_date("invalid-date")

        assert result is None

    def test_parse_date_with_iso_format(self):
        """Test _parse_date method with ISO format string."""
        transformer = CBFTransformer()

        result = transformer._parse_date("2025-01-19T10:30:00Z")

        assert isinstance(result, datetime)
        assert result.year == 2025