173 lines
4.9 KiB
Python
173 lines
4.9 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(REPO_ROOT / "opax-mcp"))
|
|
|
|
import server
|
|
|
|
|
|
class TestIdentityToken(unittest.TestCase):
|
|
def test_get_identity_token_returns_fetched_token(self):
|
|
with (
|
|
patch.dict(
|
|
os.environ,
|
|
{"OPAX_LOCAL_DEV_MODE": ""},
|
|
clear=False,
|
|
),
|
|
patch(
|
|
"server.google.auth.transport.requests.Request"
|
|
) as mock_request,
|
|
patch(
|
|
"server.google.oauth2.id_token.fetch_id_token",
|
|
return_value="mock-token",
|
|
) as mock_fetch,
|
|
):
|
|
token = server._get_identity_token(
|
|
"https://mock-agent.run.app"
|
|
)
|
|
|
|
self.assertEqual(token, "mock-token")
|
|
mock_fetch.assert_called_once_with(
|
|
mock_request.return_value,
|
|
"https://mock-agent.run.app",
|
|
)
|
|
|
|
def test_get_identity_token_raises_runtime_error_on_failure(self):
|
|
with (
|
|
patch.dict(
|
|
os.environ,
|
|
{"OPAX_LOCAL_DEV_MODE": ""},
|
|
clear=False,
|
|
),
|
|
patch(
|
|
"server.google.auth.transport.requests.Request"
|
|
),
|
|
patch(
|
|
"server.google.oauth2.id_token.fetch_id_token",
|
|
side_effect=Exception("token fetch failed"),
|
|
),
|
|
):
|
|
with self.assertRaises(RuntimeError):
|
|
server._get_identity_token(
|
|
"https://mock-agent.run.app"
|
|
)
|
|
|
|
def test_local_dev_mode_returns_empty_token_without_fetching(self):
|
|
with (
|
|
patch.dict(
|
|
os.environ,
|
|
{"OPAX_LOCAL_DEV_MODE": "true"},
|
|
clear=False,
|
|
),
|
|
patch(
|
|
"server.google.oauth2.id_token.fetch_id_token"
|
|
) as mock_fetch,
|
|
):
|
|
token = server._get_identity_token(
|
|
"https://mock-agent.run.app"
|
|
)
|
|
|
|
self.assertEqual(token, "")
|
|
mock_fetch.assert_not_called()
|
|
|
|
|
|
class TestAgentHeaders(unittest.IsolatedAsyncioTestCase):
|
|
async def test_agent_headers_include_bearer_and_internal_key(self):
|
|
with (
|
|
patch.object(
|
|
server,
|
|
"OSVAUCO_AGENT_URL",
|
|
"https://mock-agent.run.app",
|
|
),
|
|
patch.object(
|
|
server,
|
|
"INTERNAL_API_KEY",
|
|
"mock-internal-key",
|
|
),
|
|
patch(
|
|
"server.asyncio.to_thread",
|
|
new_callable=AsyncMock,
|
|
return_value="mock-token",
|
|
) as mock_to_thread,
|
|
):
|
|
headers = await server._agent_headers()
|
|
|
|
self.assertEqual(
|
|
headers["Authorization"],
|
|
"Bearer mock-token",
|
|
)
|
|
self.assertEqual(
|
|
headers["X-Internal-Key"],
|
|
"mock-internal-key",
|
|
)
|
|
self.assertEqual(
|
|
headers["Content-Type"],
|
|
"application/json",
|
|
)
|
|
mock_to_thread.assert_awaited_once_with(
|
|
server._get_identity_token,
|
|
"https://mock-agent.run.app",
|
|
)
|
|
|
|
async def test_local_dev_headers_have_no_authorization_header(self):
|
|
with (
|
|
patch.object(
|
|
server,
|
|
"OSVAUCO_AGENT_URL",
|
|
"https://mock-agent.run.app",
|
|
),
|
|
patch.object(
|
|
server,
|
|
"INTERNAL_API_KEY",
|
|
"mock-internal-key",
|
|
),
|
|
patch(
|
|
"server.asyncio.to_thread",
|
|
new_callable=AsyncMock,
|
|
return_value="",
|
|
),
|
|
):
|
|
headers = await server._agent_headers()
|
|
|
|
self.assertNotIn("Authorization", headers)
|
|
self.assertEqual(
|
|
headers["X-Internal-Key"],
|
|
"mock-internal-key",
|
|
)
|
|
self.assertEqual(
|
|
headers["Content-Type"],
|
|
"application/json",
|
|
)
|
|
|
|
async def test_agent_headers_fail_closed_when_token_fetch_fails(self):
|
|
with (
|
|
patch.object(
|
|
server,
|
|
"OSVAUCO_AGENT_URL",
|
|
"https://mock-agent.run.app",
|
|
),
|
|
patch.object(
|
|
server,
|
|
"INTERNAL_API_KEY",
|
|
"mock-internal-key",
|
|
),
|
|
patch(
|
|
"server.asyncio.to_thread",
|
|
new_callable=AsyncMock,
|
|
side_effect=RuntimeError("token fetch failed"),
|
|
),
|
|
):
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"token fetch failed",
|
|
):
|
|
await server._agent_headers()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |