kenyonxu-ComfyUI_KaiUtilities/tests/test_civitai_client.py

import hashlib
import os
import sys
import tempfile
from unittest.mock import MagicMock, patch, AsyncMock

import aiohttp
import pytest


class AsyncContextManagerMock:
    """Helper to mock an async context manager (e.g. `async with session.get() as resp:`)."""

    def __init__(self, enter_result):
        self._enter_result = enter_result

    async def __aenter__(self):
        return self._enter_result

    async def __aexit__(self, *args):
        pass


class AsyncIteratorMock:
    """Helper to mock an async iterator (e.g. `async for chunk in ...`)."""

    def __init__(self, items):
        self._items = iter(items)

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            return next(self._items)
        except StopIteration:
            raise StopAsyncIteration

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from civitai_client import (
    compute_file_hash,
    compute_file_hash_sync,
    _hash_cache,
    fetch_model_by_hash,
    download_image,
    get_session,
    CIVITAI_API_URL,
)


class TestComputeFileHash:
    @pytest.fixture
    def known_file(self):
        with tempfile.NamedTemporaryFile(mode="wb", delete=False) as f:
            f.write(b"hello world")
            path = f.name
        yield path
        os.remove(path)

    @pytest.mark.asyncio
    async def test_compute_file_hash_known_content(self, known_file):
        expected = hashlib.sha256(b"hello world").hexdigest()
        result = await compute_file_hash(known_file)
        assert result == expected

    @pytest.mark.asyncio
    async def test_compute_file_hash_caches_in_memory(self, known_file):
        _hash_cache.clear()
        result1 = await compute_file_hash(known_file)
        result2 = await compute_file_hash(known_file)
        assert result1 == result2
        # Check that some cache key containing the path exists
        assert any(known_file in k for k in _hash_cache)

    @pytest.mark.asyncio
    async def test_compute_file_hash_missing_file_raises(self):
        with pytest.raises(FileNotFoundError):
            await compute_file_hash("/nonexistent/path/file.safetensors")


class TestComputeFileHashSync:
    @pytest.fixture
    def known_file(self):
        with tempfile.NamedTemporaryFile(mode="wb", delete=False) as f:
            f.write(b"sync test")
            path = f.name
        yield path
        os.remove(path)

    def test_compute_file_hash_sync_known_content(self, known_file):
        expected = hashlib.sha256(b"sync test").hexdigest()
        _hash_cache.clear()
        result = compute_file_hash_sync(known_file)
        assert result == expected

    def test_compute_file_hash_sync_uses_cache(self, known_file):
        _hash_cache.clear()
        result1 = compute_file_hash_sync(known_file)
        result2 = compute_file_hash_sync(known_file)
        assert result1 == result2
        # Check that some cache key containing the path exists
        assert any(known_file in k for k in _hash_cache)

    def test_compute_file_hash_sync_missing_file_raises(self):
        with pytest.raises(FileNotFoundError):
            compute_file_hash_sync("/nonexistent/path/file.safetensors")


class TestFetchModelByHash:
    @pytest.mark.asyncio
    async def test_fetch_model_by_hash_success(self):
        mock_response = MagicMock()
        mock_response.raise_for_status = MagicMock()
        mock_response.json = AsyncMock(return_value={"id": 1, "name": "Test Model"})

        mock_session = MagicMock()
        mock_session.get = MagicMock(return_value=AsyncContextManagerMock(mock_response))

        with patch("civitai_client.get_session", return_value=mock_session):
            result = await fetch_model_by_hash("abc123")
            assert result["name"] == "Test Model"

    @pytest.mark.asyncio
    async def test_fetch_model_by_hash_404(self):
        mock_response = MagicMock()
        mock_response.raise_for_status.side_effect = aiohttp.ClientResponseError(
            request_info=MagicMock(),
            history=(),
            status=404,
        )

        mock_session = MagicMock()
        mock_session.get = MagicMock(return_value=AsyncContextManagerMock(mock_response))

        with patch("civitai_client.get_session", return_value=mock_session):
            with pytest.raises(aiohttp.ClientResponseError):
                await fetch_model_by_hash("notfound")


class TestDownloadImage:
    @pytest.mark.asyncio
    async def test_download_image_success(self):
        mock_response = MagicMock()
        mock_response.status = 200
        mock_response.content.iter_chunked = MagicMock(return_value=AsyncIteratorMock([b"chunk1", b"chunk2"]))

        mock_session = MagicMock()
        mock_session.get = MagicMock(return_value=AsyncContextManagerMock(mock_response))

        with patch("civitai_client.get_session", return_value=mock_session):
            with tempfile.NamedTemporaryFile(delete=False) as f:
                dest = f.name
            try:
                result = await download_image("https://example.com/img.jpg", dest)
                assert result is True
                with open(dest, "rb") as f:
                    assert f.read() == b"chunk1chunk2"
            finally:
                if os.path.exists(dest):
                    os.remove(dest)

    @pytest.mark.asyncio
    async def test_download_image_404(self):
        mock_response = MagicMock()
        mock_response.status = 404

        mock_session = MagicMock()
        mock_session.get = MagicMock(return_value=AsyncContextManagerMock(mock_response))

        with patch("civitai_client.get_session", return_value=mock_session):
            result = await download_image("https://example.com/img.jpg", "/tmp/test.jpg")
            assert result is False

    @pytest.mark.asyncio
    async def test_download_image_network_error(self):
        mock_session = MagicMock()
        mock_session.get = MagicMock(side_effect=aiohttp.ClientError("network error"))

        with patch("civitai_client.get_session", return_value=mock_session):
            result = await download_image("https://example.com/img.jpg", "/tmp/test.jpg")
            assert result is False