mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-21 00:21:38 +00:00
* Add ChatGPT subscription support and responses bridge * Fix typing import for responses bridge * Guard device code timestamp parsing * add /v1/messages endpoint to chatgpt model
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
import base64
|
|
import json
|
|
import time
|
|
from unittest.mock import mock_open, patch
|
|
|
|
import pytest
|
|
|
|
from litellm.llms.chatgpt.authenticator import Authenticator
|
|
|
|
|
|
def _make_jwt(payload: dict) -> str:
|
|
header = {"alg": "none", "typ": "JWT"}
|
|
|
|
def _b64(obj: dict) -> str:
|
|
raw = json.dumps(obj, separators=(",", ":")).encode("utf-8")
|
|
return base64.urlsafe_b64encode(raw).decode("utf-8").rstrip("=")
|
|
|
|
return f"{_b64(header)}.{_b64(payload)}."
|
|
|
|
|
|
class TestChatGPTAuthenticator:
|
|
@pytest.fixture
|
|
def authenticator(self):
|
|
with patch("os.path.exists", return_value=True):
|
|
return Authenticator()
|
|
|
|
def test_get_access_token_from_file(self, authenticator):
|
|
future_time = time.time() + 3600
|
|
auth_data = json.dumps({"access_token": "token-123", "expires_at": future_time})
|
|
|
|
with patch("builtins.open", mock_open(read_data=auth_data)):
|
|
token = authenticator.get_access_token()
|
|
assert token == "token-123"
|
|
|
|
def test_get_access_token_refresh(self, authenticator):
|
|
past_time = time.time() - 10
|
|
auth_data = json.dumps(
|
|
{
|
|
"access_token": "token-old",
|
|
"refresh_token": "refresh-123",
|
|
"expires_at": past_time,
|
|
}
|
|
)
|
|
refreshed = {
|
|
"access_token": "token-new",
|
|
"refresh_token": "refresh-123",
|
|
"id_token": "id-123",
|
|
}
|
|
|
|
with patch("builtins.open", mock_open(read_data=auth_data)), patch.object(
|
|
authenticator, "_refresh_tokens", return_value=refreshed
|
|
):
|
|
token = authenticator.get_access_token()
|
|
assert token == "token-new"
|
|
|
|
def test_get_account_id_from_id_token(self, authenticator):
|
|
id_token = _make_jwt(
|
|
{"https://api.openai.com/auth": {"chatgpt_account_id": "acct-123"}}
|
|
)
|
|
auth_data = json.dumps({"id_token": id_token})
|
|
|
|
with patch("builtins.open", mock_open(read_data=auth_data)), patch.object(
|
|
authenticator, "_write_auth_file"
|
|
) as mock_write:
|
|
account_id = authenticator.get_account_id()
|
|
assert account_id == "acct-123"
|
|
mock_write.assert_called_once()
|
|
assert mock_write.call_args[0][0]["account_id"] == "acct-123"
|