TJWaterServerBinary/tests/unit/test_metadata_repository_dsn_decrypt.py

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
from uuid import uuid4

import pytest
from cryptography.fernet import InvalidToken

from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository


class _DummyResult:
    def __init__(self, record):
        self._record = record

    def scalar_one_or_none(self):
        return self._record


class _DummyEncryptor:
    def __init__(self, decrypted=None, raise_invalid_token=False):
        self._decrypted = decrypted
        self._raise_invalid_token = raise_invalid_token
        self.encrypted_values = []

    def encrypt(self, value):
        encrypted = f"encrypted::{value}"
        self.encrypted_values.append(value)
        return encrypted

    def decrypt(self, _value):
        if self._raise_invalid_token:
            raise InvalidToken()
        return self._decrypted

def _build_record(dsn_encrypted: str):
    return SimpleNamespace(
        project_id=uuid4(),
        db_role="biz_data",
        db_type="postgresql",
        dsn_encrypted=dsn_encrypted,
        pool_min_size=1,
        pool_max_size=5,
    )


def test_invalid_token_with_plaintext_dsn_value_raises_clear_error(monkeypatch):
    record = _build_record("postgresql://user:p@ss@localhost:5432/db")
    session = SimpleNamespace(
        execute=None,
        commit=None,
    )
    session.execute = AsyncMock(return_value=_DummyResult(record))
    session.commit = AsyncMock()
    encryptor = _DummyEncryptor(raise_invalid_token=True)
    repo = MetadataRepository(session)

    monkeypatch.setattr(
        "app.infra.db.metadb.repositories.metadata_repository.is_database_encryption_configured",
        lambda: True,
    )
    monkeypatch.setattr(
        "app.infra.db.metadb.repositories.metadata_repository.get_database_encryptor",
        lambda: encryptor,
    )

    with pytest.raises(
        ValueError,
        match="DATABASE_ENCRYPTION_KEY mismatch or invalid dsn_encrypted value",
    ):
        asyncio.run(repo.get_project_db_routing(record.project_id, "biz_data"))
    session.commit.assert_not_awaited()


def test_invalid_token_with_non_dsn_value_raises_clear_error(monkeypatch):
    record = _build_record("gAAAAABinvalidciphertext")
    session = SimpleNamespace(
        execute=None,
        commit=None,
    )
    session.execute = AsyncMock(return_value=_DummyResult(record))
    session.commit = AsyncMock()
    repo = MetadataRepository(session)

    monkeypatch.setattr(
        "app.infra.db.metadb.repositories.metadata_repository.is_database_encryption_configured",
        lambda: True,
    )
    monkeypatch.setattr(
        "app.infra.db.metadb.repositories.metadata_repository.get_database_encryptor",
        lambda: _DummyEncryptor(raise_invalid_token=True),
    )

    with pytest.raises(
        ValueError,
        match="DATABASE_ENCRYPTION_KEY mismatch or invalid dsn_encrypted value",
    ):
        asyncio.run(repo.get_project_db_routing(record.project_id, "biz_data"))
    session.commit.assert_not_awaited()


def test_encrypted_dsn_decrypts_without_migration(monkeypatch):
    record = _build_record("encrypted-value")
    session = SimpleNamespace(
        execute=None,
        commit=None,
    )
    session.execute = AsyncMock(return_value=_DummyResult(record))
    session.commit = AsyncMock()
    repo = MetadataRepository(session)

    monkeypatch.setattr(
        "app.infra.db.metadb.repositories.metadata_repository.is_database_encryption_configured",
        lambda: True,
    )
    monkeypatch.setattr(
        "app.infra.db.metadb.repositories.metadata_repository.get_database_encryptor",
        lambda: _DummyEncryptor(decrypted="postgresql://u:p@ss@host/db"),
    )

    routing = asyncio.run(repo.get_project_db_routing(record.project_id, "biz_data"))

    assert routing.dsn == "postgresql://u:p%40ss@host/db"
    session.commit.assert_not_awaited()


def test_upsert_project_database_config_encrypts_plaintext_dsn(monkeypatch):
    project_id = uuid4()
    session = SimpleNamespace(
        execute=None,
        add=None,
        commit=None,
        refresh=None,
    )
    added = []
    session.execute = AsyncMock(return_value=_DummyResult(None))
    session.add = lambda item: added.append(item)
    session.commit = AsyncMock()
    session.refresh = AsyncMock()
    encryptor = _DummyEncryptor()
    repo = MetadataRepository(session)

    monkeypatch.setattr(
        "app.infra.db.metadb.repositories.metadata_repository.is_database_encryption_configured",
        lambda: True,
    )
    monkeypatch.setattr(
        "app.infra.db.metadb.repositories.metadata_repository.get_database_encryptor",
        lambda: encryptor,
    )

    record = asyncio.run(
        repo.upsert_project_database_config(
            project_id,
            db_role="biz_data",
            db_type="postgresql",
            dsn="postgresql://user:secret@localhost/db",
            pool_min_size=1,
            pool_max_size=5,
        )
    )

    assert encryptor.encrypted_values == ["postgresql://user:secret@localhost/db"]
    assert record.dsn_encrypted == "encrypted::postgresql://user:secret@localhost/db"
    assert added == [record]
    session.commit.assert_awaited_once()
    session.refresh.assert_awaited_once_with(record)