From 374296372b247875420ac7ae8d538a97a1f64825 Mon Sep 17 00:00:00 2001 From: Olivier Pichon Date: Mon, 1 Dec 2025 17:35:50 +1000 Subject: [PATCH] feat(api keys): allow SDK to fetch data from multiple tenants using API_KEYS fetched via AWS Secret manager or env variable `API_KEYS_MAPPING` for local development. Once the mapping is fetched, it is cached. Providing the SDK with an `api_key` as argument to its initializer as we used to do takes precedence over the new added method. Some test cases make sure this is the case. --- poetry.lock | 119 +++++++++++++- pyproject.toml | 1 + src/cvec/cvec.py | 17 +- src/cvec/utils/api_key_service.py | 164 +++++++++++++++++++ tests/test_cvec.py | 77 ++++++++- tests/utils/test_api_key_service.py | 243 ++++++++++++++++++++++++++++ 6 files changed, 616 insertions(+), 5 deletions(-) create mode 100644 src/cvec/utils/api_key_service.py create mode 100644 tests/utils/test_api_key_service.py diff --git a/poetry.lock b/poetry.lock index 2b500cd..3dfcc70 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.1.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. [[package]] name = "annotated-types" @@ -12,6 +12,46 @@ files = [ {file = "annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89"}, ] +[[package]] +name = "boto3" +version = "1.42.0" +description = "The AWS SDK for Python" +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "boto3-1.42.0-py3-none-any.whl", hash = "sha256:af32b7f61dd6293cad728ec205bcb3611ab1bf7b7dbccfd0f2bd7b9c9af96039"}, + {file = "boto3-1.42.0.tar.gz", hash = "sha256:9c67729a6112b7dced521ea70b0369fba138e89852b029a7876041cd1460c084"}, +] + +[package.dependencies] +botocore = ">=1.41.6,<1.42.0" +jmespath = ">=0.7.1,<2.0.0" +s3transfer = ">=0.16.0,<0.17.0" + +[package.extras] +crt = ["botocore[crt] (>=1.21.0,<2.0a0)"] + +[[package]] +name = "botocore" +version = "1.41.6" +description = "Low-level, data-driven core of boto 3." +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "botocore-1.41.6-py3-none-any.whl", hash = "sha256:963cc946e885acb941c96e7d343cb6507b479812ca22566ceb3e9410d0588de0"}, + {file = "botocore-1.41.6.tar.gz", hash = "sha256:08fe47e9b306f4436f5eaf6a02cb6d55c7745d13d2d093ce5d917d3ef3d3df75"}, +] + +[package.dependencies] +jmespath = ">=0.7.1,<2.0.0" +python-dateutil = ">=2.1,<3.0.0" +urllib3 = {version = ">=1.25.4,<2.2.0 || >2.2.0,<3", markers = "python_version >= \"3.10\""} + +[package.extras] +crt = ["awscrt (==0.29.1)"] + [[package]] name = "colorama" version = "0.4.6" @@ -56,6 +96,18 @@ files = [ {file = "iniconfig-2.1.0.tar.gz", hash = "sha256:3abbd2e30b36733fee78f9c7f7308f2d0050e88f0087fd25c2645f63c773e1c7"}, ] +[[package]] +name = "jmespath" +version = "1.0.1" +description = "JSON Matching Expressions" +optional = false +python-versions = ">=3.7" +groups = ["main"] +files = [ + {file = "jmespath-1.0.1-py3-none-any.whl", hash = "sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980"}, + {file = "jmespath-1.0.1.tar.gz", hash = "sha256:90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe"}, +] + [[package]] name = "mypy" version = "1.16.1" @@ -391,6 +443,21 @@ tomli = {version = ">=1", markers = "python_version < \"3.11\""} [package.extras] dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "requests", "setuptools", "xmlschema"] +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +description = "Extensions to the standard Python datetime module" +optional = false +python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" +groups = ["main"] +files = [ + {file = "python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3"}, + {file = "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427"}, +] + +[package.dependencies] +six = ">=1.5" + [[package]] name = "ruff" version = "0.11.13" @@ -419,6 +486,36 @@ files = [ {file = "ruff-0.11.13.tar.gz", hash = "sha256:26fa247dc68d1d4e72c179e08889a25ac0c7ba4d78aecfc835d49cbfd60bf514"}, ] +[[package]] +name = "s3transfer" +version = "0.16.0" +description = "An Amazon S3 Transfer Manager" +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "s3transfer-0.16.0-py3-none-any.whl", hash = "sha256:18e25d66fed509e3868dc1572b3f427ff947dd2c56f844a5bf09481ad3f3b2fe"}, + {file = "s3transfer-0.16.0.tar.gz", hash = "sha256:8e990f13268025792229cd52fa10cb7163744bf56e719e0b9cb925ab79abf920"}, +] + +[package.dependencies] +botocore = ">=1.37.4,<2.0a.0" + +[package.extras] +crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"] + +[[package]] +name = "six" +version = "1.17.0" +description = "Python 2 and 3 compatibility utilities" +optional = false +python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" +groups = ["main"] +files = [ + {file = "six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274"}, + {file = "six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81"}, +] + [[package]] name = "tomli" version = "2.2.1" @@ -489,7 +586,25 @@ files = [ [package.dependencies] typing-extensions = ">=4.12.0" +[[package]] +name = "urllib3" +version = "2.5.0" +description = "HTTP library with thread-safe connection pooling, file post, and more." +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc"}, + {file = "urllib3-2.5.0.tar.gz", hash = "sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760"}, +] + +[package.extras] +brotli = ["brotli (>=1.0.9) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; platform_python_implementation != \"CPython\""] +h2 = ["h2 (>=4,<5)"] +socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"] +zstd = ["zstandard (>=0.18.0)"] + [metadata] lock-version = "2.1" python-versions = ">=3.10" -content-hash = "466cefb197be201d5ccff4e90dd0c8844b48e58f580817e46a5ffbf6849eecea" +content-hash = "79ed8dedb1b96d292c7e9d1db66e4a69eaa185534111acbc6f5724037e3d6b2b" diff --git a/pyproject.toml b/pyproject.toml index 8aade12..6a9ee50 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ requires-python = ">=3.10" dependencies = [ "pydantic>=2.7.0,<3.0.0", "pyarrow>=18.0.0,<20.0.0", + "boto3>=1.26.0", ] [tool.poetry] diff --git a/src/cvec/cvec.py b/src/cvec/cvec.py index eacb6a7..9a9f1f1 100644 --- a/src/cvec/cvec.py +++ b/src/cvec/cvec.py @@ -7,6 +7,7 @@ from urllib.parse import urlencode, urljoin from urllib.request import Request, urlopen +from cvec.utils.api_key_service import get_api_key_for_host from cvec.models.metric import Metric, MetricDataPoint from cvec.models.span import Span from cvec.utils.arrow_converter import ( @@ -52,9 +53,23 @@ def __init__( raise ValueError( "CVEC_HOST must be set either as an argument or environment variable" ) + + # If api_key is not provided, try to fetch it from the API key service + if not self._api_key: + try: + self._api_key = get_api_key_for_host(self.host) + if self._api_key: + logger.info( + f"API key loaded from mapping service for host: {self.host}" + ) + except ValueError as e: + # Log the error from the service, but we'll raise our own error below + logger.debug(f"Failed to load API key from service: {e}") + if not self._api_key: raise ValueError( - "CVEC_API_KEY must be set either as an argument or environment variable" + "CVEC_API_KEY must be set either as an argument, environment variable, " + "or available in the API keys mapping (AWS_API_KEYS_SECRET or API_KEYS_MAPPING)" ) # Fetch publishable key from host config diff --git a/src/cvec/utils/api_key_service.py b/src/cvec/utils/api_key_service.py new file mode 100644 index 0000000..ff5baa2 --- /dev/null +++ b/src/cvec/utils/api_key_service.py @@ -0,0 +1,164 @@ +""" +API Key Service + +Loads host-to-API-key mappings from AWS Secrets Manager or local environment variables. +Supports both AWS Secrets Manager (for production) and local JSON environment variables (for development). +""" + +import json +import logging +import os +from typing import Dict, Optional + +logger = logging.getLogger(__name__) + +# Global cache for API keys (loaded once per process/cold start) +_API_KEYS_CACHE: Optional[Dict[str, str]] = None + + +def load_api_keys() -> Optional[Dict[str, str]]: + """ + Load host-to-API-key mapping from AWS Secrets Manager or local environment variables. + Caches the mapping in a global variable to avoid repeated API calls. + + The loading priority is: + 1. AWS Secrets Manager (if AWS_API_KEYS_SECRET environment variable is set) + 2. Local JSON environment variable API_KEYS_MAPPING (for development) + + Returns: + dict: Mapping of host URLs to API keys, or None if no source is available + """ + global _API_KEYS_CACHE + + if _API_KEYS_CACHE is not None: + logger.debug("Using cached API keys") + return _API_KEYS_CACHE + + # Try AWS Secrets Manager first + aws_secret_name = os.environ.get("AWS_API_KEYS_SECRET") + if aws_secret_name: + logger.info( + f"AWS_API_KEYS_SECRET is set, loading from Secrets Manager: {aws_secret_name}" + ) + api_keys = _load_from_aws_secrets_manager(aws_secret_name) + if api_keys is not None: + _API_KEYS_CACHE = api_keys + return _API_KEYS_CACHE + + # Fall back to local environment variable + local_mapping = os.environ.get("API_KEYS_MAPPING") + if local_mapping: + logger.info("Loading API keys from local environment variable API_KEYS_MAPPING") + api_keys = _load_from_local_env(local_mapping) + if api_keys is not None: + _API_KEYS_CACHE = api_keys + return _API_KEYS_CACHE + + logger.warning( + "No API key mapping found. Set either AWS_API_KEYS_SECRET or API_KEYS_MAPPING environment variable." + ) + return None + + +def _load_from_aws_secrets_manager(secret_name: str) -> Optional[Dict[str, str]]: + """ + Load API keys from AWS Secrets Manager. + + Args: + secret_name: The name/ARN of the secret in AWS Secrets Manager + + Returns: + dict: Mapping of host URLs to API keys, or None if loading fails + """ + try: + import boto3 # type: ignore[import-untyped] + + secretsmanager = boto3.client("secretsmanager") + + response = secretsmanager.get_secret_value(SecretId=secret_name) + + # Check if SecretString exists and is not empty + if "SecretString" not in response: + logger.warning( + f"Secret '{secret_name}' exists but has no value. Please populate it." + ) + return None + + secret_string = response["SecretString"] + + # Check if the secret string is empty or whitespace + if not secret_string or not secret_string.strip(): + logger.warning(f"Secret '{secret_name}' is empty. Please populate it.") + return None + + api_keys: Dict[str, str] = json.loads(secret_string) + logger.info( + f"Loaded API keys for {len(api_keys)} hosts from AWS Secrets Manager" + ) + return api_keys + + except json.JSONDecodeError as e: + logger.error(f"Secret '{secret_name}' contains invalid JSON: {str(e)}") + return None + except Exception as e: + # Generic catch-all for any other errors (e.g., ResourceNotFoundException, network errors, etc.) + logger.error( + f"Error loading API keys from AWS Secrets Manager: {str(e)}", + exc_info=True, + ) + return None + + +def _load_from_local_env(json_string: str) -> Optional[Dict[str, str]]: + """ + Load API keys from local environment variable containing JSON. + + Args: + json_string: JSON string containing host-to-API-key mapping + + Returns: + dict: Mapping of host URLs to API keys, or None if parsing fails + """ + try: + api_keys = json.loads(json_string) + if not isinstance(api_keys, dict): + logger.error( + "API_KEYS_MAPPING must be a JSON object/dict mapping hosts to API keys" + ) + return None + + logger.info(f"Loaded API keys for {len(api_keys)} hosts from local environment") + return api_keys + + except json.JSONDecodeError as e: + logger.error(f"API_KEYS_MAPPING contains invalid JSON: {str(e)}") + return None + + +def get_api_key_for_host(host: str) -> Optional[str]: + """ + Get the API key for a specific host. + + Args: + host: The host URL + + Returns: + str: The API key for the host, or None if not found + + Raises: + ValueError: If no API keys mapping is available + """ + api_keys = load_api_keys() + + if api_keys is None: + raise ValueError( + "No API keys mapping available. Set either AWS_API_KEYS_SECRET or " + "API_KEYS_MAPPING environment variable." + ) + + api_key = api_keys.get(host) + if not api_key: + logger.warning(f"No API key found for host: {host}") + return None + + return api_key diff --git a/tests/test_cvec.py b/tests/test_cvec.py index 31378bd..4cc1b39 100644 --- a/tests/test_cvec.py +++ b/tests/test_cvec.py @@ -68,14 +68,18 @@ def test_constructor_missing_host_raises_value_error( @patch.object(CVec, "_login_with_supabase", return_value=None) @patch.object(CVec, "_fetch_publishable_key", return_value="test_publishable_key") + @patch("cvec.cvec.get_api_key_for_host") @patch.dict(os.environ, {}, clear=True) def test_constructor_missing_api_key_raises_value_error( - self, mock_fetch_key: Any, mock_login: Any + self, mock_get_api_key: Any, mock_fetch_key: Any, mock_login: Any ) -> None: """Test CVec constructor raises ValueError if api_key is missing.""" + # Mock the api_key_service to raise ValueError (no mapping available) + mock_get_api_key.side_effect = ValueError("No API keys mapping available") + with pytest.raises( ValueError, - match="CVEC_API_KEY must be set either as an argument or environment variable", + match="CVEC_API_KEY must be set either as an argument, environment variable", ): CVec(host="test_host") @@ -147,6 +151,75 @@ def test_construct_email_from_api_key_invalid_length( ): client._construct_email_from_api_key() + @patch.object(CVec, "_login_with_supabase", return_value=None) + @patch.object(CVec, "_fetch_publishable_key", return_value="test_publishable_key") + @patch("cvec.cvec.get_api_key_for_host") + @patch.dict(os.environ, {"CVEC_HOST": "https://tenant1.cvector.dev"}, clear=True) + def test_constructor_loads_api_key_from_service( + self, mock_get_api_key: Any, mock_fetch_key: Any, mock_login: Any + ) -> None: + """Test CVec constructor loads API key from api_key_service when not provided.""" + mock_get_api_key.return_value = "cva_service12345678901234567890123456789" + + client = CVec() + + assert client.host == "https://tenant1.cvector.dev" + assert client._api_key == "cva_service12345678901234567890123456789" + mock_get_api_key.assert_called_once_with("https://tenant1.cvector.dev") + + @patch.object(CVec, "_login_with_supabase", return_value=None) + @patch.object(CVec, "_fetch_publishable_key", return_value="test_publishable_key") + @patch("cvec.cvec.get_api_key_for_host") + @patch.dict(os.environ, {}, clear=True) + def test_constructor_api_key_service_returns_none( + self, mock_get_api_key: Any, mock_fetch_key: Any, mock_login: Any + ) -> None: + """Test CVec constructor raises when api_key_service returns None.""" + mock_get_api_key.return_value = None + + with pytest.raises( + ValueError, + match="CVEC_API_KEY must be set either as an argument, environment variable", + ): + CVec(host="https://tenant1.cvector.dev") + + @patch.object(CVec, "_login_with_supabase", return_value=None) + @patch.object(CVec, "_fetch_publishable_key", return_value="test_publishable_key") + @patch("cvec.cvec.get_api_key_for_host") + @patch.dict( + os.environ, + { + "CVEC_HOST": "https://tenant1.cvector.dev", + "CVEC_API_KEY": "cva_envkey123456789012345678901234567890", + }, + clear=True, + ) + def test_constructor_env_api_key_takes_precedence_over_service( + self, mock_get_api_key: Any, mock_fetch_key: Any, mock_login: Any + ) -> None: + """Test that CVEC_API_KEY env var takes precedence over api_key_service.""" + client = CVec() + + assert client._api_key == "cva_envkey123456789012345678901234567890" + # Should not call the service if env var is set + mock_get_api_key.assert_not_called() + + @patch.object(CVec, "_login_with_supabase", return_value=None) + @patch.object(CVec, "_fetch_publishable_key", return_value="test_publishable_key") + @patch("cvec.cvec.get_api_key_for_host") + def test_constructor_arg_api_key_takes_precedence_over_service( + self, mock_get_api_key: Any, mock_fetch_key: Any, mock_login: Any + ) -> None: + """Test that api_key argument takes precedence over api_key_service.""" + client = CVec( + host="https://tenant1.cvector.dev", + api_key="cva_argkey123456789012345678901234567890", + ) + + assert client._api_key == "cva_argkey123456789012345678901234567890" + # Should not call the service if api_key argument is provided + mock_get_api_key.assert_not_called() + class TestCVecGetSpans: @patch.object(CVec, "_login_with_supabase", return_value=None) diff --git a/tests/utils/test_api_key_service.py b/tests/utils/test_api_key_service.py new file mode 100644 index 0000000..e066b04 --- /dev/null +++ b/tests/utils/test_api_key_service.py @@ -0,0 +1,243 @@ +import json +import os +import pytest +from unittest.mock import MagicMock, patch +from typing import Any + +from cvec.utils import api_key_service + + +# Reset cache before each test +@pytest.fixture(autouse=True) +def reset_cache() -> Any: + """Reset the global API keys cache before each test.""" + api_key_service._API_KEYS_CACHE = None + yield + api_key_service._API_KEYS_CACHE = None + + +class TestLoadFromAwsSecretsManager: + def test_load_from_aws_success(self) -> None: + """Test successful loading from AWS Secrets Manager.""" + # Mock boto3 module and client + mock_boto3 = MagicMock() + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + + # Mock secret response + api_keys_data = { + "https://tenant1.cvector.dev": "cva_1234567890123456789012345678901234567", + "https://tenant2.cvector.dev": "cva_abcdefghijklmnopqrstuvwxyz1234567890", + } + mock_client.get_secret_value.return_value = { + "SecretString": json.dumps(api_keys_data) + } + + with patch.dict("sys.modules", {"boto3": mock_boto3}): + result = api_key_service._load_from_aws_secrets_manager("test-secret") + + assert result == api_keys_data + mock_client.get_secret_value.assert_called_once_with(SecretId="test-secret") + + def test_load_from_aws_empty_secret(self) -> None: + """Test loading from AWS with empty secret value.""" + mock_boto3 = MagicMock() + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.get_secret_value.return_value = {"SecretString": ""} + + with patch.dict("sys.modules", {"boto3": mock_boto3}): + result = api_key_service._load_from_aws_secrets_manager("test-secret") + + assert result is None + + def test_load_from_aws_no_secret_string(self) -> None: + """Test loading from AWS when SecretString is missing.""" + mock_boto3 = MagicMock() + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.get_secret_value.return_value = {} + + with patch.dict("sys.modules", {"boto3": mock_boto3}): + result = api_key_service._load_from_aws_secrets_manager("test-secret") + + assert result is None + + def test_load_from_aws_invalid_json(self) -> None: + """Test loading from AWS with invalid JSON.""" + mock_boto3 = MagicMock() + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.get_secret_value.return_value = {"SecretString": "not valid json"} + + with patch.dict("sys.modules", {"boto3": mock_boto3}): + result = api_key_service._load_from_aws_secrets_manager("test-secret") + + assert result is None + + +class TestLoadFromLocalEnv: + def test_load_from_local_env_success(self) -> None: + """Test successful loading from local environment variable.""" + api_keys_data = { + "https://tenant1.cvector.dev": "cva_1234567890123456789012345678901234567", + "https://tenant2.cvector.dev": "cva_abcdefghijklmnopqrstuvwxyz1234567890", + } + json_string = json.dumps(api_keys_data) + + result = api_key_service._load_from_local_env(json_string) + + assert result == api_keys_data + + def test_load_from_local_env_invalid_json(self) -> None: + """Test loading from local env with invalid JSON.""" + result = api_key_service._load_from_local_env("not valid json") + + assert result is None + + def test_load_from_local_env_not_dict(self) -> None: + """Test loading from local env when JSON is not a dict.""" + result = api_key_service._load_from_local_env('["array", "not", "dict"]') + + assert result is None + + +class TestLoadApiKeys: + @patch.dict(os.environ, {}, clear=True) + def test_load_api_keys_no_source(self) -> None: + """Test loading API keys when no source is available.""" + result = api_key_service.load_api_keys() + + assert result is None + + @patch("cvec.utils.api_key_service._load_from_aws_secrets_manager") + @patch.dict( + os.environ, + {"AWS_API_KEYS_SECRET": "test-secret"}, + clear=True, + ) + def test_load_api_keys_from_aws(self, mock_aws_load: Any) -> None: + """Test loading API keys from AWS Secrets Manager.""" + api_keys_data = { + "https://tenant1.cvector.dev": "cva_1234567890123456789012345678901234567", + } + mock_aws_load.return_value = api_keys_data + + result = api_key_service.load_api_keys() + + assert result == api_keys_data + mock_aws_load.assert_called_once_with("test-secret") + + @patch.dict( + os.environ, + { + "API_KEYS_MAPPING": '{"https://tenant1.cvector.dev": "cva_1234567890123456789012345678901234567"}' + }, + clear=True, + ) + def test_load_api_keys_from_local_env(self) -> None: + """Test loading API keys from local environment variable.""" + result = api_key_service.load_api_keys() + + assert result == { + "https://tenant1.cvector.dev": "cva_1234567890123456789012345678901234567" + } + + @patch("cvec.utils.api_key_service._load_from_aws_secrets_manager") + @patch.dict( + os.environ, + { + "AWS_API_KEYS_SECRET": "test-secret", + "API_KEYS_MAPPING": '{"https://local.cvector.dev": "cva_local123456789012345678901234567890"}', + }, + clear=True, + ) + def test_load_api_keys_aws_priority(self, mock_aws_load: Any) -> None: + """Test that AWS Secrets Manager has priority over local env.""" + aws_keys = {"https://aws.cvector.dev": "cva_awskey1234567890123456789012345678"} + mock_aws_load.return_value = aws_keys + + result = api_key_service.load_api_keys() + + assert result == aws_keys + mock_aws_load.assert_called_once_with("test-secret") + + @patch("cvec.utils.api_key_service._load_from_aws_secrets_manager") + @patch.dict( + os.environ, + { + "AWS_API_KEYS_SECRET": "test-secret", + "API_KEYS_MAPPING": '{"https://local.cvector.dev": "cva_local123456789012345678901234567890"}', + }, + clear=True, + ) + def test_load_api_keys_fallback_to_local(self, mock_aws_load: Any) -> None: + """Test fallback to local env when AWS fails.""" + mock_aws_load.return_value = None + + result = api_key_service.load_api_keys() + + assert result == { + "https://local.cvector.dev": "cva_local123456789012345678901234567890" + } + + @patch("cvec.utils.api_key_service._load_from_aws_secrets_manager") + @patch.dict( + os.environ, + {"AWS_API_KEYS_SECRET": "test-secret"}, + clear=True, + ) + def test_load_api_keys_caching(self, mock_aws_load: Any) -> None: + """Test that API keys are cached after first load.""" + api_keys_data = { + "https://tenant1.cvector.dev": "cva_1234567890123456789012345678901234567", + } + mock_aws_load.return_value = api_keys_data + + # First call + result1 = api_key_service.load_api_keys() + assert result1 == api_keys_data + assert mock_aws_load.call_count == 1 + + # Second call should use cache + result2 = api_key_service.load_api_keys() + assert result2 == api_keys_data + assert mock_aws_load.call_count == 1 # Still 1, not called again + + +class TestGetApiKeyForHost: + @patch("cvec.utils.api_key_service.load_api_keys") + def test_get_api_key_for_host_success(self, mock_load: Any) -> None: + """Test successfully getting API key for a host.""" + api_keys_data = { + "https://tenant1.cvector.dev": "cva_1234567890123456789012345678901234567", + "https://tenant2.cvector.dev": "cva_abcdefghijklmnopqrstuvwxyz1234567890", + } + mock_load.return_value = api_keys_data + + result = api_key_service.get_api_key_for_host("https://tenant1.cvector.dev") + + assert result == "cva_1234567890123456789012345678901234567" + + @patch("cvec.utils.api_key_service.load_api_keys") + def test_get_api_key_for_host_not_found(self, mock_load: Any) -> None: + """Test getting API key for host that doesn't exist in mapping.""" + api_keys_data = { + "https://tenant1.cvector.dev": "cva_1234567890123456789012345678901234567", + } + mock_load.return_value = api_keys_data + + result = api_key_service.get_api_key_for_host("https://unknown.cvector.dev") + + assert result is None + + @patch("cvec.utils.api_key_service.load_api_keys") + def test_get_api_key_for_host_no_mapping(self, mock_load: Any) -> None: + """Test getting API key when no mapping is available.""" + mock_load.return_value = None + + with pytest.raises( + ValueError, + match="No API keys mapping available", + ): + api_key_service.get_api_key_for_host("https://tenant1.cvector.dev")