vps-manager/tests/test_credential_migration.py

"""凭据中心化迁移测试（accounts.login_password → credentials）

模拟旧库（accounts 无 credential_id 列）→ 跑迁移 → 断言：
密文原样搬入、site 取平台显示名、账号关联回填、原字段置空、重跑幂等、
无密码账号不产生空条目。

运行：.venv/bin/pytest tests/test_credential_migration.py -v
"""

from datetime import datetime, timezone

import pytest
import sqlalchemy as sa
from cryptography.fernet import Fernet
from sqlmodel import Session, SQLModel, create_engine, select

import app.database as db
from app.core import crypto
from app.core.config import settings
from app.models.asset import AIAccount, Account, Asset, CloudflareDetail, DomainDetail, VPSDetail
from app.models.credential import Credential, LoginType
from app.models.provider import Provider
from app.models.ssl import SiteCert, Subdomain

# 旧结构 accounts 表：无 credential_id 列（迁移应自动补列）
OLD_ACCOUNTS_DDL = """
CREATE TABLE accounts (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name VARCHAR NOT NULL,
    platform VARCHAR,
    remark VARCHAR,
    login_user VARCHAR,
    login_password_encrypted VARCHAR,
    api_config_encrypted VARCHAR,
    last_synced_at DATETIME,
    created_at DATETIME NOT NULL
)
"""


@pytest.fixture()
def legacy_engine(tmp_path, monkeypatch):
    """旧结构测试库：新表齐备，但 accounts 缺 credential_id 列"""
    monkeypatch.setattr(settings, "MASTER_KEY", Fernet.generate_key().decode())
    crypto._get_fernet.cache_clear()  # 避免复用其它测试缓存的 Fernet 实例
    engine = create_engine(
        f"sqlite:///{tmp_path / 'test.db'}", connect_args={"check_same_thread": False}
    )
    # 除 Account 外按当前模型建表（含新 credentials 表；accounts 用旧 DDL）
    SQLModel.metadata.create_all(
        engine,
        tables=[
            Provider.__table__, Asset.__table__, VPSDetail.__table__, DomainDetail.__table__,
            AIAccount.__table__, CloudflareDetail.__table__, Subdomain.__table__,
            SiteCert.__table__, Credential.__table__,
        ],
    )
    with engine.begin() as conn:
        conn.execute(sa.text(OLD_ACCOUNTS_DDL))
    # 迁移函数引用模块级 assets_engine → 指向测试库
    monkeypatch.setattr(db, "assets_engine", engine)
    yield engine
    crypto._get_fernet.cache_clear()


def _seed(engine, pwd_aliyun: str, pwd_forum: str) -> None:
    """造迁移前现场：平台账号（有密码）/ 无平台账号（有密码）/ 无密码账号"""
    with Session(engine) as s:
        s.add(Provider(slug="aliyun", name="阿里云", category="vps"))
        s.commit()
    now = datetime.now(timezone.utc).replace(tzinfo=None)
    with engine.begin() as conn:
        conn.execute(
            sa.text(
                "INSERT INTO accounts (name, platform, login_user, login_password_encrypted, created_at) "
                "VALUES ('me@gmail.com', 'aliyun', 'me@gmail.com', :pwd, :now)"
            ),
            {"pwd": pwd_aliyun, "now": now},
        )
        # platform 为空串：site 应回退 '未分类'；无 login_user：username 回退 name
        conn.execute(
            sa.text(
                "INSERT INTO accounts (name, platform, login_user, login_password_encrypted, created_at) "
                "VALUES ('forum-user', '', NULL, :pwd, :now)"
            ),
            {"pwd": pwd_forum, "now": now},
        )
        # 无密码账号：不应产生凭据条目
        conn.execute(
            sa.text(
                "INSERT INTO accounts (name, platform, created_at) VALUES ('no-pwd', 'vultr', :now)"
            ),
            {"now": now},
        )


def test_migration_moves_passwords_to_vault(legacy_engine):
    engine = legacy_engine
    pwd_aliyun = crypto.encrypt("p@ss-阿里云")
    pwd_forum = crypto.encrypt("forum-pass-中文")
    _seed(engine, pwd_aliyun, pwd_forum)

    db._migrate_assets_db()

    with Session(engine) as s:
        creds = s.exec(select(Credential).order_by(Credential.id)).all()
        # 仅两条有密码的账号被迁移，无密码账号不产生空条目
        assert len(creds) == 2

        c1, c2 = creds
        # 密文原样搬入（不重新加密，避免中间态明文暴露），仍可解密还原
        assert c1.password_encrypted == pwd_aliyun
        assert crypto.decrypt(c1.password_encrypted) == "p@ss-阿里云"
        # site 取平台显示名（join providers），username 取 login_user
        assert c1.site == "阿里云"
        assert c1.username == "me@gmail.com"
        assert c1.login_type == LoginType.PASSWORD
        # 无 provider 匹配 + platform 空 → site 回退 '未分类'；username 回退 name
        assert c2.password_encrypted == pwd_forum
        assert c2.site == "未分类"
        assert c2.username == "forum-user"

        # 账号侧：credential_id 回填、原密码字段置空（唯一事实源，防双份漂移）
        accs = {a.name: a for a in s.exec(select(Account)).all()}
        assert accs["me@gmail.com"].credential_id == c1.id
        assert accs["me@gmail.com"].login_password_encrypted is None
        assert accs["forum-user"].credential_id == c2.id
        assert accs["forum-user"].login_password_encrypted is None
        assert accs["no-pwd"].credential_id is None


def test_migration_is_idempotent(legacy_engine):
    """重跑迁移不得重复建条目（幂等条件：credential_id 已非空则跳过）"""
    engine = legacy_engine
    _seed(engine, crypto.encrypt("pw1"), crypto.encrypt("pw2"))

    db._migrate_assets_db()
    db._migrate_assets_db()
    db._migrate_assets_db()

    with Session(engine) as s:
        assert len(s.exec(select(Credential)).all()) == 2


def test_migration_adds_credential_id_column(legacy_engine):
    """旧库缺 credential_id 列时迁移应自动补列"""
    engine = legacy_engine
    with engine.connect() as conn:
        cols = {c["name"] for c in sa.inspect(conn).get_columns("accounts")}
    assert "credential_id" not in cols  # 前置：确实是旧结构

    db._migrate_assets_db()

    with engine.connect() as conn:
        cols = {c["name"] for c in sa.inspect(conn).get_columns("accounts")}
        indexes = {ix["name"] for ix in sa.inspect(conn).get_indexes("accounts")}
    assert "credential_id" in cols

    db._migrate_indexes()
    with engine.connect() as conn:
        indexes = {ix["name"] for ix in sa.inspect(conn).get_indexes("accounts")}
    assert "ix_accounts_credential_id" in indexes