litellm/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py

import os
import sys
import zoneinfo
from datetime import datetime, timezone
from unittest.mock import MagicMock, Mock, patch

import httpx
import polars as pl
import pytest

sys.path.insert(0, os.path.abspath("../../../.."))

from litellm.integrations.cloudzero.cz_stream_api import CloudZeroStreamer


class TestCloudZeroStreamer:
    """Test suite for CloudZeroStreamer class."""

    def test_init_with_defaults(self):
        """Test CloudZeroStreamer initialization with default parameters."""
        streamer = CloudZeroStreamer(
            api_key="test-key",
            connection_id="test-connection"
        )
        
        assert streamer.api_key == "test-key"
        assert streamer.connection_id == "test-connection"
        assert streamer.base_url == "https://api.cloudzero.com"
        assert streamer.user_timezone == timezone.utc

    def test_init_with_valid_timezone(self):
        """Test CloudZeroStreamer initialization with valid timezone."""
        streamer = CloudZeroStreamer(
            api_key="test-key",
            connection_id="test-connection",
            user_timezone="America/New_York"
        )
        
        assert streamer.user_timezone == zoneinfo.ZoneInfo("America/New_York")
    
    def test_send_batched_with_valid_data(self):
        """Test send_batched method with valid data."""
        streamer = CloudZeroStreamer("test-key", "test-connection")
        with patch.object(streamer, '_group_by_date') as mock_group, \
             patch.object(streamer, '_send_daily_batch') as mock_send:
            
            mock_group.return_value = {
                '2025-01-19': pl.DataFrame({'test': ['data1']}),
                '2025-01-20': pl.DataFrame({'test': ['data2']})
            }
            
            data = pl.DataFrame({'test': ['data']})
            streamer.send_batched(data, "replace_hourly")
            
            assert mock_send.call_count == 2

    def test_group_by_date_valid_data(self):
        """Test _group_by_date method with valid data."""
        streamer = CloudZeroStreamer("test-key", "test-connection")
        with patch.object(streamer, '_parse_and_convert_timestamp') as mock_parse:
            mock_parse.return_value = datetime(2025, 1, 19, 10, 30, 0, tzinfo=timezone.utc)
            
            data = pl.DataFrame({
                'time/usage_start': ['2025-01-19T10:30:00Z'],
                'cost': [10.0]
            })
            
            result = streamer._group_by_date(data)
            
            assert '2025-01-19' in result
            assert len(result['2025-01-19']) == 1


    def test_parse_and_convert_timestamp_utc(self):
        """Test _parse_and_convert_timestamp method with UTC timestamp."""
        streamer = CloudZeroStreamer("test-key", "test-connection")
        
        result = streamer._parse_and_convert_timestamp('2025-01-19T10:30:00Z')
        
        assert result.year == 2025
        assert result.month == 1
        assert result.day == 19
        assert result.hour == 10
        assert result.minute == 30
        assert result.tzinfo == timezone.utc

    def test_parse_and_convert_timestamp_with_offset(self):
        """Test _parse_and_convert_timestamp method with timezone offset."""
        streamer = CloudZeroStreamer("test-key", "test-connection")
        
        result = streamer._parse_and_convert_timestamp('2025-01-19T10:30:00+05:00')
        
        assert result.tzinfo == timezone.utc
        assert result.hour == 5  # Converted to UTC

    def test_parse_and_convert_timestamp_no_timezone(self):
        """Test _parse_and_convert_timestamp method without timezone info."""
        streamer = CloudZeroStreamer("test-key", "test-connection", user_timezone="America/New_York")
        
        result = streamer._parse_and_convert_timestamp('2025-01-19T10:30:00')
        
        assert result.tzinfo == timezone.utc

    def test_parse_and_convert_timestamp_invalid(self):
        """Test _parse_and_convert_timestamp method with invalid timestamp."""
        streamer = CloudZeroStreamer("test-key", "test-connection")
        
        with pytest.raises(ValueError):
            streamer._parse_and_convert_timestamp('invalid-timestamp')

    def test_prepare_batch_payload(self):
        """Test _prepare_batch_payload method."""
        streamer = CloudZeroStreamer("test-key", "test-connection")
        with patch.object(streamer, '_convert_cbf_to_api_format') as mock_convert:
            mock_convert.return_value = {'test': 'record'}
            
            batch_data = pl.DataFrame({'cost': [10.0]})
            result = streamer._prepare_batch_payload('2025-01-19', batch_data, 'replace_hourly')
            
            assert result['month'] == '2025-01'
            assert result['operation'] == 'replace_hourly'
            assert len(result['data']) == 1



    def test_convert_cbf_to_api_format_valid_data(self):
        """Test _convert_cbf_to_api_format method with valid data."""
        streamer = CloudZeroStreamer("test-key", "test-connection")
        with patch.object(streamer, '_ensure_utc_timestamp') as mock_ensure:
            mock_ensure.return_value = '2025-01-19T10:30:00Z'
            
            row = {
                'time/usage_start': '2025-01-19T10:30:00Z',
                'cost/cost': 10.5,
                'tokens': 100,
                'text_field': 'test'
            }
            
            result = streamer._convert_cbf_to_api_format(row)
            
            assert result['cost/cost'] == '10.5'
            assert result['tokens'] == '100'
            assert result['text_field'] == 'test'

    def test_convert_cbf_to_api_format_float_precision(self):
        """Test _convert_cbf_to_api_format method handles float precision correctly."""
        streamer = CloudZeroStreamer("test-key", "test-connection")
        
        row = {
            'cost': 10.123456789012345,
            'large_float': 1234567890.0
        }
        
        result = streamer._convert_cbf_to_api_format(row)
        
        # Should avoid scientific notation
        assert 'e' not in result['cost'].lower()
        assert 'e' not in result['large_float'].lower()