47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
import tempfile
|
|
import unittest
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
from javis.memory.usage_store import ProviderEvent, SQLiteUsageStore
|
|
|
|
|
|
class UsageStoreTests(unittest.TestCase):
|
|
def test_records_only_metadata_and_counts_daily_usage(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
database = Path(directory) / "usage.sqlite3"
|
|
store = SQLiteUsageStore(database)
|
|
store.record(
|
|
ProviderEvent(
|
|
provider="gemini",
|
|
model="gemini-test",
|
|
success=False,
|
|
error_category="quota",
|
|
fallback=False,
|
|
privacy_policy="allowed",
|
|
input_tokens=12,
|
|
output_tokens=None,
|
|
)
|
|
)
|
|
store.record(
|
|
ProviderEvent(
|
|
provider="ollama",
|
|
model="local-test",
|
|
success=True,
|
|
error_category=None,
|
|
fallback=True,
|
|
privacy_policy="allowed",
|
|
)
|
|
)
|
|
|
|
today = datetime.now(UTC).date()
|
|
self.assertEqual(store.cloud_requests_on(today), 1)
|
|
self.assertEqual(store.local_fallbacks_on(today), 1)
|
|
raw_database = database.read_bytes()
|
|
self.assertNotIn(b"prompt", raw_database)
|
|
self.assertNotIn(b"answer", raw_database)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|