api/tests/test_memory_models.py

"""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

import pytest
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.