folding-studio/tests/utils/test_headers.py

"""Test headers"""
import pytest
from unittest import mock
from folding_studio.utils.headers import get_auth_headers


def test_get_auth_headers_api_key():
    result = get_auth_headers()
    assert result == {"X-API-Key": "MY_KEY"}

@pytest.fixture()
def mock_get_auth_headers():
    with mock.patch(
        "folding_studio.utils.headers.get_id_token", return_value="identity_token"
    ) as m:
        yield m


@mock.patch("folding_studio.utils.headers.FOLDING_API_KEY", None)
def test_get_auth_headers_token(mock_get_auth_headers):
    result = get_auth_headers()
    assert result == {"Authorization": "Bearer identity_token"}