diff --git a/.github/workflows/quality_check.yml b/.github/workflows/quality_check.yml index ad5d2bbc765..f1213e7a351 100644 --- a/.github/workflows/quality_check.yml +++ b/.github/workflows/quality_check.yml @@ -43,6 +43,7 @@ jobs: quality_check: runs-on: ubuntu-latest strategy: + fail-fast: false max-parallel: 5 matrix: python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] diff --git a/aws_lambda_powertools/shared/constants.py b/aws_lambda_powertools/shared/constants.py index a68b59a7c0c..bc19ff13b30 100644 --- a/aws_lambda_powertools/shared/constants.py +++ b/aws_lambda_powertools/shared/constants.py @@ -67,5 +67,12 @@ PRETTY_INDENT: int = 4 COMPACT_INDENT: None = None +# Metadata constants +LAMBDA_METADATA_API_ENV: str = "AWS_LAMBDA_METADATA_API" +LAMBDA_METADATA_TOKEN_ENV: str = "AWS_LAMBDA_METADATA_TOKEN" +METADATA_API_VERSION: str = "2026-01-15" +METADATA_PATH: str = "/metadata/execution-environment" +METADATA_DEFAULT_TIMEOUT_SECS: float = 1.0 + # Idempotency constants IDEMPOTENCY_DISABLED_ENV: str = "POWERTOOLS_IDEMPOTENCY_DISABLED" diff --git a/aws_lambda_powertools/utilities/metadata/__init__.py b/aws_lambda_powertools/utilities/metadata/__init__.py new file mode 100644 index 00000000000..1659d3d8f44 --- /dev/null +++ b/aws_lambda_powertools/utilities/metadata/__init__.py @@ -0,0 +1,17 @@ +""" +Utility to fetch data from the AWS Lambda Metadata Endpoint +""" + +from aws_lambda_powertools.utilities.metadata.exceptions import LambdaMetadataError +from aws_lambda_powertools.utilities.metadata.lambda_metadata import ( + LambdaMetadata, + clear_metadata_cache, + get_lambda_metadata, +) + +__all__ = [ + "LambdaMetadata", + "LambdaMetadataError", + "get_lambda_metadata", + "clear_metadata_cache", +] diff --git a/aws_lambda_powertools/utilities/metadata/exceptions.py b/aws_lambda_powertools/utilities/metadata/exceptions.py new file mode 100644 index 00000000000..4948be2952f --- /dev/null +++ b/aws_lambda_powertools/utilities/metadata/exceptions.py @@ -0,0 +1,11 @@ +""" +Lambda Metadata Service exceptions +""" + + +class LambdaMetadataError(Exception): + """Raised when the Lambda Metadata Endpoint is unavailable or returns an error.""" + + def __init__(self, message: str, status_code: int = -1): + self.status_code = status_code + super().__init__(message) diff --git a/aws_lambda_powertools/utilities/metadata/lambda_metadata.py b/aws_lambda_powertools/utilities/metadata/lambda_metadata.py new file mode 100644 index 00000000000..71bbece3aac --- /dev/null +++ b/aws_lambda_powertools/utilities/metadata/lambda_metadata.py @@ -0,0 +1,174 @@ +""" +Lambda Metadata Service client + +Fetches execution environment metadata from the Lambda Metadata Endpoint, +with caching for the sandbox lifetime. +""" + +from __future__ import annotations + +import logging +import os +import urllib.request +from dataclasses import dataclass, field +from json import JSONDecodeError +from json import loads as json_loads +from typing import Any + +from aws_lambda_powertools.shared.constants import ( + LAMBDA_INITIALIZATION_TYPE, + LAMBDA_METADATA_API_ENV, + LAMBDA_METADATA_TOKEN_ENV, + METADATA_API_VERSION, + METADATA_DEFAULT_TIMEOUT_SECS, + METADATA_PATH, + POWERTOOLS_DEV_ENV, +) +from aws_lambda_powertools.utilities.metadata.exceptions import LambdaMetadataError + +logger = logging.getLogger(__name__) + +_cache: dict[str, Any] = {} + + +@dataclass(frozen=True) +class LambdaMetadata: + """Lambda execution environment metadata returned by the metadata endpoint.""" + + availability_zone_id: str | None = None + """The Availability Zone ID where the function is executing (e.g. ``use1-az1``).""" + + _raw: dict[str, Any] = field(default_factory=dict, repr=False) + """Full raw response for forward-compatibility with future fields.""" + + +def _is_lambda_environment() -> bool: + """Check whether we are running inside a Lambda execution environment.""" + return os.environ.get(LAMBDA_INITIALIZATION_TYPE, "") != "" + + +def _is_dev_mode() -> bool: + """Check whether POWERTOOLS_DEV is enabled.""" + return os.environ.get(POWERTOOLS_DEV_ENV, "false").strip().lower() in ("true", "1") + + +def _build_metadata(data: dict[str, Any]) -> LambdaMetadata: + """Build a LambdaMetadata dataclass from the raw endpoint response.""" + return LambdaMetadata( + availability_zone_id=data.get("AvailabilityZoneID"), + _raw=data, + ) + + +def _fetch_metadata(timeout: float = METADATA_DEFAULT_TIMEOUT_SECS) -> dict[str, Any]: + """ + Fetch metadata from the Lambda Metadata Endpoint via HTTP. + + Parameters + ---------- + timeout : float + Request timeout in seconds. + + Returns + ------- + dict[str, Any] + Parsed JSON response from the metadata endpoint. + + Raises + ------ + LambdaMetadataError + If required environment variables are missing, the endpoint returns + a non-200 status, or the response cannot be parsed. + """ + api = os.environ.get(LAMBDA_METADATA_API_ENV) + token = os.environ.get(LAMBDA_METADATA_TOKEN_ENV) + + if not api: + raise LambdaMetadataError( + f"Environment variable {LAMBDA_METADATA_API_ENV} is not set. Ensure {LAMBDA_METADATA_API_ENV} is set.", + ) + if not token: + raise LambdaMetadataError( + f"Environment variable {LAMBDA_METADATA_TOKEN_ENV} is not set. Ensure {LAMBDA_METADATA_TOKEN_ENV} is set.", + ) + + url = f"http://{api}/{METADATA_API_VERSION}{METADATA_PATH}" + logger.debug("Fetching Lambda metadata from: %s", url) + + req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"}) + + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: # nosec B310 + status = resp.status + body = resp.read().decode("utf-8") + except urllib.error.HTTPError as exc: + raise LambdaMetadataError( + f"Metadata request failed with status {exc.code}", + status_code=exc.code, + ) from exc + except Exception as exc: + raise LambdaMetadataError(f"Failed to fetch Lambda metadata: {exc}") from exc + + if status != 200: + raise LambdaMetadataError( + f"Metadata request failed with status {status}", + status_code=status, + ) + + try: + data: dict[str, Any] = json_loads(body) + except (JSONDecodeError, TypeError) as exc: + raise LambdaMetadataError(f"Failed to parse metadata response: {exc}") from exc + + logger.debug("Lambda metadata response: %s", data) + return data + + +def get_lambda_metadata(*, timeout: float = METADATA_DEFAULT_TIMEOUT_SECS) -> LambdaMetadata: + """ + Retrieve Lambda execution environment metadata. + + Returns cached metadata on subsequent calls. When not running in a Lambda + environment (local dev, tests) or when ``POWERTOOLS_DEV`` is enabled, + returns an empty ``LambdaMetadata``. + + Parameters + ---------- + timeout : float + HTTP request timeout in seconds (default 1.0). + + Returns + ------- + LambdaMetadata + Metadata about the current execution environment. + + Raises + ------ + LambdaMetadataError + If the metadata endpoint is unavailable or returns an error. + + Example + ------- + >>> from aws_lambda_powertools.utilities.metadata import get_lambda_metadata + >>> metadata = get_lambda_metadata() + >>> metadata.availability_zone_id # e.g. "use1-az1" + """ + if _is_dev_mode() or not _is_lambda_environment(): + return LambdaMetadata() + + if _cache: + return _build_metadata(_cache) + + data = _fetch_metadata(timeout=timeout) + _cache.update(data) + return _build_metadata(_cache) + + +def clear_metadata_cache() -> None: + """ + Clear the cached metadata. + + Useful for testing or when you need to force a fresh fetch + (e.g. after SnapStart restore). + """ + _cache.clear() diff --git a/docs/utilities/metadata.md b/docs/utilities/metadata.md new file mode 100644 index 00000000000..b1154d4cecd --- /dev/null +++ b/docs/utilities/metadata.md @@ -0,0 +1,59 @@ +--- +title: Metadata +description: Utility +status: new +--- + + + +The Metadata utility allows you to fetch data from the [AWS Lambda Metadata Endpoint (LMDS)](https://docs.aws.amazon.com/lambda/latest/dg/lambda-metadata-endpoint.html){target="_blank"}. This can be useful for retrieving information about the Lambda execution environment, such as the Availability Zone ID. + +## Key features + +* Fetch execution environment metadata from the Lambda Metadata Endpoint +* Automatic caching for the duration of the Lambda sandbox +* Graceful fallback to empty metadata outside Lambda (local dev, tests) +* Forward-compatible dataclass that can be extended as new fields are added + +## Getting started + +### Usage + +You can fetch data from the Lambda Metadata Endpoint using the `get_lambda_metadata` function. + +???+ tip + Metadata is cached for the duration of the Lambda sandbox, so subsequent calls to `get_lambda_metadata` will return the cached data. + +=== "getting_started_metadata.py" + + ```python hl_lines="2 9 10" + --8<-- "examples/metadata/src/getting_started_metadata.py" + ``` + +You can also fetch metadata eagerly during cold start, so it's ready for subsequent invocations: + +=== "getting_started_metadata_eager.py" + + ```python hl_lines="2 8" + --8<-- "examples/metadata/src/getting_started_metadata_eager.py" + ``` + +### Available metadata + +| Property | Type | Description | +| ---------------------- | --------------- | -------------------------------------------------------------- | +| `availability_zone_id` | `str` or `None` | The AZ where the function is running (e.g., `use1-az1`) | + +## Testing your code + +The metadata endpoint is not available during local development or testing. To ease testing, the `get_lambda_metadata` function automatically detects when it's running in a non-Lambda environment and returns an empty `LambdaMetadata` instance. This allows you to write tests without needing to mock the endpoint. + +If you want to mock specific metadata values for testing purposes, you can patch the internal `_fetch_metadata` function and set the required environment variables: + +=== "testing_metadata.py" + + ```python hl_lines="6-8 13-18 21" + --8<-- "examples/metadata/src/testing_metadata.py" + ``` + +We also expose a `clear_metadata_cache` function that can be used to clear the cached metadata, allowing you to test different metadata values within the same execution context. diff --git a/examples/metadata/src/getting_started_metadata.py b/examples/metadata/src/getting_started_metadata.py new file mode 100644 index 00000000000..13d5593646d --- /dev/null +++ b/examples/metadata/src/getting_started_metadata.py @@ -0,0 +1,15 @@ +from aws_lambda_powertools import Logger +from aws_lambda_powertools.utilities.metadata import LambdaMetadata, get_lambda_metadata +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + metadata: LambdaMetadata = get_lambda_metadata() + az_id = metadata.availability_zone_id # e.g., "use1-az1" + + logger.append_keys(az_id=az_id) + logger.info("Processing request") + + return {"az_id": az_id} diff --git a/examples/metadata/src/getting_started_metadata_eager.py b/examples/metadata/src/getting_started_metadata_eager.py new file mode 100644 index 00000000000..5e1977a0d0d --- /dev/null +++ b/examples/metadata/src/getting_started_metadata_eager.py @@ -0,0 +1,15 @@ +from aws_lambda_powertools import Logger +from aws_lambda_powertools.utilities.metadata import LambdaMetadata, get_lambda_metadata +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() + +# Fetch during cold start — cached for subsequent invocations +metadata: LambdaMetadata = get_lambda_metadata() + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + logger.append_keys(az_id=metadata.availability_zone_id) + logger.info("Processing request") + + return {"az_id": metadata.availability_zone_id} diff --git a/examples/metadata/src/testing_metadata.py b/examples/metadata/src/testing_metadata.py new file mode 100644 index 00000000000..49e80248510 --- /dev/null +++ b/examples/metadata/src/testing_metadata.py @@ -0,0 +1,34 @@ +from unittest.mock import patch + +from aws_lambda_powertools.utilities.metadata import LambdaMetadata, clear_metadata_cache, get_lambda_metadata + + +def test_handler_uses_metadata(monkeypatch): + # GIVEN a Lambda environment with metadata env vars + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_API", "127.0.0.1:1234") + monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "test-token") + + mock_response = {"AvailabilityZoneID": "use1-az1"} + + with patch( + "aws_lambda_powertools.utilities.metadata.lambda_metadata._fetch_metadata", + return_value=mock_response, + ): + # WHEN calling get_lambda_metadata + metadata: LambdaMetadata = get_lambda_metadata() + + # THEN it returns the mocked metadata + assert metadata.availability_zone_id == "use1-az1" + + # Clean up cache between tests + clear_metadata_cache() + + +def test_handler_works_outside_lambda(): + # GIVEN no Lambda environment variables are set + # WHEN calling get_lambda_metadata + metadata: LambdaMetadata = get_lambda_metadata() + + # THEN it returns empty metadata without errors + assert metadata.availability_zone_id is None diff --git a/mkdocs.yml b/mkdocs.yml index db49e9e45f7..cc51089fab0 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -34,6 +34,7 @@ nav: - utilities/idempotency.md - utilities/data_masking.md - utilities/feature_flags.md + - utilities/metadata.md - utilities/streaming.md - utilities/middleware_factory.md - utilities/jmespath_functions.md @@ -244,6 +245,7 @@ plugins: - utilities/idempotency.md - utilities/data_masking.md - utilities/feature_flags.md + - utilities/metadata.md - utilities/streaming.md - utilities/middleware_factory.md - utilities/jmespath_functions.md diff --git a/tests/functional/metadata/__init__.py b/tests/functional/metadata/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/functional/metadata/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/functional/metadata/test_lambda_metadata.py b/tests/functional/metadata/test_lambda_metadata.py new file mode 100644 index 00000000000..ee8eafe5047 --- /dev/null +++ b/tests/functional/metadata/test_lambda_metadata.py @@ -0,0 +1,246 @@ +"""Tests for Lambda Metadata Service utility.""" + +from __future__ import annotations + +from collections import namedtuple +from unittest.mock import patch + +import pytest + +from aws_lambda_powertools.utilities.metadata import ( + LambdaMetadata, + LambdaMetadataError, + clear_metadata_cache, + get_lambda_metadata, +) + +MOCK_METADATA_RESPONSE = {"AvailabilityZoneID": "use1-az1"} + + +@pytest.fixture(autouse=True) +def _clear_cache(): + clear_metadata_cache() + yield + clear_metadata_cache() + + +@pytest.fixture +def lambda_context(): + context = { + "function_name": "test", + "memory_limit_in_mb": 128, + "invoked_function_arn": "arn:aws:lambda:eu-west-1:123456789012:function:test", + "aws_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72", + } + return namedtuple("LambdaContext", context.keys())(*context.values()) + + +@pytest.fixture +def lambda_event(): + return {"key": "value"} + + +@pytest.fixture +def mock_metadata_endpoint(monkeypatch): + """Simulate a Lambda environment with metadata env vars and mock the HTTP fetch.""" + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_API", "127.0.0.1:1234") + monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "test-token") + + with patch( + "aws_lambda_powertools.utilities.metadata.lambda_metadata._fetch_metadata", + return_value=MOCK_METADATA_RESPONSE, + ) as mock_fetch: + yield mock_fetch + + +# --------------------------------------------------------------------------- +# LambdaMetadata dataclass +# --------------------------------------------------------------------------- + + +def test_lambda_metadata_default_has_none_az(): + # GIVEN no data + # WHEN creating a default LambdaMetadata + metadata = LambdaMetadata() + + # THEN availability_zone_id is None + assert metadata.availability_zone_id is None + + +def test_lambda_metadata_is_frozen(): + # GIVEN a LambdaMetadata instance + metadata = LambdaMetadata(availability_zone_id="use1-az1") + + # WHEN trying to mutate it + # THEN it raises FrozenInstanceError + with pytest.raises(AttributeError): + metadata.availability_zone_id = "use1-az2" + + +# --------------------------------------------------------------------------- +# LambdaMetadataError +# --------------------------------------------------------------------------- + + +def test_lambda_metadata_error_defaults_status_code_to_minus_one(): + # GIVEN a message only + # WHEN creating a LambdaMetadataError + err = LambdaMetadataError("something broke") + + # THEN message is set and status_code defaults to -1 + assert str(err) == "something broke" + assert err.status_code == -1 + + +def test_lambda_metadata_error_stores_status_code(): + # GIVEN a message and a status code + # WHEN creating a LambdaMetadataError + err = LambdaMetadataError("not found", status_code=404) + + # THEN the status_code is stored + assert err.status_code == 404 + + +# --------------------------------------------------------------------------- +# get_lambda_metadata – non-Lambda / dev mode +# --------------------------------------------------------------------------- + + +def test_get_lambda_metadata_returns_empty_outside_lambda(lambda_context, lambda_event, monkeypatch): + # GIVEN AWS_LAMBDA_INITIALIZATION_TYPE is not set (local dev / tests) + monkeypatch.delenv("AWS_LAMBDA_INITIALIZATION_TYPE", raising=False) + + def handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked + result = handler(lambda_event, lambda_context) + + # THEN it returns empty metadata without calling the endpoint + assert result.availability_zone_id is None + + +def test_get_lambda_metadata_returns_empty_when_dev_mode(lambda_context, lambda_event, monkeypatch): + # GIVEN POWERTOOLS_DEV is enabled even though init type is set + monkeypatch.setenv("POWERTOOLS_DEV", "true") + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + + def handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked + result = handler(lambda_event, lambda_context) + + # THEN it returns empty metadata + assert result.availability_zone_id is None + + +# --------------------------------------------------------------------------- +# get_lambda_metadata – missing env vars +# --------------------------------------------------------------------------- + + +def test_get_lambda_metadata_raises_when_api_env_var_missing(lambda_context, lambda_event, monkeypatch): + # GIVEN a Lambda environment without AWS_LAMBDA_METADATA_API + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "tok") + monkeypatch.delenv("AWS_LAMBDA_METADATA_API", raising=False) + + def handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked + # THEN it raises LambdaMetadataError mentioning the missing var + with pytest.raises(LambdaMetadataError, match="AWS_LAMBDA_METADATA_API"): + handler(lambda_event, lambda_context) + + +def test_get_lambda_metadata_raises_when_token_env_var_missing(lambda_context, lambda_event, monkeypatch): + # GIVEN a Lambda environment without AWS_LAMBDA_METADATA_TOKEN + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_API", "127.0.0.1:9999") + monkeypatch.delenv("AWS_LAMBDA_METADATA_TOKEN", raising=False) + + def handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked + # THEN it raises LambdaMetadataError mentioning the missing var + with pytest.raises(LambdaMetadataError, match="AWS_LAMBDA_METADATA_TOKEN"): + handler(lambda_event, lambda_context) + + +# --------------------------------------------------------------------------- +# get_lambda_metadata – happy path +# --------------------------------------------------------------------------- + + +def test_get_lambda_metadata_returns_az_id(lambda_context, lambda_event, mock_metadata_endpoint): + # GIVEN a Lambda environment with metadata env vars configured + def handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked + result = handler(lambda_event, lambda_context) + + # THEN it returns metadata with the availability zone id + assert result.availability_zone_id == "use1-az1" + mock_metadata_endpoint.assert_called_once() + + +def test_get_lambda_metadata_caches_across_invocations(lambda_context, lambda_event, mock_metadata_endpoint): + # GIVEN a Lambda environment with metadata env vars configured + def handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked twice (simulating warm start) + first = handler(lambda_event, lambda_context) + second = handler(lambda_event, lambda_context) + + # THEN both return the same data and the endpoint was called only once + assert first.availability_zone_id == "use1-az1" + assert second.availability_zone_id == "use1-az1" + mock_metadata_endpoint.assert_called_once() + + +def test_get_lambda_metadata_refetches_after_cache_clear(lambda_context, lambda_event, mock_metadata_endpoint): + # GIVEN a Lambda environment with metadata env vars configured + def handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked, cache is cleared, then invoked again + first = handler(lambda_event, lambda_context) + clear_metadata_cache() + second = handler(lambda_event, lambda_context) + + # THEN the endpoint was called twice (cache was invalidated) + assert first.availability_zone_id == "use1-az1" + assert second.availability_zone_id == "use1-az1" + assert mock_metadata_endpoint.call_count == 2 + + +# --------------------------------------------------------------------------- +# get_lambda_metadata – error responses +# --------------------------------------------------------------------------- + + +def test_get_lambda_metadata_raises_on_endpoint_error(lambda_context, lambda_event, monkeypatch): + # GIVEN a Lambda environment where the endpoint returns a 500 + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_API", "127.0.0.1:1234") + monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "test-token") + + def handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked and the endpoint fails + with patch( + "aws_lambda_powertools.utilities.metadata.lambda_metadata._fetch_metadata", + side_effect=LambdaMetadataError("Metadata request failed with status 500", status_code=500), + ): + # THEN it raises LambdaMetadataError with the status code + with pytest.raises(LambdaMetadataError, match="status 500") as exc_info: + handler(lambda_event, lambda_context) + + assert exc_info.value.status_code == 500 diff --git a/tests/integration/metadata/__init__.py b/tests/integration/metadata/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/integration/metadata/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/integration/metadata/test_lambda_metadata_http.py b/tests/integration/metadata/test_lambda_metadata_http.py new file mode 100644 index 00000000000..6354179f34a --- /dev/null +++ b/tests/integration/metadata/test_lambda_metadata_http.py @@ -0,0 +1,192 @@ +"""Integration tests for Lambda Metadata Service – exercises the real HTTP path.""" + +from __future__ import annotations + +import http.server +import json +from collections import namedtuple + +import pytest + +from aws_lambda_powertools.utilities.metadata import ( + LambdaMetadataError, + clear_metadata_cache, + get_lambda_metadata, +) + + +@pytest.fixture(autouse=True) +def _clear_cache(): + clear_metadata_cache() + yield + clear_metadata_cache() + + +@pytest.fixture +def lambda_context(): + context = { + "function_name": "test", + "memory_limit_in_mb": 128, + "invoked_function_arn": "arn:aws:lambda:eu-west-1:123456789012:function:test", + "aws_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72", + } + return namedtuple("LambdaContext", context.keys())(*context.values()) + + +@pytest.fixture +def lambda_event(): + return {"key": "value"} + + +# --------------------------------------------------------------------------- +# HTTP server fixtures +# --------------------------------------------------------------------------- + + +def _make_handler(status: int, body: str): + """Create an HTTP handler that returns a fixed status and body.""" + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): # noqa: N802 + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body.encode()) + + def log_message(self, format, *args): # noqa: A002 + pass + + return Handler + + +@pytest.fixture +def metadata_server(monkeypatch): + """Start a local HTTP server returning valid metadata and set env vars.""" + body = json.dumps({"AvailabilityZoneID": "use1-az1"}) + server = http.server.HTTPServer(("127.0.0.1", 0), _make_handler(200, body)) + port = server.server_address[1] + + import threading + + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_API", f"127.0.0.1:{port}") + monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "test-token") + + yield server + server.shutdown() + + +@pytest.fixture +def error_server(monkeypatch): + """Start a local HTTP server returning 500 and set env vars.""" + server = http.server.HTTPServer(("127.0.0.1", 0), _make_handler(500, "Internal Server Error")) + port = server.server_address[1] + + import threading + + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_API", f"127.0.0.1:{port}") + monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "test-token") + + yield server + server.shutdown() + + +@pytest.fixture +def invalid_json_server(monkeypatch): + """Start a local HTTP server returning invalid JSON.""" + server = http.server.HTTPServer(("127.0.0.1", 0), _make_handler(200, "not-json")) + port = server.server_address[1] + + import threading + + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_API", f"127.0.0.1:{port}") + monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "test-token") + + yield server + server.shutdown() + + +# --------------------------------------------------------------------------- +# Tests – happy path +# --------------------------------------------------------------------------- + + +def test_fetch_metadata_returns_az_id(lambda_context, lambda_event, metadata_server): + # GIVEN a Lambda environment pointing to a local metadata endpoint + def lambda_handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked + result = lambda_handler(lambda_event, lambda_context) + + # THEN it returns metadata with the availability zone id + assert result.availability_zone_id == "use1-az1" + + +def test_fetch_metadata_caches_across_invocations(lambda_context, lambda_event, metadata_server): + # GIVEN a Lambda environment pointing to a local metadata endpoint + def lambda_handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked twice (warm start) + first = lambda_handler(lambda_event, lambda_context) + second = lambda_handler(lambda_event, lambda_context) + + # THEN both return the same data + assert first.availability_zone_id == "use1-az1" + assert second.availability_zone_id == "use1-az1" + + +# --------------------------------------------------------------------------- +# Tests – error paths +# --------------------------------------------------------------------------- + + +def test_fetch_metadata_raises_on_http_500(lambda_context, lambda_event, error_server): + # GIVEN a Lambda environment where the endpoint returns 500 + def lambda_handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked + # THEN it raises LambdaMetadataError with status code 500 + with pytest.raises(LambdaMetadataError, match="status 500") as exc_info: + lambda_handler(lambda_event, lambda_context) + + assert exc_info.value.status_code == 500 + + +def test_fetch_metadata_raises_on_invalid_json(lambda_context, lambda_event, invalid_json_server): + # GIVEN a Lambda environment where the endpoint returns invalid JSON + def lambda_handler(event, context): + return get_lambda_metadata() + + # WHEN the handler is invoked + # THEN it raises LambdaMetadataError about parsing + with pytest.raises(LambdaMetadataError, match="Failed to parse"): + lambda_handler(lambda_event, lambda_context) + + +def test_fetch_metadata_raises_on_unreachable_endpoint(lambda_context, lambda_event, monkeypatch): + # GIVEN a Lambda environment pointing to an unreachable endpoint + monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") + monkeypatch.setenv("AWS_LAMBDA_METADATA_API", "127.0.0.1:1") + monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "test-token") + + def lambda_handler(event, context): + return get_lambda_metadata(timeout=0.1) + + # WHEN the handler is invoked + # THEN it raises LambdaMetadataError about connection failure + with pytest.raises(LambdaMetadataError, match="Failed to fetch"): + lambda_handler(lambda_event, lambda_context)