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()