"""Tests for Step 6 — memory ORM models and User.encryption_key. Uses the SQLite in-memory test DB (from conftest). The pgvector embedding column is stored as JSON in tests (SQLite-compatible). """ from __future__ import annotations import uuid from datetime import datetime, timezone import pytest import pytest_asyncio from cryptography.fernet import Fernet from sqlalchemy import select from app.models import MemoryAssociative, MemoryCore, MemoryEpisodic, MemoryProactive, User from tests.conftest import TEST_USER_IDS USER_ID = TEST_USER_IDS["power"] # ── helpers ─────────────────────────────────────────────────────────────────── def _fernet_key() -> str: return Fernet.generate_key().decode() def _encrypt(key: str, plaintext: str) -> str: return Fernet(key.encode()).encrypt(plaintext.encode()).decode() def _decrypt(key: str, ciphertext: str) -> str: return Fernet(key.encode()).decrypt(ciphertext.encode()).decode() # ── User.encryption_key ─────────────────────────────────────────────────────── @pytest.mark.asyncio async def test_user_encryption_key_column_exists(db_session): """User model has encryption_key column and it can be set.""" result = await db_session.execute(select(User).where(User.id == USER_ID)) user = result.scalar_one() # Column exists (may be None for seeded users) assert hasattr(user, "encryption_key") @pytest.mark.asyncio async def test_user_encryption_key_can_be_set(db_session): key = _fernet_key() result = await db_session.execute(select(User).where(User.id == USER_ID)) user = result.scalar_one() user.encryption_key = key await db_session.commit() result2 = await db_session.execute(select(User).where(User.id == USER_ID)) user2 = result2.scalar_one() assert user2.encryption_key == key # ── MemoryCore ──────────────────────────────────────────────────────────────── @pytest.mark.asyncio async def test_memory_core_create_and_read(db_session): key = _fernet_key() encrypted_val = _encrypt(key, "UTC") row = MemoryCore( id=str(uuid.uuid4()), user_id=USER_ID, key="timezone", value_encrypted=encrypted_val, ) db_session.add(row) await db_session.commit() result = await db_session.execute( select(MemoryCore).where(MemoryCore.user_id == USER_ID) ) fetched = result.scalar_one() assert fetched.key == "timezone" assert _decrypt(key, fetched.value_encrypted) == "UTC" @pytest.mark.asyncio async def test_memory_core_cascade_delete(db_session): """Deleting a user cascades to memory_core.""" row = MemoryCore( id=str(uuid.uuid4()), user_id=USER_ID, key="lang", value_encrypted="enc", ) db_session.add(row) await db_session.commit() user = (await db_session.execute(select(User).where(User.id == USER_ID))).scalar_one() await db_session.delete(user) await db_session.commit() remaining = ( await db_session.execute(select(MemoryCore).where(MemoryCore.user_id == USER_ID)) ).scalars().all() assert remaining == [] # ── MemoryAssociative ───────────────────────────────────────────────────────── @pytest.mark.asyncio async def test_memory_associative_create_and_read(db_session): key = _fernet_key() content = _encrypt(key, "User prefers morning meetings") embedding = [0.1] * 1536 # fake embedding row = MemoryAssociative( id=str(uuid.uuid4()), user_id=USER_ID, content_encrypted=content, embedding=embedding, entity_type="preference", entity_id=None, ) db_session.add(row) await db_session.commit() result = await db_session.execute( select(MemoryAssociative).where(MemoryAssociative.user_id == USER_ID) ) fetched = result.scalar_one() assert fetched.entity_type == "preference" assert _decrypt(key, fetched.content_encrypted) == "User prefers morning meetings" assert len(fetched.embedding) == 1536 # ── MemoryEpisodic ──────────────────────────────────────────────────────────── @pytest.mark.asyncio async def test_memory_episodic_create_and_read(db_session): key = _fernet_key() session_id = str(uuid.uuid4()) summary = _encrypt(key, "User asked about Q1 tasks") row = MemoryEpisodic( id=str(uuid.uuid4()), user_id=USER_ID, summary_encrypted=summary, session_id=session_id, ) db_session.add(row) await db_session.commit() result = await db_session.execute( select(MemoryEpisodic).where(MemoryEpisodic.session_id == session_id) ) fetched = result.scalar_one() assert _decrypt(key, fetched.summary_encrypted) == "User asked about Q1 tasks" assert isinstance(fetched.created_at, datetime) # ── MemoryProactive ─────────────────────────────────────────────────────────── @pytest.mark.asyncio async def test_memory_proactive_create_and_read(db_session): key = _fernet_key() pattern = _encrypt(key, "User always assigns tasks to self") row = MemoryProactive( id=str(uuid.uuid4()), user_id=USER_ID, pattern_encrypted=pattern, confidence=0.85, source="inferred", ) db_session.add(row) await db_session.commit() result = await db_session.execute( select(MemoryProactive).where(MemoryProactive.user_id == USER_ID) ) fetched = result.scalar_one() assert fetched.confidence == pytest.approx(0.85) assert fetched.source == "inferred" assert _decrypt(key, fetched.pattern_encrypted) == "User always assigns tasks to self" # ── Auth registration generates encryption_key ─────────────────────────────── def test_register_sets_encryption_key(client): """POST /api/v1/auth/register creates a user with a valid Fernet key.""" resp = client.post( "/api/v1/auth/register", json={"email": "newuser@test.com", "password": "testpassword123"}, ) assert resp.status_code == 201 # Fetch the newly created user via the access token token = resp.json()["access_token"] me_resp = client.get( "/api/v1/auth/me", headers={"Authorization": f"Bearer {token}"}, ) assert me_resp.status_code == 200 # We can't see encryption_key in the API response (not in UserProfile), # but we verify registration didn't crash — key generation is implicit.