Spaces:
Configuration error
Configuration error
File size: 2,245 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import datetime
import json
import os
import sys
import unittest
from typing import List, Optional, Tuple
from unittest.mock import ANY, MagicMock, Mock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system-path
import litellm
@pytest.mark.asyncio
async def test_construct_request_headers_project_id_from_env(monkeypatch):
"""Test that construct_request_headers uses GCS_PUBSUB_PROJECT_ID environment variable."""
from litellm.integrations.gcs_pubsub.pub_sub import GcsPubSubLogger
# Set up test environment variable
test_project_id = "test-project-123"
monkeypatch.setenv("GCS_PUBSUB_PROJECT_ID", test_project_id)
monkeypatch.setattr(
"litellm.proxy.proxy_server.premium_user",
True,
)
try:
# Create handler with no project_id
handler = GcsPubSubLogger(
topic_id="test-topic", credentials_path="test-path.json"
)
# Mock the Vertex AI auth calls
mock_auth_header = "mock-auth-header"
mock_token = "mock-token"
with patch(
"litellm.vertex_chat_completion._ensure_access_token_async"
) as mock_ensure_token:
mock_ensure_token.return_value = (mock_auth_header, test_project_id)
with patch(
"litellm.vertex_chat_completion._get_token_and_url"
) as mock_get_token:
mock_get_token.return_value = (mock_token, "mock-url")
# Call construct_request_headers
headers = await handler.construct_request_headers()
# Verify headers
assert headers == {
"Authorization": f"Bearer {mock_token}",
"Content-Type": "application/json",
}
# Verify _ensure_access_token_async was called with correct project_id
mock_ensure_token.assert_called_once_with(
credentials="test-path.json",
project_id=test_project_id,
custom_llm_provider="vertex_ai",
)
finally:
# Clean up environment variable
del os.environ["GCS_PUBSUB_PROJECT_ID"]
|