1+ # SPDX-FileCopyrightText: Copyright (c), CloudZero, Inc. or its affiliates. All Rights Reserved.
2+ # SPDX-License-Identifier: Apache-2.0
3+
4+ """Tests for CLI commands to ensure all commands work correctly."""
5+
6+ import os
7+ import tempfile
8+ from pathlib import Path
9+ from unittest .mock import MagicMock , patch
10+
11+ import polars as pl
12+ import pytest
13+
14+ from ll2cz .cli import main
15+
16+
17+ class TestCLICommands :
18+ """Test all CLI commands to prevent regressions."""
19+
20+ @pytest .fixture
21+ def mock_config (self ):
22+ """Mock configuration."""
23+ with patch ('ll2cz.cli.Config' ) as mock :
24+ config_instance = MagicMock ()
25+ config_instance .get .side_effect = lambda key , default = None : {
26+ 'database_url' : 'postgresql://test@localhost/test' ,
27+ 'cz_api_key' : 'test-api-key' ,
28+ 'cz_connection_id' : 'test-connection-id'
29+ }.get (key , default )
30+ mock .return_value = config_instance
31+ yield config_instance
32+
33+ @pytest .fixture
34+ def mock_database (self ):
35+ """Mock database with test data."""
36+ test_data = pl .DataFrame ({
37+ 'api_key' : ['key1' , 'key2' ],
38+ 'custom_llm_provider' : ['openai' , 'anthropic' ],
39+ 'model' : ['gpt-4' , 'claude-3' ],
40+ 'api_request' : [1 , 1 ],
41+ 'cache_hits' : [0 , 0 ],
42+ 'cache_read' : [0 , 0 ],
43+ 'completion_tokens' : [100 , 200 ],
44+ 'prompt_tokens' : [50 , 100 ],
45+ 'response_cost_usd' : [0.01 , 0.02 ],
46+ 'spend' : [0.01 , 0.02 ],
47+ 'team_id' : ['team1' , 'team2' ],
48+ 'startTime' : ['2025-01-01' , '2025-01-02' ],
49+ 'request_id' : ['req1' , 'req2' ],
50+ 'user' : ['user1' , 'user2' ],
51+ 'team_alias' : [None , None ],
52+ 'team' : [None , None ],
53+ 'key_alias' : [None , None ],
54+ 'user_email' : [None , None ],
55+ 'user_alias' : [None , None ],
56+ 'project_tag' : [None , None ],
57+ 'org_tag' : [None , None ],
58+ 'custom_tag_0' : [None , None ],
59+ 'custom_tag_1' : [None , None ],
60+ 'custom_tag_2' : [None , None ],
61+ 'custom_tag_3' : [None , None ],
62+ 'custom_tag_4' : [None , None ]
63+ })
64+
65+ with patch ('ll2cz.cli.CachedLiteLLMDatabase' ) as mock_db , \
66+ patch ('ll2cz.analysis.CachedLiteLLMDatabase' ) as mock_db_analysis :
67+ db_instance = MagicMock ()
68+ db_instance .get_usage_data .return_value = test_data
69+ db_instance .get_individual_table_data .return_value = test_data
70+ db_instance .get_spend_analysis_data .return_value = test_data
71+ db_instance .is_server_available .return_value = True
72+ db_instance .get_cache_info .return_value = {
73+ 'total_records' : 2 ,
74+ 'user_records' : 2 ,
75+ 'team_records' : 0 ,
76+ 'tag_records' : 0 ,
77+ 'last_updated' : '2025-01-01T00:00:00'
78+ }
79+ mock_db .return_value = db_instance
80+ mock_db_analysis .return_value = db_instance
81+ yield db_instance
82+
83+ def test_help_command (self , capsys ):
84+ """Test that help command works."""
85+ with patch ('sys.argv' , ['ll2cz' , '--help' ]):
86+ with pytest .raises (SystemExit ) as exc_info :
87+ main ()
88+ assert exc_info .value .code == 0
89+ captured = capsys .readouterr ()
90+ assert 'Transform LiteLLM database data' in captured .out
91+
92+ def test_version_command (self , capsys ):
93+ """Test that version command works."""
94+ with patch ('sys.argv' , ['ll2cz' , '--version' ]):
95+ with pytest .raises (SystemExit ) as exc_info :
96+ main ()
97+ assert exc_info .value .code == 0
98+ captured = capsys .readouterr ()
99+ assert 'll2cz' in captured .out
100+
101+ def test_analyze_data_command (self , mock_config , mock_database , capsys ):
102+ """Test analyze data command."""
103+ with patch ('sys.argv' , ['ll2cz' , 'analyze' , 'data' , '--limit' , '5' ]):
104+ main ()
105+
106+ captured = capsys .readouterr ()
107+ assert 'Comprehensive Data Analysis' in captured .out
108+ assert 'Database Overview' in captured .out
109+
110+ def test_analyze_spend_command (self , mock_config , mock_database , capsys ):
111+ """Test analyze spend command."""
112+ with patch ('sys.argv' , ['ll2cz' , 'analyze' , 'spend' ]):
113+ main ()
114+
115+ captured = capsys .readouterr ()
116+ assert 'Spend Analysis' in captured .out or 'Analysis' in captured .out
117+
118+ def test_analyze_schema_command (self , mock_config , mock_database , capsys ):
119+ """Test analyze schema command."""
120+ # Mock database schema
121+ mock_database .get_schema_info .return_value = {
122+ 'test_table' : [
123+ {'column_name' : 'id' , 'data_type' : 'integer' , 'is_nullable' : 'NO' , 'column_default' : None }
124+ ]
125+ }
126+
127+ with patch ('sys.argv' , ['ll2cz' , 'analyze' , 'schema' ]):
128+ main ()
129+
130+ captured = capsys .readouterr ()
131+ assert 'Schema' in captured .out or 'test_table' in captured .out
132+
133+ def test_transform_command (self , mock_config , mock_database ):
134+ """Test transform command."""
135+ with tempfile .NamedTemporaryFile (mode = 'w' , suffix = '.csv' , delete = False ) as f :
136+ output_file = f .name
137+
138+ try :
139+ with patch ('sys.argv' , ['ll2cz' , 'transform' , '--output' , output_file , '--limit' , '5' ]):
140+ main ()
141+
142+ # Verify output file was created
143+ assert Path (output_file ).exists ()
144+
145+ # Verify content
146+ content = Path (output_file ).read_text ()
147+ assert content # Should not be empty
148+
149+ finally :
150+ if Path (output_file ).exists ():
151+ os .unlink (output_file )
152+
153+ def test_transmit_test_mode (self , mock_config , mock_database , capsys ):
154+ """Test transmit command in test mode."""
155+ with patch ('sys.argv' , ['ll2cz' , 'transmit' , '--mode' , 'all' , '--test' ]):
156+ main ()
157+
158+ captured = capsys .readouterr ()
159+ assert 'TEST MODE' in captured .out
160+ assert 'Sample payload' in captured .out
161+
162+ def test_transmit_day_mode (self , mock_config , mock_database , capsys ):
163+ """Test transmit command for specific day."""
164+ with patch ('sys.argv' , ['ll2cz' , 'transmit' , '--mode' , 'day' , '--date' , '01-01-2025' , '--test' ]):
165+ main ()
166+
167+ captured = capsys .readouterr ()
168+ assert 'TEST MODE' in captured .out or 'Day: 2025-01-01' in captured .out
169+
170+ def test_transmit_month_mode (self , mock_config , mock_database , capsys ):
171+ """Test transmit command for specific month."""
172+ with patch ('sys.argv' , ['ll2cz' , 'transmit' , '--mode' , 'month' , '--date' , '01-2025' , '--test' ]):
173+ main ()
174+
175+ captured = capsys .readouterr ()
176+ assert 'TEST MODE' in captured .out or 'Month: January 2025' in captured .out
177+
178+ def test_cache_status_command (self , mock_config , mock_database , capsys ):
179+ """Test cache status command."""
180+ with patch ('sys.argv' , ['ll2cz' , 'cache' , 'status' ]):
181+ main ()
182+
183+ captured = capsys .readouterr ()
184+ assert 'Cache Status' in captured .out
185+ assert 'Records cached' in captured .out
186+
187+ def test_cache_clear_command (self , mock_config , mock_database , capsys ):
188+ """Test cache clear command."""
189+ mock_database .clear_cache .return_value = None
190+
191+ with patch ('sys.argv' , ['ll2cz' , 'cache' , 'clear' ]):
192+ main ()
193+
194+ captured = capsys .readouterr ()
195+ assert 'Cache cleared' in captured .out
196+
197+ def test_cache_refresh_command (self , mock_config , mock_database , capsys ):
198+ """Test cache refresh command."""
199+ mock_database .refresh_cache .return_value = {
200+ 'user' : {'added' : 10 , 'updated' : 5 , 'unchanged' : 100 },
201+ 'team' : {'added' : 0 , 'updated' : 0 , 'unchanged' : 0 },
202+ 'tag' : {'added' : 0 , 'updated' : 0 , 'unchanged' : 0 }
203+ }
204+
205+ with patch ('sys.argv' , ['ll2cz' , 'cache' , 'refresh' ]):
206+ main ()
207+
208+ captured = capsys .readouterr ()
209+ assert 'Refresh Results' in captured .out or 'Cache refreshed' in captured .out
210+
211+ def test_config_show_command (self , mock_config , capsys ):
212+ """Test config show command."""
213+ with patch ('sys.argv' , ['ll2cz' , 'config' , 'show' ]):
214+ main ()
215+
216+ captured = capsys .readouterr ()
217+ assert 'Configuration' in captured .out or 'Configured' in captured .out
218+
219+ def test_config_example_command (self , mock_config ):
220+ """Test config example command."""
221+ with tempfile .TemporaryDirectory () as tmpdir :
222+ config_file = Path (tmpdir ) / 'config.yml'
223+
224+ with patch ('ll2cz.cli.Path.home' , return_value = Path (tmpdir )):
225+ with patch ('sys.argv' , ['ll2cz' , 'config' , 'example' ]):
226+ main ()
227+
228+ # Verify config file was created
229+ expected_path = Path (tmpdir ) / '.ll2cz' / 'config.yml'
230+ assert expected_path .exists ()
231+
232+ def test_invalid_command (self , capsys ):
233+ """Test that invalid commands show error."""
234+ with patch ('sys.argv' , ['ll2cz' , 'invalid-command' ]):
235+ with pytest .raises (SystemExit ) as exc_info :
236+ main ()
237+ assert exc_info .value .code != 0
238+
239+ captured = capsys .readouterr ()
240+ assert 'invalid choice' in captured .err or 'error' in captured .err .lower ()
241+
242+ def test_transform_with_source_logs (self , mock_config , mock_database ):
243+ """Test transform command with logs source."""
244+ with tempfile .NamedTemporaryFile (mode = 'w' , suffix = '.csv' , delete = False ) as f :
245+ output_file = f .name
246+
247+ try :
248+ # Mock spend logs data
249+ mock_database .get_spend_logs_data .return_value = mock_database .get_usage_data ()
250+
251+ with patch ('sys.argv' , ['ll2cz' , 'transform' , '--source' , 'logs' , '--output' , output_file ]):
252+ main ()
253+
254+ assert Path (output_file ).exists ()
255+
256+ finally :
257+ if Path (output_file ).exists ():
258+ os .unlink (output_file )
259+
260+ def test_transmit_with_append_mode (self , mock_config , mock_database , capsys ):
261+ """Test transmit command with append mode."""
262+ with patch ('sys.argv' , ['ll2cz' , 'transmit' , '--mode' , 'all' , '--append' , '--test' ]):
263+ main ()
264+
265+ captured = capsys .readouterr ()
266+ assert 'TEST MODE' in captured .out
267+ # In append mode, operation should be 'sum' instead of 'replace_hourly'
268+ assert 'sum' in captured .out or 'append' in captured .out .lower ()
0 commit comments