ComfyUI-Lora-Manager/tests/services/test_downloaded_version_history_service.py

from pathlib import Path

import pytest

from py.services.downloaded_version_history_service import (
    DownloadedVersionHistoryService,
)


class DummySettings:
    def get_active_library_name(self) -> str:
        return "alpha"


@pytest.mark.asyncio
async def test_download_history_roundtrip_and_manual_override(tmp_path: Path) -> None:
    db_path = tmp_path / "download-history.sqlite"
    service = DownloadedVersionHistoryService(
        str(db_path),
        settings_manager=DummySettings(),
    )

    await service.mark_downloaded(
        "lora",
        101,
        model_id=11,
        source="scan",
        file_path="/models/a.safetensors",
    )
    assert await service.has_been_downloaded("lora", 101) is True
    assert await service.get_downloaded_version_ids("lora", 11) == [101]

    await service.mark_as_deleted("lora", 101)
    assert await service.has_been_downloaded("lora", 101) is False
    assert await service.get_downloaded_version_ids("lora", 11) == []

    await service.mark_downloaded(
        "lora",
        101,
        model_id=11,
        source="download",
        file_path="/models/a.safetensors",
    )
    assert await service.has_been_downloaded("lora", 101) is True
    assert await service.get_downloaded_version_ids("lora", 11) == [101]


@pytest.mark.asyncio
async def test_download_history_bulk_lookup(tmp_path: Path) -> None:
    db_path = tmp_path / "download-history.sqlite"
    service = DownloadedVersionHistoryService(
        str(db_path),
        settings_manager=DummySettings(),
    )

    await service.mark_downloaded_bulk(
        "checkpoint",
        [
            {"model_id": 5, "version_id": 501, "file_path": "/m/one.safetensors"},
            {"model_id": 5, "version_id": 502, "file_path": "/m/two.safetensors"},
            {"model_id": 6, "version_id": 601, "file_path": "/m/three.safetensors"},
        ],
        source="scan",
    )

    assert await service.get_downloaded_version_ids("checkpoint", 5) == [501, 502]
    assert await service.get_downloaded_version_ids_bulk("checkpoint", [5, 6, 7]) == {
        5: {501, 502},
        6: {601},
    }


@pytest.mark.asyncio
async def test_per_file_history_tracking(tmp_path: Path) -> None:
    """Per-file records coexist with the version-level row (#1058)."""
    db_path = tmp_path / "download-history.sqlite"
    service = DownloadedVersionHistoryService(
        str(db_path),
        settings_manager=DummySettings(),
    )

    await service.mark_downloaded(
        "lora", 101, model_id=11, source="download",
        file_path="/models/a.safetensors", file_id=1001, file_name="a.safetensors",
    )
    await service.mark_downloaded(
        "lora", 101, model_id=11, source="download",
        file_path="/models/b.safetensors", file_id=1002, file_name="b.safetensors",
    )

    assert await service.get_downloaded_file_ids("lora", 101) == [1001, 1002]
    # Version-level tracking remains single-row per version
    assert await service.get_downloaded_version_ids("lora", 11) == [101]

    # Re-downloading the same file updates in place, no duplicate
    await service.mark_downloaded(
        "lora", 101, source="download", file_id=1001, file_name="a.safetensors",
    )
    assert await service.get_downloaded_file_ids("lora", 101) == [1001, 1002]

    # Single-file deletion keeps the sibling record
    await service.mark_file_deleted("lora", 101, 1001)
    assert await service.get_downloaded_file_ids("lora", 101) == [1002]

    # Whole-version deletion clears per-file records
    await service.mark_as_deleted("lora", 101)
    assert await service.get_downloaded_file_ids("lora", 101) == []
    assert await service.has_been_downloaded("lora", 101) is False


@pytest.mark.asyncio
async def test_per_file_history_ignores_invalid_ids(tmp_path: Path) -> None:
    service = DownloadedVersionHistoryService(
        str(tmp_path / "download-history.sqlite"),
        settings_manager=DummySettings(),
    )

    # mark_downloaded without a file id only touches the version-level table
    await service.mark_downloaded("lora", 201, model_id=21, source="scan")
    assert await service.get_downloaded_file_ids("lora", 201) == []

    # Invalid inputs are no-ops
    await service.mark_file_deleted("lora", 201, None)  # type: ignore[arg-type]
    assert await service.get_downloaded_file_ids("unknown-type", 201) == []


@pytest.mark.asyncio
async def test_file_history_table_created_for_legacy_db(tmp_path: Path) -> None:
    """Existing databases gain the per-file table via CREATE IF NOT EXISTS."""
    import sqlite3

    db_path = tmp_path / "download-history.sqlite"
    conn = sqlite3.connect(db_path)
    conn.executescript(
        """
        CREATE TABLE downloaded_model_versions (
            model_type TEXT NOT NULL,
            version_id INTEGER NOT NULL,
            model_id INTEGER,
            first_seen_at REAL NOT NULL,
            last_seen_at REAL NOT NULL,
            source TEXT NOT NULL,
            last_file_path TEXT,
            last_library_name TEXT,
            is_deleted_override INTEGER NOT NULL DEFAULT 0,
            PRIMARY KEY (model_type, version_id)
        );
        """
    )
    conn.close()

    service = DownloadedVersionHistoryService(
        str(db_path),
        settings_manager=DummySettings(),
    )
    await service.mark_downloaded(
        "lora", 301, model_id=31, source="download",
        file_id=9001, file_name="file.safetensors",
    )
    assert await service.get_downloaded_file_ids("lora", 301) == [9001]
    assert await service.has_been_downloaded("lora", 301) is True