Skip to content

Commit b1bd4c6

Browse files
AndrewGavrilAndrewGavril
authored
Merge pull request #61 from moevm/load_tests
Load test
2 parents 2e2718b + e07a10b commit b1bd4c6

File tree

5 files changed

+253
-4
lines changed

5 files changed

+253
-4
lines changed

‎env/.grpc.env

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
GRPC_SERVER_HOST=0.0.0.0
22
GRPC_CLIENT_HOST=localhost
3-
GRPC_PORT=50052
3+
GRPC_PORT=50051
44
GRPC_MAX_MESSAGE_LENGTH=104857600

‎load_test.py

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
import time
2+
import random
3+
import threading
4+
import statistics
5+
from concurrent.futures import ThreadPoolExecutor
6+
from prometheus_client import start_http_server, Counter, Histogram
7+
import sys
8+
import os
9+
from dotenv import load_dotenv
10+
11+
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))
12+
from client.client import FileClient
13+
from utils.payload_generator import PayloadGenerator
14+
15+
load_dotenv(os.path.join(os.path.dirname(__file__), "env", ".grpc.env"))
16+
17+
REQUEST_COUNT = Counter('grpc_requests_total', 'Total gRPC requests')
18+
REQUEST_LATENCY = Histogram('grpc_request_latency_seconds', 'gRPC request latency')
19+
REQUEST_SIZE = Histogram('grpc_request_size_bytes', 'Size of uploaded files')
20+
REQUEST_ERRORS = Counter('grpc_request_errors_total', 'Total gRPC request errors')
21+
22+
def parse_file_size(size_str: str) -> tuple[int, str]:
23+
"""Parse file size string (e.g., '1MB')"""
24+
size_str = size_str.upper()
25+
if size_str.endswith('MB'):
26+
return int(size_str[:-2]), 'MB'
27+
elif size_str.endswith('KB'):
28+
return int(size_str[:-2]), 'KB'
29+
elif size_str.endswith('B'):
30+
return int(size_str[:-1]), 'B'
31+
else:
32+
return int(size_str), 'B'
33+
34+
class LoadTest:
35+
def __init__(self, num_users=10, duration=300, file_size="1MB", file_type="text"):
36+
self.num_users = num_users
37+
self.duration = duration
38+
self.file_size = file_size
39+
self.file_type = file_type
40+
self.results = []
41+
self.lock = threading.Lock()
42+
43+
self.client = FileClient()
44+
45+
size, unit = parse_file_size(file_size)
46+
47+
self.generator = PayloadGenerator(
48+
size=size,
49+
file_type=self.file_type,
50+
unit=unit
51+
)
52+
53+
def make_request(self):
54+
start_time = time.time()
55+
try:
56+
file_path = self.generator.generate()
57+
58+
response = self.client.upload_and_validate(file_path, self.file_type)
59+
60+
latency = time.time() - start_time
61+
REQUEST_COUNT.inc()
62+
REQUEST_LATENCY.observe(latency)
63+
REQUEST_SIZE.observe(response.size)
64+
65+
with self.lock:
66+
self.results.append(latency)
67+
68+
os.remove(file_path)
69+
70+
except Exception as e:
71+
print(f"Error making request: {e}")
72+
REQUEST_ERRORS.inc()
73+
74+
def user_loop(self):
75+
end_time = time.time() + self.duration
76+
while time.time() < end_time:
77+
self.make_request()
78+
time.sleep(1.0 / self.num_users) # Rate limiting
79+
80+
def run(self):
81+
print(f"Starting load test with {self.num_users} users for {self.duration} seconds")
82+
print(f"Test file size: {self.file_size}")
83+
print(f"Test file type: {self.file_type}")
84+
start_time = time.time()
85+
86+
#start Prometheus metrics server
87+
start_http_server(8000)
88+
89+
#start user threads
90+
with ThreadPoolExecutor(max_workers=self.num_users) as executor:
91+
futures = [executor.submit(self.user_loop) for _ in range(self.num_users)]
92+
for future in futures:
93+
future.result()
94+
95+
if self.results:
96+
avg_latency = statistics.mean(self.results)
97+
p95_latency = statistics.quantiles(self.results, n=20)[18] # 95th percentile
98+
p99_latency = statistics.quantiles(self.results, n=100)[98] # 99th percentile
99+
100+
print("\nTest Results:")
101+
print(f"Total Requests: {len(self.results)}")
102+
print(f"Average Latency: {avg_latency:.3f}s")
103+
print(f"95th Percentile: {p95_latency:.3f}s")
104+
print(f"99th Percentile: {p99_latency:.3f}s")
105+
print(f"Total Duration: {time.time() - start_time:.2f}s")
106+
107+
if __name__ == '__main__':
108+
num_users = int(sys.argv[1]) if len(sys.argv) > 1 else 10
109+
duration = int(sys.argv[2]) if len(sys.argv) > 2 else 300
110+
file_size = sys.argv[3] if len(sys.argv) > 3 else "1MB"
111+
file_type = sys.argv[4] if len(sys.argv) > 4 else "text"
112+
113+
test = LoadTest(num_users, duration, file_size, file_type)
114+
test.run()

‎log_and_metric/.env

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1-
GF_ADMIN_USER=<user>
2-
GF_ADMIN_PASSWORD=<your_password>
3-
1+
GF_ADMIN_USER=user
2+
GF_ADMIN_PASSWORD=your_password
3+
GRAFANA_URL="http://localhost:3000"
4+
PROMETHEUS_URL="http://localhost:9090"
5+
LOKI_URL="http://localhost:3100"

