locust-cloud 1.9.0__py3-none-any.whl → 1.11.0__py3-none-any.whl

locust_cloud/__init__.py

@@ -1,10 +1,10 @@
 import importlib.metadata
 import os
+import sys
 
 os.environ["LOCUST_SKIP_MONKEY_PATCH"] = "1"
 __version__ = importlib.metadata.version("locust-cloud")
 
-import argparse
 import logging
 
 import configargparse
@@ -19,38 +19,29 @@ from locust_cloud.timescale.query import register_query
 from psycopg.conninfo import make_conninfo
 from psycopg_pool import ConnectionPool
 
-GRAPH_VIEWER = os.environ.get("GRAPH_VIEWER")
 logger = logging.getLogger(__name__)
 
 
 @events.init_command_line_parser.add_listener
 def add_arguments(parser: LocustArgumentParser):
-    if not (os.environ.get("PGHOST") or GRAPH_VIEWER):
+    if not (os.environ.get("PGHOST")):
         parser.add_argument_group(
             "locust-cloud",
             "locust-cloud disabled, because PGHOST was not set - this is normal for local runs",
         )
         return
 
+    try:
+        REGION = os.environ["AWS_DEFAULT_REGION"]
+    except KeyError:
+        logger.fatal("Missing AWS_DEFAULT_REGION env var")
+        sys.exit(1)
+
     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",
     )
-    locust_cloud.add_argument(
-        "--exporter",
-        default=True,
-        action=argparse.BooleanOptionalAction,
-        env_var="LOCUST_EXPORTER",
-        help="Exports Locust stats to Timescale",
-    )
-    locust_cloud.add_argument(
-        "--description",
-        type=str,
-        env_var="LOCUST_DESCRIPTION",
-        default="",
-        help="Description of the test being run",
-    )
     # do not set
     # used for sending the run id from master to workers
     locust_cloud.add_argument(
@@ -59,6 +50,27 @@ def add_arguments(parser: LocustArgumentParser):
         env_var="LOCUSTCLOUD_RUN_ID",
         help=configargparse.SUPPRESS,
     )
+    locust_cloud.add_argument(
+        "--allow-signup",
+        env_var="LOCUSTCLOUD_ALLOW_SIGNUP",
+        help=configargparse.SUPPRESS,
+        default=False,
+        action="store_true",
+    )
+    locust_cloud.add_argument(
+        "--graph-viewer",
+        env_var="LOCUSTCLOUD_GRAPH_VIEWER",
+        help=configargparse.SUPPRESS,
+        default=False,
+        action="store_true",
+    )
+    locust_cloud.add_argument(
+        "--deployer-url",
+        type=str,
+        env_var="LOCUSTCLOUD_DEPLOYER_URL",
+        help=configargparse.SUPPRESS,
+        default=f"https://api.{REGION}.locust.cloud/1",
+    )
 
 
 def set_autocommit(conn: psycopg.Connection):
@@ -86,14 +98,12 @@ def on_locust_init(environment: locust.env.Environment, **_args):
         logger.exception(e)
         raise
 
-    if not GRAPH_VIEWER:
+    if not environment.parsed_options.graph_viewer:
         IdleExit(environment)
-
-    if not GRAPH_VIEWER and environment.parsed_options and environment.parsed_options.exporter:
         Exporter(environment, pool)
 
     if environment.web_ui:
-        if GRAPH_VIEWER:
+        if environment.parsed_options.graph_viewer:
             environment.web_ui.template_args["isGraphViewer"] = True
 
         register_auth(environment)

locust_cloud/auth.py

@@ -11,11 +11,6 @@ from flask_login import UserMixin, login_user
 from locust.html import render_template_from
 from locust_cloud import __version__
 
-REGION = os.environ.get("AWS_DEFAULT_REGION")
-DEPLOYER_URL = f"https://api.{REGION}.locust.cloud/1"
-ALLOW_SIGNUP = os.environ.get("ALLOW_SIGNUP", True)
-
-
 logger = logging.getLogger(__name__)
 
 
