locust-cloud 1.0.6__py3-none-any.whl → 1.0.8__py3-none-any.whl

locust_cloud/__init__.py

@@ -1,10 +1,11 @@
-from locust_cloud.auth import register_auth
-from locust_cloud.timescale.exporter import Timescale
-
 import os
 
+os.environ["LOCUST_SKIP_MONKEY_PATCH"] = "1"
+
 from locust import events
 from locust.argument_parser import LocustArgumentParser
+from locust_cloud.auth import register_auth
+from locust_cloud.timescale.exporter import Timescale
 
 PG_USER = os.environ.get("PG_USER")
 PG_HOST = os.environ.get("PG_HOST")
@@ -15,11 +16,18 @@ PG_PORT = os.environ.get("PG_PORT", 5432)
 
 @events.init_command_line_parser.add_listener
 def add_arguments(parser: LocustArgumentParser):
+    if not PG_HOST:
+        parser.add_argument_group(
+            "locust-cloud",
+            "locust-cloud disabled, because PG_HOST was not set - this is normal for local runs",
+        )
+        return
+
+    os.environ["LOCUST_BUILD_PATH"] = os.path.join(os.path.dirname(__file__), "webui/dist")
     locust_cloud = parser.add_argument_group(
         "locust-cloud",
-        "Arguments for use with Locust cloud!",
+        "Arguments for use with Locust cloud",
     )
-
     locust_cloud.add_argument(
         "--exporter",
         default=True,
@@ -38,7 +46,8 @@ def add_arguments(parser: LocustArgumentParser):
 
 @events.init.add_listener
 def on_locust_init(environment, **args):
-    os.environ["LOCUST_BUILD_PATH"] = os.path.join(os.path.dirname(__file__), "webui/dist")
+    if not PG_HOST:
+        return
 
     if environment.parsed_options.exporter:
         Timescale(

locust_cloud/auth.py

@@ -1,5 +1,5 @@
-import datetime
 import os
+from datetime import UTC, datetime, timedelta
 
 import requests
 from flask import redirect, request, url_for
@@ -21,9 +21,9 @@ def set_credentials(credentials, response):
     user_sub_id = credentials["user_sub_id"]
     refresh_token = credentials["refresh_token"]
 
-    response.set_cookie("cognito_token", id_token, expires=datetime.datetime.utcnow() + datetime.timedelta(days=1))
-    response.set_cookie("user_token", refresh_token, expires=datetime.datetime.utcnow() + datetime.timedelta(days=365))
-    response.set_cookie("user_sub_id", user_sub_id, expires=datetime.datetime.utcnow() + datetime.timedelta(days=365))
+    response.set_cookie("cognito_token", id_token, expires=datetime.now(tz=UTC) + timedelta(days=1))
+    response.set_cookie("user_token", refresh_token, expires=datetime.now(tz=UTC) + timedelta(days=365))
+    response.set_cookie("user_sub_id", user_sub_id, expires=datetime.now(tz=UTC) + timedelta(days=365))
 
     return response
 

locust_cloud/cloud.py

@@ -15,6 +15,7 @@ from botocore.exceptions import ClientError
 LAMBDA = "https://deployer.locust.cloud/1"
 DEFAULT_CLUSTER_NAME = "locust"
 DEFAULT_REGION_NAME = "eu-north-1"
+LOCUST_ENV_VARIABLE_IGNORE_LIST = ["LOCUST_BUILD_PATH", "LOCUST_SKIP_MONKEY_PATCH"]
 
 
 class LocustTomlConfigParser(configargparse.TomlConfigParser):
@@ -275,7 +276,7 @@ def deploy(
     locust_env_variables = [
         {"name": env_variable, "value": str(os.environ[env_variable])}
         for env_variable in os.environ
-        if env_variable.startswith("LOCUST_")
+        if env_variable.startswith("LOCUST_") and env_variable not in LOCUST_ENV_VARIABLE_IGNORE_LIST
     ]
 
     response = requests.post(