ImpTune/tests/test_db.py

"""Tests for database initialization and driver storage."""
import hashlib
from pathlib import Path

import pytest


@pytest.fixture
def db_env(tmp_path, monkeypatch):
    """Set up temp DATA_DIR and configure db to use temp paths."""
    data_dir = tmp_path / "data"
    data_dir.mkdir()
    drivers_dir = data_dir / "drivers"
    drivers_dir.mkdir()
    db_path = str(data_dir / "imptune.db")

    monkeypatch.setenv("DATA_DIR", str(data_dir))

    import imptune.config as cfg
    cfg.DATA_DIR = str(data_dir)
    cfg.DB_PATH = db_path
    cfg.DRIVERS_DIR = str(drivers_dir)

    # Close and re-init the db with the temp path
    from imptune.db.database import db
    if not db.is_closed():
        db.close()

    yield {
        "data_dir": data_dir,
        "drivers_dir": drivers_dir,
        "db_path": db_path,
    }

    # A connection left open here stays bound to *this* test's file, and the next
    # test's ORM writes would land in it instead of its own tmp DB — the app
    # thread would then 404 on rows the test just created.
    if not db.is_closed():
        db.close()


def test_create_tables(db_env):
    """init_db() creates all 4 tables in a fresh SQLite file."""
    from imptune.db.database import db, init_db

    if not db.is_closed():
        db.close()

    init_db()

    tables = db.get_tables()
    assert "owner" in tables
    assert "client" in tables
    assert "driver" in tables
    assert "printer" in tables
    assert "icon" in tables

    db.close()


def test_wal_mode(db_env):
    """After init_db(), PRAGMA journal_mode returns 'wal'."""
    from imptune.db.database import db, init_db

    if not db.is_closed():
        db.close()

    init_db()

    cursor = db.execute_sql("PRAGMA journal_mode;")
    mode = cursor.fetchone()[0]
    assert mode == "wal"

    db.close()


def test_foreign_keys(db_env):
    """After init_db(), PRAGMA foreign_keys returns 1."""
    from imptune.db.database import db, init_db

    if not db.is_closed():
        db.close()

    init_db()

    cursor = db.execute_sql("PRAGMA foreign_keys;")
    value = cursor.fetchone()[0]
    assert value == 1

    db.close()


def test_idempotent(db_env):
    """Calling init_db() twice does not raise an error."""
    from imptune.db.database import db, init_db

    if not db.is_closed():
        db.close()

    init_db()
    db.close()
    init_db()  # second call — must not raise

    db.close()


def test_driver_store_save(db_env, tmp_path):
    """Saving bytes returns their SHA256 hex digest and creates the file."""
    from imptune.storage.driver_store import DriverStore

    drivers_dir = db_env["drivers_dir"]
    store = DriverStore(str(drivers_dir))

    data = b"test driver package content"
    expected_sha256 = hashlib.sha256(data).hexdigest()

    result = store.save(data)

    assert result == expected_sha256
    assert (drivers_dir / f"{expected_sha256}.zip").exists()


def test_driver_store_dedup(db_env):
    """Saving the same bytes twice results in one file on disk."""
    from imptune.storage.driver_store import DriverStore

    drivers_dir = db_env["drivers_dir"]
    store = DriverStore(str(drivers_dir))

    data = b"duplicate driver data"
    store.save(data)
    store.save(data)

    files = list(drivers_dir.iterdir())
    assert len(files) == 1


def test_driver_store_get_path(db_env):
    """get_path(sha256) returns the correct file path."""
    from imptune.storage.driver_store import DriverStore

    drivers_dir = db_env["drivers_dir"]
    store = DriverStore(str(drivers_dir))

    sha256 = "abcdef1234567890" * 4  # 64 hex chars
    path = store.get_path(sha256)

    assert path == Path(str(drivers_dir)) / f"{sha256}.zip"


def test_init_db_adds_display_name_to_a_preexisting_driver_table(db_env):
    """A DB created before the rename feature gains the column, keeping its rows.

    `create_tables(safe=True)` skips a table that already exists, so a new field
    on an old model only lands through `_add_missing_columns`.
    """
    from imptune.db.database import db, init_db

    init_db()
    db.execute_sql("ALTER TABLE driver DROP COLUMN display_name")
    db.execute_sql(
        "INSERT INTO driver (sha256, original_filename, size_bytes, uploaded_at, "
        "has_cat_file) VALUES ('legacy', 'old.zip', 10, '2024-01-01 00:00:00', 0)"
    )
    db.close()

    init_db()

    columns = {row[1] for row in db.execute_sql("PRAGMA table_info(driver)")}
    assert "display_name" in columns

    from imptune.db.models import Driver

    legacy = Driver.get(Driver.sha256 == "legacy")
    assert legacy.display_name is None
    assert legacy.label == "old.zip"


def test_init_db_creates_the_driver_icon_table(db_env):
    from imptune.db.database import db, init_db

    init_db()
    tables = {
        row[0]
        for row in db.execute_sql("SELECT name FROM sqlite_master WHERE type='table'")
    }
    assert "driver_icon" in tables