""" tests/test_context.py — C1.4 Tests for context(), export_data(), and token budget enforcement. """ import sys import os import datetime import pytest sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_kv(tmp_path): from db import StateKV os.environ.pop("AGENTCACHE_SECRET", None) return StateKV(db_path=str(tmp_path / "test.db")) def _now(): return datetime.datetime.utcnow().isoformat() + "Z" # --------------------------------------------------------------------------- # context() tests # --------------------------------------------------------------------------- class TestContext: def test_empty_db_returns_minimal_context(self, tmp_path): """Empty DB should return a well-formed but empty context.""" from functions import context kv = _make_kv(tmp_path) result = context(kv, { "sessionId": "sess_test_empty", "project": "/home/user/my-project", "budget": 2000, }) assert isinstance(result, dict) assert "context" in result assert "blocks" in result assert "tokens" in result assert result["blocks"] == 0 assert result["tokens"] == 0 def test_raises_on_missing_session_id(self, tmp_path): from functions import context kv = _make_kv(tmp_path) with pytest.raises(ValueError): context(kv, {"project": "/home/user/proj"}) def test_raises_on_missing_project(self, tmp_path): from functions import context kv = _make_kv(tmp_path) with pytest.raises(ValueError): context(kv, {"sessionId": "sess_x"}) def test_respects_token_budget(self, tmp_path): """Context output tokens should not exceed the requested budget.""" from functions import context, lesson_save, KV kv = _make_kv(tmp_path) project = "/home/user/budget-test" # Add many lessons to push towards the budget for i in range(20): lesson_save(kv, { "content": f"Lesson {i}: " + ("x " * 100), "project": project, "confidence": 0.9, }) budget = 500 result = context(kv, { "sessionId": "sess_budget", "project": project, "budget": budget, }) # Token estimate is len/3 — check that total tokens respects budget assert result["tokens"] <= budget + 50 # small headroom for header/footer def test_context_includes_xml_wrapper(self, tmp_path): """Non-empty context should be wrapped in .""" from functions import context, lesson_save kv = _make_kv(tmp_path) project = "/home/user/xml-test" lesson_save(kv, { "content": "Always validate user input before processing", "project": project, "confidence": 0.8, }) result = context(kv, { "sessionId": "sess_xml", "project": project, "budget": 2000, }) if result["blocks"] > 0: assert "" in result["context"] def test_token_budget_env_var_respected(self, tmp_path, monkeypatch): """TOKEN_BUDGET env var should be used when no budget param given.""" from functions import context, lesson_save kv = _make_kv(tmp_path) project = "/home/user/env-budget" monkeypatch.setenv("TOKEN_BUDGET", "100") for i in range(10): lesson_save(kv, { "content": f"Important lesson {i}: " + ("word " * 50), "project": project, "confidence": 0.9, }) result = context(kv, {"sessionId": "sess_env_budget", "project": project}) # Should use TOKEN_BUDGET=100 from env assert result["tokens"] <= 150 # with some headroom for XML wrapper # --------------------------------------------------------------------------- # export_data() tests # --------------------------------------------------------------------------- class TestExportData: def test_export_returns_folders_and_memories(self, tmp_path): from functions import folder_observe, remember, export_data kv = _make_kv(tmp_path) folder_observe(kv, { "folderPath": "/home/user/export-test", "agentId": "kiro", "text": "Working on export feature", "timestamp": _now(), }) remember(kv, {"content": "Export data uses v2 format"}) result = export_data(kv, {}) assert isinstance(result, dict) assert "folders" in result or "observations" in result or "memories" in result def test_export_empty_db(self, tmp_path): from functions import export_data kv = _make_kv(tmp_path) result = export_data(kv, {}) assert isinstance(result, dict) # Should not crash on empty DB # --------------------------------------------------------------------------- # estimate_tokens() # --------------------------------------------------------------------------- class TestEstimateTokens: def test_empty_string(self): from functions import estimate_tokens assert estimate_tokens("") == 0 def test_typical_text(self): from functions import estimate_tokens text = "hello world this is a test" * 10 tokens = estimate_tokens(text) # Should be approximately len/3 assert tokens == len(text) // 3