@@ -47,7 +42,7 @@ def set_credentials(username: str, credentials: Credentials, response: werkzeug.
 
 
 def register_auth(environment: locust.env.Environment):
-    environment.web_ui.app.config["SECRET_KEY"] = os.getenv("SECRET_KEY")
+    environment.web_ui.app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "") + os.getenv("CUSTOMER_ID", "")
     environment.web_ui.app.debug = False
 
     web_base_path = environment.parsed_options.web_base_path
@@ -71,7 +66,7 @@ def register_auth(environment: locust.env.Environment):
         },
     )
 
-    if ALLOW_SIGNUP:
+    if environment.parsed_options.allow_signup:
         environment.web_ui.auth_args["auth_providers"] = [
             {"label": "Sign Up", "callback_url": f"{web_base_path}/signup"}
         ]
@@ -83,7 +78,7 @@ def register_auth(environment: locust.env.Environment):
 
         try:
             auth_response = requests.post(
-                f"{DEPLOYER_URL}/auth/login",
+                f"{environment.parsed_options.deployer_url}/auth/login",
                 json={"username": username, "password": password},
                 headers={"X-Client-Version": __version__},
             )
@@ -91,6 +86,11 @@ def register_auth(environment: locust.env.Environment):
             auth_response.raise_for_status()
 
             credentials = auth_response.json()
+
+            if os.getenv("CUSTOMER_ID", "") and credentials["customer_id"] != os.getenv("CUSTOMER_ID", ""):
+                session["auth_error"] = "Invalid login for this deployment"
+                return redirect(url_for("locust.login"))
+
             response = redirect(url_for("locust.index"))
             response = set_credentials(username, credentials, response)
             login_user(AuthUser(credentials["user_sub_id"]))
@@ -108,7 +108,7 @@ def register_auth(environment: locust.env.Environment):
 
     @auth_blueprint.route("/signup")
     def signup():
-        if not ALLOW_SIGNUP:
+        if not environment.parsed_options.allow_signup:
             return redirect(url_for("locust.login"))
 
         if session.get("username"):
@@ -118,6 +118,7 @@ def register_auth(environment: locust.env.Environment):
                         {
                             "label": "Confirmation Code",
                             "name": "confirmation_code",
+                            "is_required": True,
                         },
                     ],
                     "callback_url": f"{web_base_path}/confirm-signup",
@@ -131,20 +132,30 @@ def register_auth(environment: locust.env.Environment):
                         {
                             "label": "Username",
                             "name": "username",
+                            "is_required": True,
+                            "type": "email",
+                        },
+                        {
+                            "label": "Full Name",
+                            "name": "customer_name",
+                            "is_required": True,
                         },
                         {
                             "label": "Password",
                             "name": "password",
                             "is_secret": True,
+                            "is_required": True,
                         },
                         {
                             "label": "Access Code",
                             "name": "access_code",
+                            "is_required": True,
                         },
                         {
-                            "label": "I consent to:\n\n1. Only test your own website/service or our example example target\n\n2. Only use locust-cloud for its intended purpose: to load test other sites/services.\n\n3. Not attempt to circumvent your account limitations (e.g. max user count or max request count)\n\n4. Not use personal data (real names, addresses etc) in your tests.",
+                            "label": "I consent to:\n\n1. Only test your own website/service or our example target\n\n2. Only use locust-cloud for its intended purpose: to load test other sites/services.\n\n3. Not attempt to circumvent your account limitations (e.g. max user count or max request count)\n\n4. Not use personal data (real names, addresses etc) in your tests.",
                             "name": "consent",
                             "default_value": False,
+                            "is_required": True,
                         },
                     ],
                     "callback_url": f"{web_base_path}/create-account",
@@ -164,32 +175,28 @@ def register_auth(environment: locust.env.Environment):
 
     @auth_blueprint.route("/create-account", methods=["POST"])
     def create_account():
-        if not ALLOW_SIGNUP:
+        if not environment.parsed_options.allow_signup:
             return redirect(url_for("locust.login"))
 
         session["auth_sign_up_error"] = ""
         session["auth_info"] = ""
 
         username = request.form.get("username", "")
+        customer_name = request.form.get("customer_name", "")
         password = request.form.get("password")
         access_code = request.form.get("access_code")
