argus-netconf-exporter/tests/test_sqlite_vendor_column.py

import sqlite3
from pathlib import Path

import pytest

from exporter.config import DeviceConfig
from exporter.sqlite_store import PasswordEncryptor, SQLiteDeviceStore


@pytest.fixture
def encryptor() -> PasswordEncryptor:
    # 生成一个有效 Fernet key
    from cryptography.fernet import Fernet

    key = Fernet.generate_key().decode()
    return PasswordEncryptor(key)


def test_init_db_creates_vendor_column_on_fresh_db(tmp_path: Path, encryptor: PasswordEncryptor):
    db_path = tmp_path / "test_vendor.db"
    store = SQLiteDeviceStore(str(db_path), encryptor)
    store.init_db()

    conn = sqlite3.connect(str(db_path))
    cols = [row[1] for row in conn.execute("PRAGMA table_info(devices)").fetchall()]
    conn.close()

    assert "vendor" in cols


def test_init_db_alter_table_vendor_preserves_existing_rows(tmp_path: Path, encryptor: PasswordEncryptor):
    db_path = tmp_path / "legacy.db"
    conn = sqlite3.connect(str(db_path))
    # 创建旧版本 devices 表（无 vendor 列）
    conn.execute(
        """
        CREATE TABLE devices (
          name TEXT PRIMARY KEY,
          host TEXT NOT NULL,
          port INTEGER NOT NULL,
          username TEXT NOT NULL,
          password_cipher BLOB NOT NULL,
          enabled INTEGER NOT NULL,
          scrape_interval_seconds INTEGER,
          supports_xpath INTEGER NOT NULL DEFAULT 0,
          created_at INTEGER NOT NULL,
          updated_at INTEGER NOT NULL
        );
        """
    )
    conn.execute(
        "INSERT INTO devices (name, host, port, username, password_cipher, enabled, "
        "scrape_interval_seconds, supports_xpath, created_at, updated_at) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        ("old-dev", "h", 830, "u", b"cipher", 1, None, 0, 1, 1),
    )
    conn.commit()
    conn.close()

    store = SQLiteDeviceStore(str(db_path), encryptor)
    store.init_db()

    conn2 = sqlite3.connect(str(db_path))
    row = conn2.execute("SELECT name, vendor FROM devices WHERE name = 'old-dev'").fetchone()
    conn2.close()

    assert row is not None
    assert row[0] == "old-dev"
    # 旧数据 vendor 应为空
    assert row[1] is None


def test_save_and_load_device_persists_vendor(tmp_path: Path, encryptor: PasswordEncryptor):
    db_path = tmp_path / "vendor_persist.db"
    store = SQLiteDeviceStore(str(db_path), encryptor)
    store.init_db()

    dev = DeviceConfig(
        name="dev-vendor",
        host="h",
        port=830,
        username="u",
        password="p",
        enabled=True,
        vendor="ruijie",
        source="runtime",
    )
    store.save_device(dev)
    loaded = store.load_runtime_devices()

    assert len(loaded) == 1
    assert loaded[0].name == "dev-vendor"
    assert loaded[0].vendor == "ruijie"


def test_save_and_load_device_vendor_none_roundtrip(tmp_path: Path, encryptor: PasswordEncryptor):
    db_path = tmp_path / "vendor_none.db"
    store = SQLiteDeviceStore(str(db_path), encryptor)
    store.init_db()

    dev = DeviceConfig(
        name="dev-no-vendor",
        host="h",
        port=830,
        username="u",
        password="p",
        enabled=True,
        vendor=None,
        source="runtime",
    )
    store.save_device(dev)
    loaded = store.load_runtime_devices()

    assert len(loaded) == 1
    assert loaded[0].name == "dev-no-vendor"
    assert loaded[0].vendor is None


def test_init_db_alter_table_silently_ignores_duplicate_column_error(tmp_path: Path, encryptor: PasswordEncryptor):
    db_path = tmp_path / "dup_col.db"
    store = SQLiteDeviceStore(str(db_path), encryptor)
    # 第一次初始化，创建带 vendor 列的表
    store.init_db()
    # 第二次调用，不应抛异常
    store.init_db()


def test_init_db_alter_table_raises_on_non_duplicate_errors(tmp_path: Path, encryptor: PasswordEncryptor, monkeypatch):
    db_path = tmp_path / "locked.db"

    real_connect = sqlite3.connect

    class ConnWrapper:
        def __init__(self, inner: sqlite3.Connection) -> None:
            self._inner = inner
            self._alter_attempted = False

        def execute(self, sql: str, *args, **kwargs):
            # 在第一次尝试 ALTER TABLE 时注入错误
            if "ALTER TABLE devices ADD COLUMN vendor" in sql and not self._alter_attempted:
                self._alter_attempted = True
                raise sqlite3.OperationalError("database is locked")
            return self._inner.execute(sql, *args, **kwargs)

        def commit(self) -> None:
            return self._inner.commit()

        def close(self) -> None:
            return self._inner.close()

    def wrapped_connect(*args, **kwargs):
        conn = real_connect(*args, **kwargs)
        return ConnWrapper(conn)

    monkeypatch.setattr(sqlite3, "connect", wrapped_connect)

    store = SQLiteDeviceStore(str(db_path), encryptor)
    with pytest.raises(sqlite3.OperationalError):
        store.init_db()