locust-cloud 1.2.11__py3-none-any.whl → 1.3.1__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
@@ -17,7 +19,7 @@ GRAPH_VIEWER = os.environ.get("GRAPH_VIEWER")
 
 @events.init_command_line_parser.add_listener
 def add_arguments(parser: LocustArgumentParser):
-    if not PG_HOST:
+    if not (PG_HOST or GRAPH_VIEWER):
         parser.add_argument_group(
             "locust-cloud",
             "locust-cloud disabled, because PG_HOST was not set - this is normal for local runs",
@@ -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__":