-        has_consented = request.form.get("consent")
-
-        if not has_consented:
-            session["auth_sign_up_error"] = "Please accept the terms and conditions to create an account."
-
-            return redirect(url_for("locust_cloud_auth.signup"))
 
         try:
             auth_response = requests.post(
-                f"{DEPLOYER_URL}/auth/signup",
+                f"{environment.parsed_options.deployer_url}/auth/signup",
                 json={"username": username, "password": password, "access_code": access_code},
             )
 
             auth_response.raise_for_status()
 
-            session["user_sub"] = auth_response.json().get("user_sub")
+            session["user_sub_id"] = auth_response.json().get("user_sub_id")
             session["username"] = username
+            session["customer_name"] = customer_name
             session["auth_info"] = (
                 "Please check your email and enter the confirmation code. If you didn't get a code after one minute, you can [request a new one](/resend-code)"
             )
@@ -206,7 +213,7 @@ def register_auth(environment: locust.env.Environment):
     def resend_code():
         try:
             auth_response = requests.post(
-                "http://localhost:8000/1/auth/resend-confirmation",
+                f"{environment.parsed_options.deployer_url}/auth/resend-confirmation",
                 json={"username": session["username"]},
             )
 
@@ -225,7 +232,7 @@ def register_auth(environment: locust.env.Environment):
 
     @auth_blueprint.route("/confirm-signup", methods=["POST"])
     def confirm_signup():
-        if not ALLOW_SIGNUP:
+        if not environment.parsed_options.allow_signup:
             return redirect(url_for("locust.login"))
 
         session["auth_sign_up_error"] = ""
@@ -233,10 +240,11 @@ def register_auth(environment: locust.env.Environment):
 
         try:
             auth_response = requests.post(
-                f"{DEPLOYER_URL}/auth/confirm-signup",
+                f"{environment.parsed_options.deployer_url}/auth/confirm-signup",
                 json={
                     "username": session.get("username"),
-                    "user_sub": session["user_sub"],
+                    "customer_name": session.get("customer_name"),
+                    "user_sub_id": session["user_sub_id"],
                     "confirmation_code": confirmation_code,
                 },
             )
@@ -247,7 +255,7 @@ def register_auth(environment: locust.env.Environment):
             session["auth_info"] = "Account created successfully!"
             session["auth_sign_up_error"] = ""
 
-            return redirect(url_for("locust.login"))
+            return redirect("https://docs.locust.cloud/")
         except requests.exceptions.HTTPError as e:
             message = e.response.json().get("Message", "An unexpected error occured. Please try again.")
             session["auth_info"] = ""

locust_cloud/cloud.py

@@ -148,6 +148,12 @@ parser.add_argument(
     default="latest",
     help=configargparse.SUPPRESS,  # overrides the locust-cloud docker image tag. for internal use
 )
+parser.add_argument(
+    "--mock-server",
+    action="store_true",
+    default=False,
+    help=configargparse.SUPPRESS,
+)
 
 options, locust_options = parser.parse_known_args()
 options: Namespace
@@ -166,7 +172,7 @@ logging.getLogger("requests").setLevel(logging.INFO)
 logging.getLogger("urllib3").setLevel(logging.INFO)
 
 
-api_url = f"https://api.{options.region}.locust.cloud/1"
+api_url = os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", f"https://api.{options.region}.locust.cloud/1")
 
 
 def main() -> None:
@@ -175,6 +181,8 @@ def main() -> None:
             "Setting a region is required to use Locust Cloud. Please ensure the AWS_DEFAULT_REGION env variable or the --region flag is set."
         )
         sys.exit(1)
+    if options.region:
+        os.environ["AWS_DEFAULT_REGION"] = options.region
 
     s3_bucket = "dmdb-default" if options.region == "us-east-1" else "locust-default"
     deployments: list[Any] = []
@@ -203,17 +211,18 @@ def main() -> None:
         aws_session_token = credentials.get("token", "")
 
         if options.delete:
-            delete(s3_bucket, credential_manager)
+            delete(credential_manager)
             return
 
         logger.info(f"Uploading {options.locustfile}")
         logger.debug(f"... to {s3_bucket}")
         s3 = credential_manager.session.client("s3")
         try:
