feat: add IdP, API and form submit performance tests (#853)

Add locust scripts to performance test the IdP, API and form submission flow:

- IdP: request and introspect (validate) an access token.
- API: request new form submissions and retrieve a form submission.
- Form submit: generate a random response and send through the Reliability 
queue using the Submission lambda function.

.devcontainer/devcontainer.json

@@ -19,6 +19,10 @@
     },
     "node": {
       "version": "lts"
+    },
+    "ghcr.io/devcontainers/features/python:1": {
+      "version": "3.12",
+      "installTools": false
     }
   },
 

.github/workflows/test-lambda-code/action.yml

@@ -10,7 +10,7 @@ runs:
   using: "composite"
   steps:
     - run: |
-        if [ -d tests ]; then
+        if [ -d tests ] && [ -f yarn.lock ]; then
           yarn install
           yarn test
         else

.gitignore

@@ -57,6 +57,11 @@ tf.plan
 # Ignore coverage reports
 **/coverage/
 
-# Ignore IdP certificates
+# Ignore IdP certificates and private keys
 idp/docker/*.crt
-idp/docker/*.key
+idp/docker/*.key
+*_private_api_key.json
+
+# Python
+*.pyc
+__pychache__

lambda-code/load-testing/Makefile

@@ -0,0 +1,13 @@
+fmt:
+	black .
+
+install:
+	pip3 install --user -r requirements.txt
+
+locust:
+	locust -f tests/locust_test_file.py --host=https://forms-staging.cdssandbox.xyz
+
+.PHONY: \
+	fmt \
+	install	\
+	locust

lambda-code/load-testing/main.py

@@ -1,47 +1,51 @@
 import logging
-import json
 import os
+import boto3
 from invokust.aws_lambda import get_lambda_runtime_info
 from invokust import LocustLoadTest, create_settings
 
 logging.basicConfig(level=logging.INFO)
 
-class LoadTest(LocustLoadTest):
-    def getFormInfo(self):
-        for cls in self.env.user_classes:
-            cls.on_test_stop()
+ssm_client = boto3.client("ssm")
+
+
+def get_ssm_parameter(client, parameter_name):
+    response = client.get_parameter(Name=parameter_name, WithDecryption=True)
+    return response["Parameter"]["Value"]
+
+# Load required environment variables from AWS SSM
+os.environ["FORM_ID"] = get_ssm_parameter(ssm_client, "load-testing/form-id")
+os.environ["PRIVATE_API_KEY_APP_JSON"] = get_ssm_parameter(
+    ssm_client, "load-testing/private-api-key-app"
+)
+os.environ["PRIVATE_API_KEY_USER_JSON"] = get_ssm_parameter(
+    ssm_client, "load-testing/private-api-key-user"
+)
 
 
 def handler(event=None, context=None):
-    try:
-        if event:
-            settings = create_settings(**event)
-        else:
-            settings = create_settings(from_environment=True)
 
-        if os.path.exists("/tmp/form_completion.json"):
-            os.remove("/tmp/form_completion.json")
+    # Check for required environment variables
+    required_env_vars = [
+        "FORM_ID",
+        "PRIVATE_API_KEY_APP_JSON",
+        "PRIVATE_API_KEY_USER_JSON",
+    ]
+    for env_var in required_env_vars:
+        if env_var not in os.environ:
+            raise ValueError(f"Missing required environment variable: {env_var}")
 
-        loadtest = LoadTest(settings)
+    try:
+        settings = (
+            create_settings(**event)
+            if event
+            else create_settings(from_environment=True)
+        )
+        loadtest = LocustLoadTest(settings)
         loadtest.run()
-
     except Exception as e:
-        logging.error("Locust exception {0}".format(repr(e)))
-
+        logging.error("Exception running locust tests {0}".format(repr(e)))
     else:
-        loadtest.getFormInfo()
         locust_stats = loadtest.stats()
-        lambda_runtime_info = get_lambda_runtime_info(context)
-        loadtest_results = locust_stats.copy()
-        loadtest_results.update(lambda_runtime_info)
-
-        form_input_file = open("/tmp/form_completion.json", "r")
-        form_input = json.load(form_input_file)
-        loadtest_results.update({"form_input":form_input})
-        json_results = json.dumps(loadtest_results)
-
-        ### Clean up
-        if os.path.exists("/tmp/form_completion.json"):
-            os.remove("/tmp/form_completion.json")
-
-        return json_results
+        locust_stats.update(get_lambda_runtime_info(context))
+        return locust_stats

lambda-code/load-testing/requirements.txt

@@ -1 +1,7 @@
+black==24.10.0
+boto3==1.35.39
+cryptography==43.0.1
+httpx==0.27.2
 invokust==0.77
+locust==2.31.7
+PyJWT==2.9.0

lambda-code/load-testing/tests/behaviours/api.py

@@ -0,0 +1,89 @@
+"""
+Tests the API's retrieval of new and specific responses.
+"""
+
+import os
+import json
+
+from locust import HttpUser, task
+
+from utils.data_structures import EncryptedFormSubmission
+from utils.form_submission_decrypter import FormSubmissionDecrypter
+from utils.jwt_generator import JwtGenerator
+from utils.task_set import SequentialTaskSetWithFailure
+
+
+class RetrieveResponseBehaviour(SequentialTaskSetWithFailure):
+    def __init__(self, parent: HttpUser) -> None:
+        super().__init__(parent)
+        self.form_id = os.getenv("FORM_ID")
+        self.form_decrypted_submissions = {}
+        self.form_new_submissions = None
+        self.headers = None
+        self.jwt_user = None
+
+    def on_start(self) -> None:
+        self.jwt_user = JwtGenerator.generate(self.idp_url, self.private_api_key_user)
+        data = {
+            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+            "assertion": self.jwt_user,
+            "scope": f"openid profile urn:zitadel:iam:org:project:id:{self.idp_project_id}:aud",
+        }
+        response = self.request_with_failure_check(
+            "post", f"{self.idp_url}/oauth/v2/token", 200, data=data
+        )
+        self.headers = {
+            "Authorization": f"Bearer {response["access_token"]}",
+            "Content-Type": "application/json",
+        }
+
+    @task
+    def get_new_submissions(self) -> None:
+        self.form_new_submissions = self.request_with_failure_check(
+            "get",
+            f"{self.api_url}/forms/{self.form_id}/submission/new",
+            200,
+            headers=self.headers,
+            name=f"/forms/submission/new",
+        )
+
+    @task
+    def get_form_template(self) -> None:
+        self.request_with_failure_check(
+            "get",
+            f"{self.api_url}/forms/{self.form_id}/template",
+            200,
+            headers=self.headers,
+            name=f"/forms/template",
+        )
+
+    @task
+    def get_submission_by_name(self) -> None:
+        submission = self.form_new_submissions.pop()
+        response = self.request_with_failure_check(
+            "get",
+            f"{self.api_url}/forms/{self.form_id}/submission/{submission["name"]}",
+            200,
+            headers=self.headers,
+            name=f"/forms/submission/retrieve",
+        )
+        encrypted_submission = EncryptedFormSubmission.from_json(response)
+        decrypted_submission = FormSubmissionDecrypter.decrypt(
+            encrypted_submission, self.private_api_key_user
+        )
+        self.form_decrypted_submissions[submission["name"]] = json.loads(
+            decrypted_submission
+        )
+
+    @task
+    def confirm_submission(self) -> None:
+        submission = self.form_decrypted_submissions.popitem()
+        submission_name = submission[0]
+        submission_data = submission[1]
+        self.request_with_failure_check(
+            "put",
+            f"{self.api_url}/forms/{self.form_id}/submission/{submission_name}/confirm/{submission_data['confirmationCode']}",
+            200,
+            headers=self.headers,
+            name=f"/forms/submission/confirm",
+        )

lambda-code/load-testing/tests/behaviours/idp.py

@@ -0,0 +1,42 @@
+"""
+Tests the IdP's access token generation and introspection endpoints.
+"""
+
+from locust import HttpUser, task
+from utils.jwt_generator import JwtGenerator
+from utils.task_set import SequentialTaskSetWithFailure
+
+
+class AccessTokenBehaviour(SequentialTaskSetWithFailure):
+    def __init__(self, parent: HttpUser) -> None:
+        super().__init__(parent)
+        self.jwt_app = None
+        self.jwt_user = None
+        self.access_token = None
+
+    def on_start(self) -> None:
+        self.jwt_app = JwtGenerator.generate(self.idp_url, self.private_api_key_app)
+        self.jwt_user = JwtGenerator.generate(self.idp_url, self.private_api_key_user)
+
+    @task
+    def request_access_token(self) -> None:
+        data = {
+            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+            "assertion": self.jwt_user,
+            "scope": f"openid profile urn:zitadel:iam:org:project:id:{self.idp_project_id}:aud",
+        }
+        response = self.request_with_failure_check(
+            "post", f"{self.idp_url}/oauth/v2/token", 200, data=data
+        )
+        self.access_token = response["access_token"]
+
+    @task
+    def introspect_access_token(self) -> None:
+        data = {
+            "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+            "client_assertion": self.jwt_app,
+            "token": self.access_token,
+        }
+        self.request_with_failure_check(
+            "post", f"{self.idp_url}/oauth/v2/introspect", 200, data=data
+        )

lambda-code/load-testing/tests/behaviours/submit.py

@@ -0,0 +1,54 @@
+"""
+Tests the Form submission flow through the reliability queue
+"""
+
+import os
+
+from locust import HttpUser, task
+
+from utils.form_submission_generator import FormSubmissionGenerator
+from utils.jwt_generator import JwtGenerator
+from utils.task_set import SequentialTaskSetWithFailure
+
+
+class FormSubmitBehaviour(SequentialTaskSetWithFailure):
+    def __init__(self, parent: HttpUser) -> None:
+        super().__init__(parent)
+        self.access_token = None
+        self.jwt_user = None
+        self.form_id = os.getenv("FORM_ID")
+        self.form_template = None
+        self.form_submission_generator = None
+
+    def on_start(self) -> None:
+        self.jwt_user = JwtGenerator.generate(self.idp_url, self.private_api_key_user)
+        data = {
+            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+            "assertion": self.jwt_user,
+            "scope": f"openid profile urn:zitadel:iam:org:project:id:{self.idp_project_id}:aud",
+        }
+        response = self.request_with_failure_check(
+            "post", f"{self.idp_url}/oauth/v2/token", 200, data=data
+        )
+        self.access_token = response["access_token"]
+        headers = {
+            "Authorization": f"Bearer {self.access_token}",
+            "Content-Type": "application/json",
+        }
+        self.form_template = self.request_with_failure_check(
+            "get",
+            f"{self.api_url}/forms/{self.form_id}/template",
+            200,
+            headers=headers,
+            name=f"/forms/template",
+        )
+        self.form_submission_generator = FormSubmissionGenerator(
+            self.form_id, self.form_template
+        )
+
+    @task
+    def submit_response(self) -> None:
+        with self.parent.environment.events.request.measure(
+            "task", "/forms/submit"
+        ) as request_meta:
+            self.form_submission_generator.submit_response()