from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from src.models.resolve_metrics_url import ResolveMetricsUrlResponse from src.services.observability_resolve import ( get_cluster_info, normalize_bearer_token, resolve_biz_type, resolve_metrics_url, ) from src.utils.errors import AuthServiceError def _mock_httpx_get_response(response: httpx.Response): mock_client = MagicMock() mock_client.get = AsyncMock(return_value=response) mock_cm = MagicMock() mock_cm.__aenter__ = AsyncMock(return_value=mock_client) mock_cm.__aexit__ = AsyncMock(return_value=None) return mock_cm class TestResolveBizType: def test_dedicated_prod_returns_tidbcloud(self): assert resolve_biz_type("dedicated", "prod") == "tidbcloud" def test_premium_prod_returns_prod_nextgen(self): assert resolve_biz_type("premium", "prod") == "prod-tidbcloud-nextgen" def test_resource_pool_dev_returns_dev_nextgen(self): assert resolve_biz_type("resource_pool", "dev") == "dev-tidbcloud-nextgen" def test_starter_default_env_returns_prod_devtier(self): assert resolve_biz_type("starter") == "prod-devtier" def test_empty_env_returns_base(self): assert resolve_biz_type("essential", "") == "prod-devtier" def test_unknown_deploy_type_raises(self): with pytest.raises(AuthServiceError): resolve_biz_type("unknown", "dev") class TestNormalizeBearerToken: def test_removes_bearer_prefix(self): assert normalize_bearer_token("Bearer abc") == "abc" def test_keeps_raw_token(self): assert normalize_bearer_token("abc") == "abc" class TestResolveMetricsUrl: @pytest.fixture(autouse=True) def _mock_gs_base_url(self): with patch( "src.services.observability_resolve.get_observability_global_service_url", return_value="https://gs.example.com", ): yield @pytest.mark.asyncio async def test_resolves_using_internal_read_uri_and_auth_header(self): payload = { "code": 0, "message": "", "data": { "clusters": [ { "internalReadUri": "https://ds.example.com/internal/metrics/abc/api/v1/", "openReadUri": "https://ds.example.com/external/metrics/abc/api/v1/", } ] }, } response = httpx.Response(200, json=payload) with patch( "src.services.observability_resolve.get_server_auth0_token", return_value="auth0-token" ): with patch( "src.services.observability_resolve.httpx.AsyncClient", return_value=_mock_httpx_get_response(response), ) as mock_client_cls: resolved = await resolve_metrics_url( cluster_id="cluster1", env="dev", deploy_type="starter", ) assert isinstance(resolved, ResolveMetricsUrlResponse) assert resolved.url == "https://ds.example.com/internal/metrics/abc/api/v1/" assert resolved.token == "auth0-token" get_call = mock_client_cls.return_value.__aenter__.return_value.get _, kwargs = get_call.call_args assert kwargs["params"]["bizType"] == "dev-devtier" assert kwargs["params"]["applicationID"] == "cluster1" assert "vendor" not in kwargs["params"] assert "region" not in kwargs["params"] assert kwargs["headers"]["Authorization"] == "Bearer auth0-token" @pytest.mark.asyncio async def test_falls_back_to_open_read_uri(self): payload = { "code": 0, "data": { "clusters": [ { "openReadUri": "https://ds.example.com/external/metrics/abc/api/v1/", } ] }, } response = httpx.Response(200, json=payload) with patch( "src.services.observability_resolve.get_server_auth0_token", return_value="Bearer auth0-token", ): with patch( "src.services.observability_resolve.httpx.AsyncClient", return_value=_mock_httpx_get_response(response), ): resolved = await resolve_metrics_url( cluster_id="cluster1", env="dev", deploy_type="starter", ) assert resolved.url == "https://ds.example.com/external/metrics/abc/api/v1/" assert resolved.token == "auth0-token" @pytest.mark.asyncio async def test_missing_auth0_token_raises(self): with patch("src.services.observability_resolve.get_server_auth0_token", return_value=None): with patch( "src.services.observability_resolve.get_valid_server_auth", new_callable=AsyncMock, ): with pytest.raises(AuthServiceError) as exc_info: await resolve_metrics_url( cluster_id="cluster1", env="dev", deploy_type="starter", ) assert "Auth0 token" in exc_info.value.message @pytest.mark.asyncio async def test_retries_once_on_401_with_refreshed_token(self): payload = { "code": 0, "data": { "clusters": [ { "internalReadUri": "https://ds.example.com/internal/metrics/abc/api/v1/", } ] }, } unauthorized = httpx.Response(401, json={"code": 401}) ok = httpx.Response(200, json=payload) mock_client = MagicMock() mock_client.get = AsyncMock(side_effect=[unauthorized, ok]) mock_cm = MagicMock() mock_cm.__aenter__ = AsyncMock(return_value=mock_client) mock_cm.__aexit__ = AsyncMock(return_value=None) with patch( "src.services.observability_resolve.get_server_auth0_token", side_effect=["stale-token", "fresh-token"], ): with patch( "src.services.observability_resolve.get_valid_server_auth", new_callable=AsyncMock, ) as mock_refresh: with patch( "src.services.observability_resolve.httpx.AsyncClient", return_value=mock_cm, ): resolved = await resolve_metrics_url( cluster_id="cluster1", env="dev", deploy_type="starter", ) mock_refresh.assert_awaited_once() assert mock_client.get.await_count == 2 assert resolved.token == "fresh-token" @pytest.mark.asyncio async def test_missing_gs_base_url_raises(self): with patch( "src.services.observability_resolve.get_observability_global_service_url", return_value=None, ): with patch( "src.services.observability_resolve.get_server_auth0_token", return_value="auth0-token", ): with pytest.raises(AuthServiceError) as exc_info: await resolve_metrics_url( cluster_id="cluster1", env="dev", deploy_type="starter", ) assert "Observability Global Service URL" in exc_info.value.message @pytest.mark.asyncio async def test_preserves_cluster_id_case_in_application_id(self): payload = { "code": 0, "data": { "clusters": [ { "internalReadUri": "https://ds.example.com/internal/metrics/abc/api/v1/", } ] }, } response = httpx.Response(200, json=payload) with patch( "src.services.observability_resolve.get_server_auth0_token", return_value="auth0-token" ): with patch( "src.services.observability_resolve.httpx.AsyncClient", return_value=_mock_httpx_get_response(response), ) as mock_client_cls: await resolve_metrics_url( cluster_id="CLuStEr-A1", env="dev", deploy_type="starter", ) get_call = mock_client_cls.return_value.__aenter__.return_value.get _, kwargs = get_call.call_args assert kwargs["params"]["applicationID"] == "CLuStEr-A1" class TestClusterIdNormalizationAcrossModels: def test_list_metrics_cluster_id_float_rejected(self): from src.models.list_metrics import ListMetricsRequest with pytest.raises(ValueError): ListMetricsRequest.model_validate( { "cluster_id": 10161575535533068800.0, "env": "prod", "deploy_type": "premium", } ) def test_query_range_cluster_id_float_rejected(self): from src.models.query_range import RangeQueryRequest with pytest.raises(ValueError): RangeQueryRequest.model_validate( { "query": "up", "start": "2024-01-01T00:00:00Z", "end": "2024-01-01T00:01:00Z", "step": "1m", "cluster_id": 10161575535533068800.0, "env": "prod", "deploy_type": "premium", } ) def test_query_instant_cluster_id_float_rejected(self): from src.models.query_instant import InstantQueryRequest with pytest.raises(ValueError): InstantQueryRequest.model_validate( { "query": "up", "cluster_id": 10161575535533068800.0, "env": "prod", "deploy_type": "premium", } ) class TestGetClusterInfo: @pytest.fixture(autouse=True) def _mock_gs_base_url(self): with patch( "src.services.observability_resolve.get_observability_global_service_url", return_value="https://gs.example.com", ): yield @pytest.mark.asyncio async def test_gets_cluster_info_from_gs_response(self): payload = { "code": 0, "message": "", "data": { "clusters": [ { "id": "019cd1e0-a8ff-7b6f-9e90-209d4946bb14", "tenantID": "1369847559691397282", "internalReadUri": ( "https://www.ds.us-west-2-staging.aws.observability.tidbcloud.com/" "internal/metrics/019cd1e0-a8ff-7b6f-9e90-209d4946bb14/api/v1/" ), "vendor": "aws", "region": "us-west-2", } ] }, } response = httpx.Response(200, json=payload) with patch( "src.services.observability_resolve.get_server_auth0_token", return_value="auth0-token", ): with patch( "src.services.observability_resolve.httpx.AsyncClient", return_value=_mock_httpx_get_response(response), ): cluster, token = await get_cluster_info( cluster_id="10289018618901654118", env="staging", deploy_type="dedicated", ) assert cluster["id"] == "019cd1e0-a8ff-7b6f-9e90-209d4946bb14" assert cluster["tenantID"] == "1369847559691397282" assert cluster["vendor"] == "aws" assert cluster["region"] == "us-west-2" assert token == "auth0-token" @pytest.mark.asyncio async def test_gets_cluster_info_with_biz_type_mapping(self): payload = { "code": 0, "message": "", "data": { "clusters": [ { "id": "019cd1e0-a8ff-7b6f-9e90-209d4946bb14", "tenantID": "1369847559691397282", "internalReadUri": ( "https://www.ds.us-west-2-staging.aws.observability.tidbcloud.com/" "internal/metrics/019cd1e0-a8ff-7b6f-9e90-209d4946bb14/api/v1/" ), } ] }, } response = httpx.Response(200, json=payload) with patch( "src.services.observability_resolve.get_server_auth0_token", return_value="auth0-token", ): with patch( "src.services.observability_resolve.httpx.AsyncClient", return_value=_mock_httpx_get_response(response), ): cluster, token = await get_cluster_info( cluster_id="10289018618901654118", env="staging", deploy_type="dedicated", ) assert cluster["id"] == "019cd1e0-a8ff-7b6f-9e90-209d4946bb14" assert cluster["tenantID"] == "1369847559691397282" @pytest.mark.asyncio async def test_raises_error_when_no_clusters_found(self): payload = { "code": 0, "message": "", "data": {"clusters": []}, } response = httpx.Response(200, json=payload) with patch( "src.services.observability_resolve.get_server_auth0_token", return_value="auth0-token", ): with patch( "src.services.observability_resolve.httpx.AsyncClient", return_value=_mock_httpx_get_response(response), ): with pytest.raises(AuthServiceError) as exc_info: await get_cluster_info( cluster_id="10289018618901654118", env="staging", deploy_type="dedicated", ) assert "no clusters" in exc_info.value.message.lower()