From a67e36870ef2cb95e4ca649cd091a0ca9e6c5da5 Mon Sep 17 00:00:00 2001 From: BradNicolle Date: Mon, 19 Jan 2026 07:02:47 +0000 Subject: [PATCH] Add a Locust loadtest for SCD to the existing Locust infrastructure. --- .basedpyright/baseline.json | 34 ----- monitoring/loadtest/locust_files/ISA.py | 10 +- monitoring/loadtest/locust_files/SCD.py | 136 ++++++++++++++++++ monitoring/loadtest/locust_files/Sub.py | 12 +- monitoring/loadtest/locust_files/client.py | 84 ++++------- monitoring/loadtest/locust_files/geo_utils.py | 97 +++++++++++++ 6 files changed, 275 insertions(+), 98 deletions(-) create mode 100644 monitoring/loadtest/locust_files/SCD.py create mode 100644 monitoring/loadtest/locust_files/geo_utils.py diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index b890589bd1..3baa7756d6 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -998,40 +998,6 @@ } } ], - "./monitoring/loadtest/locust_files/client.py": [ - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 15, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 41, - "endColumn": 47, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 33, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 32, - "endColumn": 41, - "lineCount": 1 - } - } - ], "./monitoring/mock_uss/auth.py": [ { "code": "reportArgumentType", diff --git a/monitoring/loadtest/locust_files/ISA.py b/monitoring/loadtest/locust_files/ISA.py index 7b23e00ec2..632734b38a 100644 --- a/monitoring/loadtest/locust_files/ISA.py +++ b/monitoring/loadtest/locust_files/ISA.py @@ -45,6 +45,7 @@ def create_isa(self): }, "flights_url": make_fake_url(), }, + name="/identification_service_areas/[isa_uuid]", ) if resp.status_code == 200: self.isa_dict[isa_uuid] = resp.json()["service_area"]["version"] @@ -74,6 +75,7 @@ def update_isa(self): }, "flights_url": make_fake_url(), }, + name="/identification_service_areas/[target_isa]/[target_version]", ) if resp.status_code == 200: self.isa_dict[target_isa] = resp.json()["service_area"]["version"] @@ -86,7 +88,10 @@ def get_isa(self): if not target_isa: print("Nothing to pick from isa_dict for GET") return - self.client.get(f"/identification_service_areas/{target_isa}") + self.client.get( + f"/identification_service_areas/{target_isa}", + name="/identification_service_areas/[target_isa]", + ) @task(1) def delete_isa(self): @@ -95,7 +100,8 @@ def delete_isa(self): print("Nothing to pick from isa_dict for DELETE") return self.client.delete( - f"/identification_service_areas/{target_isa}/{target_version}" + f"/identification_service_areas/{target_isa}/{target_version}", + name="/identification_service_areas/[target_isa]/[target_version]", ) def checkout_isa(self): diff --git a/monitoring/loadtest/locust_files/SCD.py b/monitoring/loadtest/locust_files/SCD.py new file mode 100644 index 0000000000..213853f9b4 --- /dev/null +++ b/monitoring/loadtest/locust_files/SCD.py @@ -0,0 +1,136 @@ +"""Loaded by default by the Locust testing framework.""" + +import argparse +import datetime +import random +import uuid + +import client +import geo_utils +import locust +import shapely + + +@locust.events.init_command_line_parser.add_listener +def init_parser(parser: argparse.ArgumentParser): + """Setup config params, populated by locust.conf.""" + + parser.add_argument( + "--uss-base-url", + type=str, + help="Base URL of the Token Exchanger from which to request JWTs", + required=True, + ) + parser.add_argument( + "--area-lat", + type=float, + help="Latitude of the center of the area in which to create flights", + required=True, + ) + parser.add_argument( + "--area-lng", + type=float, + help="Longitude of the center of the area in which to create flights", + required=True, + ) + parser.add_argument( + "--area-radius", + type=int, + help="Radius (in meters) of the area in which to create flights", + required=True, + ) + parser.add_argument( + "--max-flight-distance", + type=int, + help="Maximum distance to cover for an individual flight", + required=True, + ) + + +def _format_time(time: datetime.datetime) -> str: + return time.astimezone(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _create_volume( + polygon: shapely.Polygon, + altitude_lower: float, + altitude_upper: float, + time_start: datetime.datetime, + time_end: datetime.datetime, +): + return { + "volume": { + "outline_polygon": { + "vertices": [ + {"lat": v[1], "lng": v[0]} for v in polygon.exterior.coords[:-1] + ] + }, + "altitude_lower": { + "value": altitude_lower, + "reference": "W84", + "units": "M", + }, + "altitude_upper": { + "value": altitude_upper, + "reference": "W84", + "units": "M", + }, + }, + "time_start": { + "value": _format_time(time_start), + "format": "RFC3339", + }, + "time_end": { + "value": _format_time(time_end), + "format": "RFC3339", + }, + } + + +def _create_random_flight_path( + lat: float, lng: float, radius: int, max_flight_distance_meters: int +): + altitude_lower = random.randint(0, 10000) + altitude_upper = altitude_lower + 1 + + start_time = datetime.datetime.now() + end_time = start_time + datetime.timedelta(seconds=10) + + rects = geo_utils.create_random_flight_path( + lat, lng, radius, max_flight_distance_meters + ) + return [ + _create_volume(r, altitude_lower, altitude_upper, start_time, end_time) + for r in rects.geoms + ] + + +class SCD(client.USS): + wait_time = locust.between(0.01, 0.1) + + def on_start(self): + self.uss_base_url = self.environment.parsed_options.uss_base_url + self.lat = self.environment.parsed_options.area_lat + self.lng = self.environment.parsed_options.area_lng + self.radius = self.environment.parsed_options.area_radius + self.max_flight_distance = self.environment.parsed_options.max_flight_distance + + @locust.task + def task_put_intent(self): + entity_id = uuid.uuid4().hex + + body = { + "state": "Accepted", + "uss_base_url": self.uss_base_url, + "new_subscription": { + "uss_base_url": self.uss_base_url, + }, + "extents": _create_random_flight_path( + self.lat, self.lng, self.radius, self.max_flight_distance + ), + } + self.client.put( + f"/dss/v1/operational_intent_references/{entity_id}", + json=body, + name="/dss/v1/operational_intent_references/[id]", + ) diff --git a/monitoring/loadtest/locust_files/Sub.py b/monitoring/loadtest/locust_files/Sub.py index d2582f7682..d935389cec 100644 --- a/monitoring/loadtest/locust_files/Sub.py +++ b/monitoring/loadtest/locust_files/Sub.py @@ -47,6 +47,7 @@ def create_sub(self): }, "callbacks": {"identification_service_area_url": make_fake_url()}, }, + name="/subscriptions/[sub_uuid]", ) if resp.status_code == 200: self.sub_dict[sub_uuid] = resp.json()["subscription"]["version"] @@ -59,7 +60,10 @@ def get_sub(self): if not target_sub: print("Nothing to pick from sub_dict for GET") return - self.client.get(f"/subscriptions/{target_sub}") + self.client.get( + f"/subscriptions/{target_sub}", + name="/subscriptions/[target_sub]", + ) @task(50) def update_sub(self): @@ -86,6 +90,7 @@ def update_sub(self): }, "callbacks": {"identification_service_area_url": make_fake_url()}, }, + name="/subscriptions/[target_sub]/[target_version]", ) if resp.status_code == 200: self.sub_dict[target_sub] = resp.json()["subscription"]["version"] @@ -96,7 +101,10 @@ def delete_sub(self): if not target_sub: print("Nothing to pick from sub_dict for DELETE") return - self.client.delete(f"/subscriptions/{target_sub}/{target_version}") + self.client.delete( + f"/subscriptions/{target_sub}/{target_version}", + name="/subscriptions/[target_sub]/[target_version]", + ) def checkout_sub(self): self.lock.acquire() diff --git a/monitoring/loadtest/locust_files/client.py b/monitoring/loadtest/locust_files/client.py index e94d6f66ad..701c6b0489 100644 --- a/monitoring/loadtest/locust_files/client.py +++ b/monitoring/loadtest/locust_files/client.py @@ -1,60 +1,16 @@ #!env/bin/python3 import os -import time -from locust import User -from uas_standards.astm.f3411.v19.constants import Scope +import requests +from locust import HttpUser +from uas_standards.astm.f3411.v19.constants import Scope as f3411_scope +from uas_standards.astm.f3548.v21.constants import Scope as f3548_scope -from monitoring.monitorlib import auth, infrastructure +from monitoring.monitorlib import auth -class DSSClient(infrastructure.UTMClientSession): - _locust_environment = None - - def request(self, method: str, url: str, **kwargs): - if (method == "PUT" and len(url.split("/")) > 3) or method == "PATCH": - real_method = "UPDATE" - else: - real_method = method - name = url.split("/")[1] - start_time = time.time() - result = None - try: - result = super().request(method, url, **kwargs) - except Exception as e: - self.log_exception(real_method, name, start_time, e) - else: - if result is None or result.status_code != 200: - if result is None: - msg = "Got None for Response" - else: - msg = result.text - self.log_exception(real_method, name, start_time, Exception(msg)) - else: - total_time = int((time.time() - start_time) * 1000) - self._locust_environment.events.request_success.fire( - request_type=real_method, - name=name, - response_time=total_time, - response_length=0, - ) - return result - - def log_exception( - self, real_method: str, name: str, start_time: float, e: Exception - ): - total_time = int((time.time() - start_time) * 1000) - self._locust_environment.events.request_failure.fire( - request_type=real_method, - name=name, - response_time=total_time, - exception=e, - response_length=0, - ) - - -class USS(User): +class USS(HttpUser): # Suggested by Locust 1.2.2 API Docs https://docs.locust.io/en/stable/api.html#locust.User.abstract abstract = True isa_dict: dict[str, str] = {} @@ -63,17 +19,25 @@ class USS(User): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) auth_spec = os.environ.get("AUTH_SPEC") - oauth_adapter = auth.make_auth_adapter(auth_spec) if auth_spec else None - self.client = DSSClient(self.host, oauth_adapter) - self.client._locust_environment = self.environment + if not auth_spec: - # logging after creation of client so that we surface the error in the UI - e = Exception("Missing AUTH_SPEC environment variable, please check README") - self.client.log_exception( - "Initialization", "Create DSS Client", time.time(), e + raise Exception( + "Missing AUTH_SPEC environment variable, please check README" ) - # raising exception to not allow things to proceed further - raise e + # This is a load tester its acceptable to have all the scopes required to operate anything. # We are not testing if the scope is incorrect. We are testing if it can handle the load. - self.client.default_scopes = [Scope.Write, Scope.Read] + scopes = [ + f3411_scope.Read, + f3411_scope.Write, + f3548_scope.StrategicCoordination, + ] + oauth_adapter = auth.make_auth_adapter(auth_spec) + + def _auth( + prepared_request: requests.PreparedRequest, + ) -> requests.PreparedRequest: + oauth_adapter.add_headers(prepared_request, scopes) + return prepared_request + + self.client.auth = _auth diff --git a/monitoring/loadtest/locust_files/geo_utils.py b/monitoring/loadtest/locust_files/geo_utils.py new file mode 100644 index 0000000000..4d24fb9042 --- /dev/null +++ b/monitoring/loadtest/locust_files/geo_utils.py @@ -0,0 +1,97 @@ +import math +import random + +import shapely + + +def _create_rectangle( + center: shapely.Point, + width: float, + height: float, + rotation_angle_deg: float, +) -> shapely.Polygon: + rect = shapely.geometry.box( + center.x - width / 2, + center.y - height / 2, + center.x + width / 2, + center.y + height / 2, + ) + rotated_rect = shapely.affinity.rotate(rect, rotation_angle_deg) + return rotated_rect + + +def _create_rectangles_on_path( + start_point: shapely.geometry.Point, + bearing_deg: float, + distance: float, + rect_width: float, + rect_height: float, +) -> shapely.MultiPolygon: + bearing_rad = math.radians(bearing_deg) + + # Calculate the cartesian step between each rectangle. + dx = rect_width * math.cos(bearing_rad) + dy = rect_width * math.sin(bearing_rad) + + # Add one so that we always have at least one rect. + num_rects = 1 + math.floor(distance / rect_width) + rectangles = [] + + for _ in range(num_rects): + rectangles.append( + _create_rectangle( + start_point, + rect_width, + rect_height, + bearing_deg, + ) + ) + start_point = shapely.affinity.translate(start_point, dx, dy) + + return shapely.geometry.MultiPolygon(rectangles) + + +def _random_point_within_circle( + center: shapely.Point, + radius: float, +) -> shapely.Point: + # Take sqrt of random to ensure uniform distribution of points throughout + # circular area: + random_radius = radius * math.sqrt(random.random()) + random_angle = 2 * math.pi * random.random() + + x = random_radius * math.cos(random_angle) + center.x + y = random_radius * math.sin(random_angle) + center.y + + return shapely.geometry.Point(x, y) + + +def _meters_to_angle(distance: float) -> float: + # Rough lat/lng angle of one meter at the equator - longitude gets distorted + # towards the poles. + return distance / 111320 + + +def create_random_flight_path( + lat: float, lng: float, radius: int, max_flight_distance_meters: float +) -> shapely.geometry.MultiPolygon: + bearing_deg = random.random() * 360 + distance_meters = random.random() * max_flight_distance_meters + + # Roughly scale all distance measurements to deg lat/lng: + radius_angle = _meters_to_angle(radius) + distance_angle = _meters_to_angle(distance_meters) + rect_width = _meters_to_angle(100) + rect_height = _meters_to_angle(10) + + # Create a random start point within the circle of given radius and center: + center = shapely.geometry.Point(lng, lat) + start_point = _random_point_within_circle(center, radius_angle) + + return _create_rectangles_on_path( + start_point, + bearing_deg, + distance_angle, + rect_width, + rect_height, + )