-            s3.upload_file(options.locustfile, s3_bucket, os.path.basename(options.locustfile))
+            filename = options.username + "__" + os.path.basename(options.locustfile)
+            s3.upload_file(options.locustfile, s3_bucket, filename)
             locustfile_url = s3.generate_presigned_url(
                 ClientMethod="get_object",
-                Params={"Bucket": s3_bucket, "Key": os.path.basename(options.locustfile)},
+                Params={"Bucket": s3_bucket, "Key": filename},
                 ExpiresIn=3600,
             )
             logger.debug(f"Uploaded {options.locustfile} successfully")
@@ -228,10 +237,11 @@ def main() -> None:
         if options.requirements:
             logger.info(f"Uploading {options.requirements}")
             try:
-                s3.upload_file(options.requirements, s3_bucket, "requirements.txt")
+                filename = options.username + "__" + "requirements.txt"
+                s3.upload_file(options.requirements, s3_bucket, filename)
                 requirements_url = s3.generate_presigned_url(
                     ClientMethod="get_object",
-                    Params={"Bucket": s3_bucket, "Key": "requirements.txt"},
+                    Params={"Bucket": s3_bucket, "Key": filename},
                     ExpiresIn=3600,
                 )
                 logger.debug(f"Uploaded {options.requirements} successfully")
@@ -268,6 +278,7 @@ def main() -> None:
             ],
             "user_count": options.users,
             "image_tag": options.image_tag,
+            "mock_server": options.mock_server,
         }
         if options.workers is not None:
             payload["worker_count"] = options.workers
@@ -377,10 +388,10 @@ def main() -> None:
         logger.exception(e)
         sys.exit(1)
     finally:
-        delete(s3_bucket, credential_manager)
+        delete(credential_manager)
 
 
-def delete(s3_bucket, credential_manager):
+def delete(credential_manager):
     try:
         logger.info("Tearing down Locust cloud...")
         credential_manager.refresh_credentials()
@@ -412,9 +423,9 @@ def delete(s3_bucket, credential_manager):
 
     try:
         logger.debug("Cleaning up locustfiles")
-        s3 = credential_manager.session.resource("s3")
-        bucket = s3.Bucket(s3_bucket)
-        bucket.objects.all().delete()
+        # s3 = credential_manager.session.resource("s3")
+        # bucket = s3.Bucket(s3_bucket)
+        # bucket.objects.all().delete()
     except ClientError as e:
         logger.debug(f"Failed to clean up locust files: {e}")
         # sys.exit(1)

locust_cloud/timescale/exporter.py

@@ -217,7 +217,7 @@ class Exporter:
         cmd = sys.argv[1:]
         with self.pool.connection() as conn:
             conn.execute(
-                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, description, arguments, customer) VALUES (%s,%s,%s,%s,%s,%s,%s,current_user)",
+                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, arguments, customer) VALUES (%s,%s,%s,%s,%s,%s,current_user)",
                 (
                     self._run_id,
                     self.env.runner.target_user_count if self.env.runner else 1,
@@ -228,8 +228,7 @@ class Exporter:
                     )
                     else 0,
                     self.env.web_ui.template_args.get("username", "") if self.env.web_ui else "",
-                    self.env.parsed_locustfiles[0].split("/")[-1],
-                    self.env.parsed_options.description,
+                    self.env.parsed_locustfiles[0].split("/")[-1].split("__")[-1],
                     " ".join(cmd),
                 ),
             )

locust_cloud/timescale/queries.py

@@ -262,7 +262,7 @@ ORDER BY a.time
 
 total_vuh = """
 SELECT
-  SUM((end_time - id) * num_users) AS "totalVuh"
+  COALESCE(SUM((end_time - id) * num_users), '0') AS "totalVuh"
 FROM testruns
 WHERE id >= date_trunc('month', NOW()) AND NOT refund
 """
@@ -274,7 +274,7 @@ SELECT
   max_users as "maxUsers",
   users_per_worker as "usersPerWorker"
 FROM customers
-WHERE customer = current_user
+WHERE id = current_user
 """
 
 queries: dict["str", LiteralString] = {