imageforge/backend/tests/test_a1111_provider.py

import base64

from backend.app.providers.a1111_provider import A1111Provider


class _FakeResponse:
    def __init__(self, payload: bytes) -> None:
        self._payload = payload

    def read(self) -> bytes:
        return self._payload

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


def test_a1111_provider_builds_auth_header_from_compact_auth(monkeypatch) -> None:
    monkeypatch.setattr("backend.app.providers.a1111_provider.A1111_API_AUTH", "user:secret")
    monkeypatch.setattr("backend.app.providers.a1111_provider.A1111_API_USER", "")
    monkeypatch.setattr("backend.app.providers.a1111_provider.A1111_API_PASSWORD", "")

    provider = A1111Provider()

    assert provider.auth_header.startswith("Basic ")
    encoded = provider.auth_header.split(" ", maxsplit=1)[1]
    decoded = base64.b64decode(encoded.encode("ascii")).decode("utf-8")
    assert decoded == "user:secret"


def test_a1111_provider_retries_transient_failure(monkeypatch) -> None:
    monkeypatch.setattr("backend.app.providers.a1111_provider.A1111_RETRY_COUNT", 2)
    monkeypatch.setattr("backend.app.providers.a1111_provider.A1111_RETRY_BACKOFF_SECONDS", 0.0)

    attempts = {"count": 0}

    def fake_urlopen(req, timeout):  # noqa: ANN001
        attempts["count"] += 1
        if attempts["count"] == 1:
            raise OSError("temporary network issue")
        return _FakeResponse(b'{"ok": true}')

    monkeypatch.setattr("backend.app.providers.a1111_provider.request.urlopen", fake_urlopen)

    provider = A1111Provider()
    result = provider._request_json("GET", "/sdapi/v1/sd-models")

    assert result == {"ok": True}
    assert attempts["count"] == 2