locust-cloud 1.2.12__py3-none-any.whl → 1.3.2__py3-none-any.whl

locust_cloud/__init__.py

@@ -2,6 +2,8 @@ import os
 
 os.environ["LOCUST_SKIP_MONKEY_PATCH"] = "1"
 
+import argparse
+
 from locust import events
 from locust.argument_parser import LocustArgumentParser
 from locust_cloud.auth import register_auth
@@ -32,7 +34,7 @@ def add_arguments(parser: LocustArgumentParser):
     locust_cloud.add_argument(
         "--exporter",
         default=True,
-        action="store_true",
+        action=argparse.BooleanOptionalAction,
         env_var="LOCUST_EXPORTER",
         help="Exports Locust stats to Timescale",
     )

locust_cloud/cloud.py

@@ -92,7 +92,6 @@ parser.add_argument(
     env_var="LOCUST_USERS",
 )
 parser.add_argument(
-    "-r",
     "--requirements",
     type=str,
     help="Optional requirements.txt file that contains your external libraries.",
@@ -169,6 +168,11 @@ parser.add_argument(
     env_var="LOCUST_CLOUD_WORKERS",
     default=None,
 )
+parser.add_argument(
+    "--delete",
+    action="store_true",
+    help="Delete a running cluster. Useful if locust-cloud was killed/disconnected or if there was an error.",
+)
 
 options, locust_options = parser.parse_known_args()
 options: Namespace
@@ -218,6 +222,10 @@ def main() -> None:
         aws_secret_access_key = credentials.get("secret_key")
         aws_session_token = credentials.get("token", "")
 
+        if options.delete:
+            delete(s3_bucket, credential_manager)
+            return
+
         logger.info(f"Uploading {options.locustfile}")
         s3 = credential_manager.session.client("s3")
         try:
@@ -257,7 +265,18 @@ def main() -> None:
         locust_env_variables = [
             {"name": env_variable, "value": str(os.environ[env_variable])}
             for env_variable in os.environ
-            if env_variable.startswith("LOCUST_") and os.environ[env_variable]
+            if env_variable.startswith("LOCUST_")
+            and not env_variable.startswith("LOCUST_CLOUD")
+            and not env_variable
+            in [
+                "LOCUST_LOCUSTFILE",
+                "LOCUST_USERS",
+                "LOCUST_WEB_HOST_DISPLAY_NAME",
+                "LOCUST_API_BASE_URL",
+                "LOCUST_SKIP_MONKEY_PATCH",
+                "LOCUST_REQUIREMENTS_URL",
+            ]
+            and os.environ[env_variable]
         ]
         deploy_endpoint = f"{options.lambda_url}/{options.kube_cluster_name}"
         payload = {
@@ -289,9 +308,12 @@ def main() -> None:
         if response.status_code == 200:
             deployed_pods = response.json().get("pods", [])
         else:
-            logger.error(
-                f"HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
-            )
+            try:
+                logger.error(f"Error when deploying: {response.json()['Message']}")
+            except Exception:
+                logger.error(
+                    f"HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
+                )
             sys.exit(1)
     except CredentialError as ce:
         logger.error(f"Credential error: {ce}")
@@ -362,43 +384,47 @@ def main() -> None:
         logger.exception(e)
         sys.exit(1)
     finally:
-        try:
-            logger.info("Tearing down Locust cloud...")
-            credential_manager.refresh_credentials()
-            refreshed_credentials = credential_manager.get_current_credentials()
-
-            headers = {
-                "AWS_ACCESS_KEY_ID": refreshed_credentials.get("access_key", ""),
-                "AWS_SECRET_ACCESS_KEY": refreshed_credentials.get("secret_key", ""),
-                "Authorization": f"Bearer {refreshed_credentials.get('cognito_client_id_token', '')}",
-            }
-
-            token = refreshed_credentials.get("token")
-            if token:
-                headers["AWS_SESSION_TOKEN"] = token
-
-            response = requests.delete(
-                f"{options.lambda_url}/{options.kube_cluster_name}",
-                headers=headers,
-                params={"namespace": options.kube_namespace} if options.kube_namespace else {},
-            )
+        delete(s3_bucket, credential_manager)
 
-            if response.status_code != 200:
-                logger.error(
-                    f"Could not automatically tear down Locust Cloud: HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
-                )
-        except Exception as e:
-            logger.error(f"Could not automatically tear down Locust Cloud: {e}")
 
-        try:
-            logger.info("Cleaning up locust files")
-            s3 = credential_manager.session.resource("s3")
-            bucket = s3.Bucket(s3_bucket)
-            bucket.objects.all().delete()
-            logger.info("Done! ✨")
-        except ClientError as e:
-            logger.error(f"Failed to clean up locust files: {e}")
-            sys.exit(1)
+def delete(s3_bucket, credential_manager):
+    try:
+        logger.info("Tearing down Locust cloud...")
+        credential_manager.refresh_credentials()
+        refreshed_credentials = credential_manager.get_current_credentials()
+
+        headers = {
+            "AWS_ACCESS_KEY_ID": refreshed_credentials.get("access_key", ""),
+            "AWS_SECRET_ACCESS_KEY": refreshed_credentials.get("secret_key", ""),
+            "Authorization": f"Bearer {refreshed_credentials.get('cognito_client_id_token', '')}",
+        }
+
+        token = refreshed_credentials.get("token")
+        if token:
+            headers["AWS_SESSION_TOKEN"] = token
+
+        response = requests.delete(
+            f"{options.lambda_url}/{options.kube_cluster_name}",
+            headers=headers,
+            params={"namespace": options.kube_namespace} if options.kube_namespace else {},
+        )
+
+        if response.status_code != 200:
+            logger.error(
+                f"Could not automatically tear down Locust Cloud: HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
+            )
+    except Exception as e:
+        logger.error(f"Could not automatically tear down Locust Cloud: {e}")
+
+    try:
+        logger.debug("Cleaning up locustfiles")
+        s3 = credential_manager.session.resource("s3")
+        bucket = s3.Bucket(s3_bucket)
+        bucket.objects.all().delete()
+        logger.info("Done! ✨")
+    except ClientError as e:
+        logger.error(f"Failed to clean up locust files: {e}")
+        sys.exit(1)
 
 
 if __name__ == "__main__":

locust_cloud/timescale/exporter.py

@@ -195,7 +195,7 @@ VALUES (%(time)s, %(run_id)s, %(greenlet_id)s, %(loadgen)s, %(name)s, %(request_
                 "INSERT INTO testruns (id, num_users, description, arguments) VALUES (%s,%s,%s,%s)",
                 (
                     self._run_id,
-                    self.env.parsed_options.num_users or self.env.runner.user_count or 0,
+                    self.env.runner.target_user_count if self.env.runner else 1,
                     "self.env.parsed_options.description",
                     " ".join(cmd),
                 ),

{locust_cloud-1.2.12.dist-info → locust_cloud-1.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: locust-cloud
-Version: 1.2.12
+Version: 1.3.2
 Summary: Locust Cloud
 Project-URL: Homepage, https://locust.cloud
 Requires-Python: >=3.11

{locust_cloud-1.2.12.dist-info → locust_cloud-1.3.2.dist-info}/RECORD

@@ -1,9 +1,9 @@
-locust_cloud/__init__.py,sha256=d-vcft9odW37rtliWSquU4h-BQNlZftSKDoZTk90HsM,1970
+locust_cloud/__init__.py,sha256=GS_nha_uwtF2MTYEpkeiK-YmbY9ZlYBgDIyvPzQ0iGQ,2005
 locust_cloud/auth.py,sha256=2fcTIq0tYDckDsZZ7yFTHZn5nG7Ir1dlV1QIu_tp3JA,3164
-locust_cloud/cloud.py,sha256=gEiP2Lq7iQbWYiuuDbLmMgv5CUTUrYJ1uRzAgFv2cHw,15062
+locust_cloud/cloud.py,sha256=3i91W4FLTij5919tnn2VKOtFbMnzlMwsDgrNBcXwpGo,15809
 locust_cloud/constants.py,sha256=YGgIEnaOT9cAA8F8a8Z8lOmlI78D2fXiVpP0RF213MA,167
 locust_cloud/credential_manager.py,sha256=JbuS22RTGAmngVB9DFLxHIgaonCSImemMGvtWnxBZRE,5515
-locust_cloud/timescale/exporter.py,sha256=j1K-czMUdOLFt5Sop0PTF6TL52vYdngEEI5OEM3Xl60,11128
+locust_cloud/timescale/exporter.py,sha256=0VodnUUFHMDLH3Ma0dz2lfOhgjaisU9rn48fpckBYhI,11119
 locust_cloud/webui/.eslintrc,sha256=huuuTZNwbBGkhfelEOMhLQ-tgC9OYnnX0DRYgBn83Ig,1146
 locust_cloud/webui/.gitignore,sha256=i8EHIqDlVm1-Dkvf_GTZmP_Wu99GE7ABfbYzHXOU734,54
 locust_cloud/webui/.prettierrc,sha256=qZpzAJPlHUQ3--NHYK8WXgc_TlIqq-eoz4pbkTKK3nk,167
@@ -15,7 +15,7 @@ locust_cloud/webui/vite.config.ts,sha256=cqxPMkbwEA3H9mGGbuPulQUhIHCosUqm_1usxzs
 locust_cloud/webui/yarn.lock,sha256=HQipc2d3thDdUieYtmTLSOxMlgpN-nQjVWKusAlvf2g,227106
 locust_cloud/webui/dist/index.html,sha256=uRX9dnfQEMYFWaKl0czXN_2JbtQF0WRgP4gEOWG0WpQ,664
 locust_cloud/webui/dist/assets/index-BzflLAyP.js,sha256=GR85cK07EBnN6PBKImgYlHdsBl40y0Tay9J8w9LZKqc,2794820
-locust_cloud-1.2.12.dist-info/METADATA,sha256=rOtFHGBE1-1QiHrAutxSdcvr5AH9wIef8XvRUelNpxA,1664
-locust_cloud-1.2.12.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-locust_cloud-1.2.12.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
-locust_cloud-1.2.12.dist-info/RECORD,,
+locust_cloud-1.3.2.dist-info/METADATA,sha256=bK1z7s7Tk-pkEItcPpn2-OQo-CSvvV0WjqrbAM6ROSU,1663
+locust_cloud-1.3.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+locust_cloud-1.3.2.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
+locust_cloud-1.3.2.dist-info/RECORD,,