Merge pull request #218 from aws/maxday/unit-test-invocation-id

test: assert the invocation ID reaches the wire

Author: maxday

.github/workflows/test-on-push-and-pr.yml

@@ -18,6 +18,26 @@ jobs:
     - name: Run 'pr' target
       run: make pr
 
+  test-runtime-client:
+    name: test runtime client
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+
+    # The tests in this job drive the compiled extension against a stub RAPID,
+    # so unlike the unit tests they need the native code built first.
+    - name: Install aws-lambda-cpp build dependencies
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y g++ make cmake libcurl4-openssl-dev
+
+    - name: Install requirements
+      run: make init
+
+    - name: Compile the runtime client extension and run its tests
+      run: make test-runtime-client
+
   integration-test:
     runs-on: ubuntu-latest
     strategy:

Makefile

@@ -8,8 +8,16 @@ init:
 	pip3 install -r requirements/base.txt -r requirements/dev.txt
 
 .PHONY: test
+# test_runtime_client_headers.py needs the compiled extension, so it runs in its
+# own target (test-runtime-client) rather than here.
 test: check-format
-	pytest --cov awslambdaric --cov-report term-missing --cov-fail-under 90 tests
+	pytest --cov awslambdaric --cov-report term-missing --cov-fail-under 90 \
+		--ignore tests/test_runtime_client_headers.py tests
+
+.PHONY: test-runtime-client
+test-runtime-client:
+	python3 setup.py build_ext --inplace
+	pytest -v tests/test_runtime_client_headers.py
 
 .PHONY: test-integ
 test-integ:

tests/invocation_id_probe.py

@@ -0,0 +1,51 @@
+"""
+Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+Client half of test_runtime_client_headers.py, run as a subprocess.
+
+It is a separate process rather than a thread because the native
+post_invocation_result/post_error entry points hold the GIL for the duration of
+the blocking curl call so an in-process stub server thread could never be scheduled to answer the POST.
+"""
+
+import json
+import sys
+
+import runtime_client
+
+INVOCATION_ID_HEADER = "Lambda-Runtime-Invocation-Id"
+REQUEST_ID_HEADER = "Lambda-Runtime-Aws-Request-Id"
+
+
+def main(mode):
+    runtime_client.initialize_client("test-agent")
+
+    _, headers = runtime_client.next()
+    invocation_id = headers[INVOCATION_ID_HEADER]
+
+    # Thread the invocation id straight back, exactly as bootstrap does. Passing
+    # it through untouched is what makes the POST assertions meaningful.
+    if mode == "response":
+        runtime_client.post_invocation_result(
+            headers[REQUEST_ID_HEADER],
+            b'{"ok": true}',
+            "application/json",
+            invocation_id,
+        )
+    elif mode == "error":
+        runtime_client.post_error(
+            headers[REQUEST_ID_HEADER],
+            '{"errorMessage": "boom"}',
+            "{}",
+            invocation_id,
+        )
+    else:
+        raise SystemExit(f"unknown mode: {mode!r}")
+
+    print(json.dumps({"next_invocation_id": invocation_id}))
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        raise SystemExit(__doc__)
+    main(sys.argv[1])

tests/test_runtime_client_headers.py

