Skip to content

feat: add load testing infrastructure with Locust #164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aanyacloud wants to merge 2 commits into
base: main
Choose a base branch
from
+116 −1
Open

feat: add load testing infrastructure with Locust #164

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
57544ee
feat: add load testing infrastructure with Locust
aanyacloud Jun 12, 2026
6a6442e
Merge branch 'main' into issue-131-load-testing
S3DFX-CYBER Jun 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion requirements-dev.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,4 +27,5 @@ watchdog>=6.0.0

# Type Stubs
types-requests>=2.31.0
types-redis>=4.6.0.20241004
types-redis>=4.6.0.20241004
locust>=2.31.0
58 changes: 58 additions & 0 deletions tests/load/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
# TENET AI Load Testing

## Overview

This directory contains load testing scripts for TENET AI APIs using Locust.

## Tested Endpoints

- GET /health
- GET /v1/stats
- GET /v1/events
- POST /v1/events/llm

## Installation

```bash
pip install locust
```

## Running Tests

```bash
locust -f tests/load/locustfile.py
```

Then open:

http://localhost:8089
Comment on lines +20 to +28

@coderabbitai coderabbitai Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Document required host configuration and API key setup.

The running instructions omit critical setup details: users must specify the target host URL via --host and ensure a valid API key is configured. Without these, the load tests will fail.

📝 Suggested documentation addition 
 ## Running Tests
 
+### Prerequisites
+
+1. Ensure the TENET AI ingest service is running
+2. Configure a valid API key with `ingest` and `read` permissions (see Authentication section below)
+
+### Execute
+
 ```bash
-locust -f tests/load/locustfile.py
+locust -f tests/load/locustfile.py --host=http://localhost:8000

Then open:

http://localhost:8089
+
+### Authentication
+
+Update the API_KEY constant in locustfile.py with a valid test API key that has:
+- ingest permission for POST /v1/events/llm
+- read permission for GET /v1/stats and GET /v1/events

</details>

<details>
<summary>🤖 Prompt for AI Agents</summary>

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In @tests/load/README.md around lines 20 - 28, Update the tests/load/README.md
to include required host and auth setup: instruct users to pass --host (example:
--host=http://localhost:8000) when running locust with tests/load/locustfile.py
and add an "Authentication" section telling them to set the API_KEY constant in
locustfile.py to a valid test API key with the needed permissions (ingest for
POST /v1/events/llm and read for GET /v1/stats and GET /v1/events); reference
locustfile.py and the API_KEY constant so it's clear where to configure these
values.


</details>

<!-- fingerprinting:phantom:poseidon:puma -->

<!-- cr-comment:v1:fc7fe9e7eef0cf670bb72c46 -->

<!-- This is an auto-generated comment by CodeRabbit -->


## Test Scenarios

### Sustained Load Test

- Target: 100 requests/second
- Duration: 10 minutes

### Spike Load Test

- Simulate sudden bursts of traffic
- Increase users rapidly

### Concurrent Connection Test

- Run multiple concurrent users
- Observe latency and error rates

## Metrics to Measure

- P50 latency
- P95 latency
- P99 latency
- Throughput
- Error rate

## Acceptance Criteria

- P95 latency < 500 ms under sustained load
- Error rate < 0.1% under normal load
56 changes: 56 additions & 0 deletions tests/load/locustfile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
from locust import HttpUser, task, between

API_KEY = "test-api-key"

@coderabbitai coderabbitai Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Replace hardcoded API key with environment variable.

The hardcoded test-api-key constant poses security and configuration risks: it may be committed to version control and cannot be changed per environment without modifying code. From the relevant code snippets (services/ingest/app.py), the API requires valid authenticated keys with specific permissions (ingest for POST /v1/events/llm, read for GET endpoints).

🔐 Proposed fix using environment variable 
+import os
 from locust import HttpUser, task, between
 
-API_KEY = "test-api-key"
+API_KEY = os.getenv("TENET_LOAD_TEST_API_KEY", "test-api-key")

Then document in README.md:

export TENET_LOAD_TEST_API_KEY="your-valid-test-key"
locust -f tests/load/locustfile.py --host=http://localhost:8000
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
API_KEY = "test-api-key"
import os
from locust import HttpUser, task, between
API_KEY = os.getenv("TENET_LOAD_TEST_API_KEY", "test-api-key")
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/load/locustfile.py` at line 3, Replace the hardcoded API_KEY constant
in tests/load/locustfile.py with a value read from an environment variable
(e.g., TENET_LOAD_TEST_API_KEY); update the code that references API_KEY so it
raises a clear error or exits if the env var is missing, and ensure any locust
task functions (where API_KEY is used) continue to reference the same symbol
(API_KEY) so no other changes are required; also add a brief README note or
comment showing how to export TENET_LOAD_TEST_API_KEY before running locust.

