import os
from unittest.mock import AsyncMock, patch


def test_auth_middleware_missing_auth(test_client):
    response = test_client.post("/api/generate", json={"model": "llama3", "prompt": "test"})
    assert response.status_code == 401


def test_auth_middleware_invalid_key(test_client):
    response = test_client.post(
        "/api/generate",
        headers={"Authorization": "sk-invalid-key"},
        json={"model": "llama3", "prompt": "test"},
    )
    assert response.status_code == 401


@patch("main.proxy_request", new_callable=AsyncMock)
def test_auth_middleware_valid_key(mock_proxy, test_client):
    mock_proxy.return_value.status_code = 200
    mock_proxy.return_value.json = lambda: {"response": "success"}
    mock_proxy.return_value.headers = {}

    response = test_client.post(
        "/api/generate",
        headers={"Authorization": os.environ.get("TEST_API_KEY", "")},
        json={"model": "llama3", "prompt": "test"},
    )
    assert response.status_code == 200