84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
# generates a file to be used by the vegeta load testing tool
|
|
# it generates TOTAL_TARGETS amount of requests that will be run
|
|
|
|
# This is a worst case scenario for the JWT cache:
|
|
# - all requests will have a unique JWT so no cache hits
|
|
# - all jwts have an expiration that will be long enough to be
|
|
# valid at time of request but short enough that already
|
|
# validated jwts will expire later during the loadtest run
|
|
# - the above guarantees JWT cache purging will happen
|
|
#
|
|
# We want this to track resource consumption in the worst case
|
|
import time
|
|
import argparse
|
|
import sys
|
|
import random
|
|
import jwt
|
|
|
|
SECRET = b"reallyreallyreallyreallyverysafe"
|
|
URL = "http://postgrest"
|
|
TOTAL_TARGETS = 200000 # tuned by hand to reduce result variance
|
|
|
|
|
|
def generate_jwt(exp_inc: int) -> str:
|
|
"""Generate an HS256 JWT"""
|
|
now = int(time.time())
|
|
payload = {
|
|
"sub": f"user_{random.getrandbits(32)}",
|
|
"iat": now,
|
|
"exp": now + exp_inc,
|
|
"role": "postgrest_test_author",
|
|
}
|
|
|
|
return jwt.encode(payload, SECRET, "HS256")
|
|
|
|
|
|
# We want to ensure 401 Unauthorized responses don't happen during
|
|
# JWT validation, this can happen when the jwt `exp` is too short.
|
|
# At the same time, we want to ensure the `exp` is not too big,
|
|
# so expires will occur and postgREST will have to clean cached expired JWTs.
|
|
def estimate_adequate_jwt_exp_increase(iteration: int) -> int:
|
|
# estimated time takes to build and run postgrest itself
|
|
build_run_postgrest_time = 2
|
|
# estimated time it takes to generate the targets file
|
|
file_generation_time = TOTAL_TARGETS // (10**-5)
|
|
# estimated exp time so some JWTs will expire
|
|
dynamic_exp_inc = iteration // 1000
|
|
|
|
return build_run_postgrest_time + file_generation_time + dynamic_exp_inc
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate Vegeta targets with unique JWTs"
|
|
)
|
|
parser.add_argument(
|
|
"output",
|
|
help="Path to write the generated targets file",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
lines = []
|
|
start_time = time.time()
|
|
|
|
for i in range(TOTAL_TARGETS):
|
|
token = generate_jwt(estimate_adequate_jwt_exp_increase(i))
|
|
lines.append(f"OPTIONS {URL}/authors_only")
|
|
lines.append(f"Authorization: Bearer {token}")
|
|
lines.append("") # blank line to separate requests
|
|
|
|
try:
|
|
with open(args.output, "w") as f:
|
|
f.write("\n".join(lines))
|
|
except IOError as e:
|
|
print(f"Error writing to {args.output}: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
elapsed = time.time() - start_time
|
|
print(f"Created {TOTAL_TARGETS} targets with unique JWTs", end=" ")
|
|
print(f"in {args.output} ({elapsed:.2f}s)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|