@cubic-dev-ai cubic-dev-ai Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1: Hardcoding test-api-key here makes the load test brittle and likely unauthorized under the project’s expected API-key configuration, so the script may only measure 401/403 responses instead of actual ingest performance.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At tests/load/locustfile.py, line 3:

<comment>Hardcoding `test-api-key` here makes the load test brittle and likely unauthorized under the project’s expected API-key configuration, so the script may only measure 401/403 responses instead of actual ingest performance.</comment>

<file context>
@@ -0,0 +1,56 @@
+from locust import HttpUser, task, between
+
+API_KEY = "test-api-key"
+
+
</file context>



class TENETLoadUser(HttpUser):
wait_time = between(1, 3)

headers = {
"x-api-key": API_KEY,
"Content-Type": "application/json",
}
Comment on lines +9 to +12

@coderabbitai coderabbitai Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Convert mutable class attribute to property or instance attribute.

Ruff correctly flags that defining headers as a mutable dict at class level means all instances share the same dict object. If any task modifies self.headers, it affects all concurrent users.

🛡️ Recommended fix using on_start 
-    headers = {
-        "x-api-key": API_KEY,
-        "Content-Type": "application/json",
-    }
+
+    def on_start(self):
+        """Initialize headers per user instance."""
+        self.headers = {
+            "x-api-key": API_KEY,
+            "Content-Type": "application/json",
+        }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
headers = {
"x-api-key": API_KEY,
"Content-Type": "application/json",
}
def on_start(self):
"""Initialize headers per user instance."""
self.headers = {
"x-api-key": API_KEY,
"Content-Type": "application/json",
}
🧰 Tools
🪛 Ruff (0.15.15)

[warning] 9-12: Mutable default value for class attribute

(RUF012)

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/load/locustfile.py` around lines 9 - 12, The class-level mutable dict
headers should be converted to an instance attribute or property to avoid
sharing it across Locust users; update the Locust user class (where headers is
declared) to initialize self.headers in an on_start method (or implement a
`@property` that returns a fresh dict) using API_KEY and "Content-Type":
"application/json", and update any references from headers to self.headers so
each user gets its own dict instance and concurrent modifications won't affect
other users.

Source: Linters/SAST tools


@task(3)
def health_check(self):
self.client.get(
"/health",
headers=self.headers,
name="GET /health"
)

@task(2)
def get_stats(self):
self.client.get(
"/v1/stats",
headers=self.headers,
name="GET /v1/stats"
)

@task(2)
def list_events(self):
self.client.get(
"/v1/events",
headers=self.headers,
name="GET /v1/events"
)
Comment on lines +22 to +36

@coderabbitai coderabbitai Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Add response validation and error handling to authenticated endpoints.

The get_stats and list_events tasks don't validate response status codes or catch exceptions. For load testing, tracking success/failure rates is essential to meet the acceptance criteria (error rate < 0.1%).

✅ Proposed enhancement with response validation 
     `@task`(2)
     def get_stats(self):
-        self.client.get(
+        with self.client.get(
             "/v1/stats",
             headers=self.headers,
-            name="GET /v1/stats"
-        )
+            name="GET /v1/stats",
+            catch_response=True
+        ) as response:
+            if response.status_code != 200:
+                response.failure(f"Got status code {response.status_code}")
+            elif "total_events" not in response.json():
+                response.failure("Missing expected field: total_events")
+            else:
+                response.success()
 
     `@task`(2)
     def list_events(self):
-        self.client.get(
+        with self.client.get(
             "/v1/events",
             headers=self.headers,
-            name="GET /v1/events"
-        )
+            name="GET /v1/events",
+            catch_response=True
+        ) as response:
+            if response.status_code != 200:
+                response.failure(f"Got status code {response.status_code}")
+            elif "events" not in response.json():
+                response.failure("Missing expected field: events")
+            else:
+                response.success()
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/load/locustfile.py` around lines 22 - 36, Wrap the GET calls in
get_stats and list_events with client.get(..., catch_response=True) and validate
the response inside that context: check resp.status_code (expect 200) and call
resp.success() on valid responses or resp.failure(...) with an informative
message on non-200 or unexpected body; also catch exceptions around the request
and call resp.failure(...) (or mark failure) so Locust records errors for those
authenticated endpoints (functions get_stats and list_events, and the
self.client.get calls using self.headers).


@task(1)
def submit_llm_event(self):
payload = {
"source_type": "chat",
"source_id": "load-test-user",
"model": "gpt-4",
"prompt": "This is a load testing request",
"system_prompt": "You are a helpful assistant",
"metadata": {
"environment": "load-test"
}
}

self.client.post(
"/v1/events/llm",
json=payload,
headers=self.headers,
name="POST /v1/events/llm"
)
Loading