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, 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.