from __future__ import annotations import json import tempfile import unittest import zipfile from pathlib import Path from kimu.artifact_storage import EncryptedArtifactStore from kimu.crypto import EncryptionError, decrypt_file_bytes, encrypt_file, file_aad, generate_master_key, is_encrypted_file, parse_master_key from kimu.db import connect, initialise from kimu.fileops import ( FileValidationError, attachment_content, inspect_stored_file, safe_filename, ) from kimu.models import build_profiles, normalise_preferences from kimu.provider import build_payload, parse_sse_lines from kimu.security import hash_password, verify_password from kimu.tci import fake_execute class CoreTests(unittest.TestCase): def setUp(self): self.cfg = { "K26_MODEL_ID": "moonshotai/Kimi-K2.6", "K27_MODEL_ID": "moonshotai/Kimi-K2.7-Code", "K27_CONFIGURED": True, "TOGETHER_API_KEY": "x", "FAKE_PROVIDER": False, } def test_profiles(self): profiles = build_profiles(self.cfg) self.assertTrue(profiles["k2.6"].reasoning_toggle) self.assertFalse(profiles["k2.7-code"].reasoning_toggle) self.assertTrue(profiles["k2.7-code"].preserve_reasoning) def test_preferences_lock(self): profile = build_profiles(self.cfg)["k2.7-code"] value = normalise_preferences( profile, {"reasoning_enabled": False, "temperature": 0.2} ) self.assertTrue(value["reasoning_enabled"]) self.assertEqual(value["temperature"], 1.0) def test_payloads(self): payload = build_payload( "moonshotai/Kimi-K2.6", [{"role": "user", "content": "x"}], { "temperature": 1.0, "top_p": 0.95, "max_output_tokens": 1024, "reasoning_enabled": False, }, "k2.6", ) self.assertEqual(payload["reasoning"], {"enabled": False}) code_payload = build_payload( "endpoint", [], { "temperature": 1.0, "top_p": 0.95, "max_output_tokens": 1024, "reasoning_enabled": True, }, "k2.7-code", ) self.assertFalse(code_payload["chat_template_kwargs"]["clear_thinking"]) def test_sse(self): lines = [ b'data: {"choices":[{"delta":{"reasoning":"r","content":"c"}}]}\n', b"data: [DONE]\n", ] self.assertEqual( [item["type"] for item in parse_sse_lines(lines)], ["reasoning", "content", "done"], ) def test_password(self): value = hash_password("correct horse battery staple") self.assertTrue(verify_password("correct horse battery staple", value)) self.assertFalse(verify_password("wrong password here", value)) def test_database_schema_and_wal(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "db.sqlite" initialise(path) db = connect(path) self.assertEqual(db.execute("PRAGMA integrity_check").fetchone()[0], "ok") self.assertEqual(db.execute("PRAGMA journal_mode").fetchone()[0], "wal") tables = { row[0] for row in db.execute( "SELECT name FROM sqlite_master WHERE type='table'" ).fetchall() } self.assertTrue( { "files", "message_attachments", "jobs", "job_events", "memories", "skills", "audit_log", "sync_operations", "sync_changes", }.issubset(tables) ) db.close() def test_authenticated_artifact_encryption_round_trip(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) plain = root / "plain.bin" encrypted = root / "encrypted.kimuenc" body = b"KIMU encrypted artifact evidence" * 128 plain.write_bytes(body) key = parse_master_key(generate_master_key()) aad = file_aad("fil_test", "a" * 64, 1) encrypt_file(plain, encrypted, key, aad) self.assertNotIn(body[:20], encrypted.read_bytes()) self.assertEqual(decrypt_file_bytes(encrypted, key, aad), body) tampered = bytearray(encrypted.read_bytes()) tampered[-17] ^= 1 encrypted.write_bytes(tampered) with self.assertRaises(EncryptionError): decrypt_file_bytes(encrypted, key, aad) def test_legacy_plaintext_file_migrates_to_encrypted_store(self): import hashlib with tempfile.TemporaryDirectory() as directory: root = Path(directory) database = root / "db.sqlite3" storage = root / "files" storage.mkdir() initialise(database) db = connect(database) db.execute( "INSERT INTO users(id,username,password_hash,role) VALUES('usr_test','owner','x','owner')" ) body = b"legacy plaintext artifact" digest = hashlib.sha256(body).hexdigest() stored_name = "legacy.bin" (storage / stored_name).write_bytes(body) db.execute( """ INSERT INTO files( id,user_id,original_name,stored_name,mime_type,size_bytes,sha256, kind,status,text_content,zip_inventory_json,warnings_json ) VALUES('fil_legacy','usr_test','legacy.txt',?,'text/plain',?,?, 'text','ready','preview text','[]','[]') """, (stored_name, len(body), digest), ) store = EncryptedArtifactStore(storage, parse_master_key(generate_master_key()), 1) result = store.migrate_legacy_files(db) row = db.execute("SELECT * FROM files WHERE id='fil_legacy'").fetchone() self.assertEqual(result["migrated"], 1) self.assertTrue(is_encrypted_file(storage / stored_name)) self.assertEqual(store.body_bytes(row), body) self.assertEqual(store.private_metadata(row)["text_content"], "preview text") self.assertEqual(row["text_content"], "") db.close() def test_text_and_safe_zip_inspection(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) text_file = root / "sample.py" text_file.write_text("print('hello')\n", encoding="utf-8") text = inspect_stored_file(text_file, "sample.py", "text/x-python") self.assertEqual(text.kind, "text") self.assertIn("hello", text.text_content) self.assertEqual(len(text.sha256), 64) archive = root / "project.zip" with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as handle: handle.writestr("src/main.py", "print('zip works')\n") handle.writestr("README.md", "# Project\n") inspected = inspect_stored_file(archive, "project.zip", "application/zip") self.assertEqual(inspected.kind, "zip") self.assertEqual(len(inspected.zip_inventory), 2) self.assertIn("ZIP member: src/main.py", inspected.text_content) def test_unsafe_zip_path_is_blocked(self): with tempfile.TemporaryDirectory() as directory: archive = Path(directory) / "unsafe.zip" with zipfile.ZipFile(archive, "w") as handle: handle.writestr("../escape.txt", "blocked") with self.assertRaises(FileValidationError) as context: inspect_stored_file(archive, "unsafe.zip", "application/zip") self.assertEqual(context.exception.code, "KIMU_ZIP_UNSAFE_PATH") def test_attachment_blocks_and_filename_normalisation(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "notes.txt" path.write_text("important context", encoding="utf-8") row = { "kind": "text", "size_bytes": path.stat().st_size, "mime_type": "text/plain", "text_content": "important context", "zip_inventory_json": "[]", "original_name": "notes.txt", "sha256": "a" * 64, } block = attachment_content(row, path) self.assertEqual(block["type"], "text") self.assertIn("important context", block["text"]) self.assertNotIn("..", safe_filename("../../unsafe name?.txt")) def test_fake_code_interpreter(self): result = fake_execute("print(1)") self.assertEqual(result["status"], "completed") self.assertTrue(result["session_id"]) self.assertIn("Received 8 characters", result["outputs"][0]["data"]) if __name__ == "__main__": unittest.main()