|
import pytest
|
|
import tempfile
|
|
import os
|
|
import json
|
|
from pr.core.usage_tracker import UsageTracker
|
|
|
|
@pytest.fixture
|
|
def temp_usage_file(tmp_path, monkeypatch):
|
|
from pr.core import usage_tracker
|
|
original_file = usage_tracker.USAGE_DB_FILE
|
|
temp_file = str(tmp_path / "usage.json")
|
|
monkeypatch.setattr(usage_tracker, 'USAGE_DB_FILE', temp_file)
|
|
yield temp_file
|
|
if os.path.exists(temp_file):
|
|
os.remove(temp_file)
|
|
monkeypatch.setattr(usage_tracker, 'USAGE_DB_FILE', original_file)
|
|
|
|
def test_usage_tracker_init():
|
|
tracker = UsageTracker()
|
|
summary = tracker.get_session_summary()
|
|
assert summary['requests'] == 0
|
|
assert summary['total_tokens'] == 0
|
|
assert summary['estimated_cost'] == 0.0
|
|
|
|
def test_track_request_known_model():
|
|
tracker = UsageTracker()
|
|
tracker.track_request('gpt-3.5-turbo', 100, 50)
|
|
|
|
summary = tracker.get_session_summary()
|
|
assert summary['requests'] == 1
|
|
assert summary['input_tokens'] == 100
|
|
assert summary['output_tokens'] == 50
|
|
assert summary['total_tokens'] == 150
|
|
assert 'gpt-3.5-turbo' in summary['models_used']
|
|
# Cost: (100/1000)*0.0005 + (50/1000)*0.0015 = 0.00005 + 0.000075 = 0.000125
|
|
assert abs(summary['estimated_cost'] - 0.000125) < 1e-6
|
|
|
|
def test_track_request_unknown_model():
|
|
tracker = UsageTracker()
|
|
tracker.track_request('unknown-model', 100, 50)
|
|
|
|
summary = tracker.get_session_summary()
|
|
assert summary['requests'] == 1
|
|
assert summary['estimated_cost'] == 0.0 # Unknown model, cost 0
|
|
|
|
def test_track_request_multiple():
|
|
tracker = UsageTracker()
|
|
tracker.track_request('gpt-3.5-turbo', 100, 50)
|
|
tracker.track_request('gpt-4', 200, 100)
|
|
|
|
summary = tracker.get_session_summary()
|
|
assert summary['requests'] == 2
|
|
assert summary['input_tokens'] == 300
|
|
assert summary['output_tokens'] == 150
|
|
assert summary['total_tokens'] == 450
|
|
assert len(summary['models_used']) == 2
|
|
|
|
def test_get_formatted_summary():
|
|
tracker = UsageTracker()
|
|
tracker.track_request('gpt-3.5-turbo', 100, 50)
|
|
|
|
formatted = tracker.get_formatted_summary()
|
|
assert "Total Requests: 1" in formatted
|
|
assert "Total Tokens: 150" in formatted
|
|
assert "Estimated Cost: $0.0001" in formatted
|
|
assert "gpt-3.5-turbo" in formatted
|
|
|
|
def test_get_total_usage_no_file(temp_usage_file):
|
|
total = UsageTracker.get_total_usage()
|
|
assert total['total_requests'] == 0
|
|
assert total['total_tokens'] == 0
|
|
assert total['total_cost'] == 0.0
|
|
|
|
def test_get_total_usage_with_data(temp_usage_file):
|
|
# Manually create history file
|
|
history = [
|
|
{'timestamp': '2023-01-01', 'model': 'gpt-3.5-turbo', 'input_tokens': 100, 'output_tokens': 50, 'total_tokens': 150, 'cost': 0.000125},
|
|
{'timestamp': '2023-01-02', 'model': 'gpt-4', 'input_tokens': 200, 'output_tokens': 100, 'total_tokens': 300, 'cost': 0.008}
|
|
]
|
|
with open(temp_usage_file, 'w') as f:
|
|
json.dump(history, f)
|
|
|
|
total = UsageTracker.get_total_usage()
|
|
assert total['total_requests'] == 2
|
|
assert total['total_tokens'] == 450
|
|
assert abs(total['total_cost'] - 0.008125) < 1e-6
|