Added LiteLLM to the stack

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

@@ -0,0 +1,160 @@
+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()
+
+