‎requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ pathspec==0.12.1
1515
platformdirs==4.3.7
1616
pluggy==1.5.0
1717
pre_commit==4.2.0
18+
prometheus_client==0.20.0
1819
protobuf==5.29.3
1920
pytest==8.3.5
2021
pytest-cov==6.0.0

‎run_load_test.sh

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
#!/bin/bash
2+
set -e
3+
4+
# Load environment variables
5+
if [ -f log_and_metric/.env ]; then
6+
source log_and_metric/.env
7+
else
8+
echo "Error: .env file not found"
9+
exit 1
10+
fi
11+
12+
# Default configuration
13+
TEST_DURATION=60
14+
CONCURRENT_USERS=3
15+
TEST_FILE_SIZE="10KB"
16+
TEST_FILE_TYPE="text"
17+
18+
# Parse CLI arguments
19+
while [[ "$#" -gt 0 ]]; do
20+
case $1 in
21+
-d|--duration) TEST_DURATION="$2"; shift ;;
22+
-c|--concurrent) CONCURRENT_USERS="$2"; shift ;;
23+
-s|--size) TEST_FILE_SIZE="$2"; shift ;;
24+
-t|--type) TEST_FILE_TYPE="$2"; shift ;;
25+
*) echo "Unknown parameter: $1"; exit 1 ;;
26+
esac
27+
shift
28+
done
29+
30+
function check_dependencies() {
31+
echo "Checking dependencies..."
32+
for dep in python3 pip3 curl; do
33+
if ! command -v $dep &> /dev/null; then
34+
echo "Missing dependency: $dep"
35+
exit 1
36+
fi
37+
done
38+
}
39+
40+
function check_service() {
41+
local name=$1
42+
local url=$2
43+
local retries=30
44+
local wait=2
45+
46+
echo "Checking $name at $url..."
47+
for ((i=1; i<=retries; i++)); do
48+
if curl -s "$url" > /dev/null; then
49+
echo "$name is available."
50+
return 0
51+
fi
52+
echo "Waiting for $name... ($i/$retries)"
53+
sleep $wait
54+
done
55+
echo "Failed to reach $name"
56+
return 1
57+
}
58+
59+
function check_services() {
60+
check_service "Grafana" "$GRAFANA_URL" || return 1
61+
check_service "Prometheus" "$PROMETHEUS_URL" || return 1
62+
check_service "Loki" "$LOKI_URL" || return 1
63+
}
64+
65+
function setup_virtualenv() {
66+
if [ ! -d "venv" ]; then
67+
echo "Creating virtual environment..."
68+
python3 -m venv venv
69+
source venv/bin/activate
70+
pip install --upgrade pip
71+
pip install -r requirements.txt
72+
else
73+
echo "Using existing virtual environment."
74+
source venv/bin/activate
75+
76+
if [ requirements.txt -nt venv/requirements.timestamp ]; then
77+
echo "Requirements changed. Updating..."
78+
pip install -r requirements.txt
79+
touch venv/requirements.timestamp
80+
fi
81+
fi
82+
}
83+
84+
function generate_protobuf() {
85+
if [ ! -f "src/protobuf/file_service_pb2.py" ]; then
86+
echo "Generating gRPC code..."
87+
python -m grpc_tools.protoc -Iprotos/ --python_out=src/protobuf/ --grpc_python_out=src/protobuf/ protos/file_service.proto
88+
fi
89+
export PYTHONPATH="src/protobuf:$PYTHONPATH"
90+
}
91+
92+
function generate_test_file() {
93+
case "$TEST_FILE_TYPE" in
94+
text)
95+
base64 /dev/urandom | head -c "$TEST_FILE_SIZE" > testfile.txt ;;
96+
binary)
97+
head -c "$TEST_FILE_SIZE" /dev/urandom > testfile.bin ;;
98+
*)
99+
echo "Unknown file type: $TEST_FILE_TYPE"; exit 1 ;;
100+
esac
101+
}
102+
103+
function run_load_test() {
104+
echo "Running load test for $TEST_DURATION sec with $CONCURRENT_USERS users on $TEST_FILE_TYPE file of $TEST_FILE_SIZE"
105+
python3 load_test.py "$CONCURRENT_USERS" "$TEST_DURATION" "$TEST_FILE_SIZE" "$TEST_FILE_TYPE"
106+
}
107+
108+
function collect_metrics() {
109+
echo "Collecting Prometheus metrics..."
110+
curl -s "$PROMETHEUS_URL/api/v1/query?query=grpc_requests_total" > metrics_requests.json
111+
curl -s "$PROMETHEUS_URL/api/v1/query?query=grpc_request_latency_seconds_count" > metrics_latency.json
112+
curl -s "$PROMETHEUS_URL/api/v1/query?query=grpc_request_size_bytes_sum" > metrics_size.json
113+
curl -s "$PROMETHEUS_URL/api/v1/query?query=grpc_request_errors_total" > metrics_errors.json
114+
}
115+
116+
function cleanup_files() {
117+
rm -f testfile.txt testfile.bin
118+
}
119+
120+
function main() {
121+
check_dependencies
122+
check_services || { echo "One or more services are unavailable."; exit 1; }
123+
setup_virtualenv
124+
generate_protobuf
125+
generate_test_file
126+
run_load_test
127+
collect_metrics
128+
cleanup_files
129+
echo "Test finished. Check Grafana dashboard for results."
130+
}
131+
132+
main

0 commit comments

Comments
 (0)