@@ -0,0 +1,141 @@
+"""
+Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+"""
+
+import http.server
+import json
+import os
+import subprocess
+import sys
+import threading
+import unittest
+
+INVOCATION_ID_HEADER = "Lambda-Runtime-Invocation-Id"
+REQUEST_ID = "request-id-1"
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+CLIENT_PROBE = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "invocation_id_probe.py"
+)
+
+
+class StubRapidHandler(http.server.BaseHTTPRequestHandler):
+    """Minimal stand-in for the RAPID /next, /response and /error endpoints.
+
+    Each request's raw header block is appended to the server's `received` list
+    so that tests can assert on what actually went over the wire rather than on
+    what the Python layer believed it was sending.
+    """
+
+    protocol_version = "HTTP/1.1"
+
+    def log_message(self, format, *args):
+        pass  # keep the test output clean
+
+    def _record(self):
+        self.server.received.append((self.path, self.headers))
+
+    def do_GET(self):
+        self._record()
+        body = b"{}"
+        self.send_response(200)
+        self.send_header("Lambda-Runtime-Aws-Request-Id", REQUEST_ID)
+        self.send_header("Lambda-Runtime-Invoked-Function-Arn", "arn:aws:lambda:::f")
+        self.send_header("Lambda-Runtime-Deadline-Ms", "1735689600000")
+        self.send_header("Content-Type", "application/json")
+        if self.server.invocation_id is not None:
+            self.send_header(INVOCATION_ID_HEADER, self.server.invocation_id)
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_POST(self):
+        self._record()
+        length = int(self.headers.get("Content-Length") or 0)
+        if length:
+            self.rfile.read(length)
+        self.send_response(202)
+        self.send_header("Content-Length", "0")
+        self.end_headers()
+
+
+class StubRapid(http.server.ThreadingHTTPServer):
+
+    def __init__(self):
+        super().__init__(("127.0.0.1", 0), StubRapidHandler)
+        self.invocation_id = None
+        self.received = []
+
+    @property
+    def endpoint(self):
+        return "{}:{}".format(*self.server_address)
+
+
+class TestInvocationIdOnTheWire(unittest.TestCase):
+    """End-to-end coverage of the Lambda-Runtime-Invocation-Id echo.
+
+    These tests drive the compiled extension against a real socket, so they
+    cover the header-emitting C++ code in aws-lambda-cpp that the mocked unit
+    tests in test_lambda_runtime_client.py cannot reach. If
+    aws-lambda-cpp-add-invocation-id.patch is ever dropped from the vendored
+    tarball, these are the tests that fail.
+    """
+
+    def _run_client(self, invocation_id, mode="response"):
+        """Serve one invocation to a real client and return its POST headers."""
+        server = StubRapid()
+        server.invocation_id = invocation_id
+        thread = threading.Thread(target=server.serve_forever, daemon=True)
+        thread.start()
+        self.addCleanup(thread.join, 5)
+        self.addCleanup(server.shutdown)
+        self.addCleanup(server.server_close)
+
+        env = dict(
+            os.environ,
+            AWS_LAMBDA_RUNTIME_API=server.endpoint,
+            PYTHONPATH=os.pathsep.join(
+                [REPO_ROOT]
+                + ([os.environ["PYTHONPATH"]] if os.environ.get("PYTHONPATH") else [])
+            ),
+        )
+        completed = subprocess.run(
+            [sys.executable, CLIENT_PROBE, mode],
+            env=env,
+            cwd=REPO_ROOT,
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        self.assertEqual(
+            0,
+            completed.returncode,
+            f"client failed:\nstdout={completed.stdout}\nstderr={completed.stderr}",
+        )
+
+        posted = [h for path, h in server.received if path.endswith("/" + mode)]
+        self.assertEqual(1, len(posted), f"expected exactly one POST to /{mode}")
+        return json.loads(completed.stdout), posted[0]
+
+    def test_invocation_id_is_echoed_on_response(self):
+        seen, posted = self._run_client("inv-uuid-round-trip")
+
+        self.assertEqual("inv-uuid-round-trip", seen["next_invocation_id"])
+        # get_all, not get: a duplicated header must be visible, not normalized.
+        self.assertEqual(["inv-uuid-round-trip"], posted.get_all(INVOCATION_ID_HEADER))
+
+    def test_invocation_id_is_echoed_on_error(self):
+        seen, posted = self._run_client("inv-uuid-error-path", mode="error")
+
+        self.assertEqual("inv-uuid-error-path", seen["next_invocation_id"])
+        self.assertEqual(["inv-uuid-error-path"], posted.get_all(INVOCATION_ID_HEADER))
+
+    def test_no_header_is_sent_when_rapid_did_not_send_one(self):
+        seen, posted = self._run_client(None)
+
+        self.assertIsNone(seen["next_invocation_id"])
+        self.assertIsNone(posted.get_all(INVOCATION_ID_HEADER))
+
+
+if __name__ == "__main__":
+    unittest.main()