locust-cloud 1.14.0__py3-none-any.whl → 1.14.2__py3-none-any.whl

locust_cloud/cloud.py

@@ -1,6 +1,7 @@
 import base64
 import gzip
 import importlib.metadata
+import io
 import json
 import logging
 import os
@@ -11,10 +12,12 @@ import time
 import tomllib
 import urllib.parse
 import webbrowser
-from argparse import Namespace
+from argparse import ArgumentTypeError, Namespace
 from collections import OrderedDict
+from collections.abc import Generator
 from dataclasses import dataclass
 from typing import IO, Any
+from zipfile import ZipFile
 
 import configargparse
 import jwt
@@ -24,6 +27,7 @@ import socketio
 import socketio.exceptions
 
 __version__ = importlib.metadata.version("locust-cloud")
+CWD = pathlib.Path.cwd()
 
 
 class LocustTomlConfigParser(configargparse.TomlConfigParser):
@@ -48,6 +52,51 @@ class LocustTomlConfigParser(configargparse.TomlConfigParser):
         return result
 
 
+def valid_extra_files_path(file_path: str) -> pathlib.Path:
+    p = pathlib.Path(file_path).resolve()
+
+    if not CWD in p.parents:
+        raise ArgumentTypeError(f"Can only reference files under current working directory: {CWD}")
+    if not p.exists():
+        raise ArgumentTypeError(f"File not found: {file_path}")
+    return p
+
+
+def transfer_encode(file_name: str, stream: IO[bytes]) -> dict[str, str]:
+    return {
+        "filename": file_name,
+        "data": base64.b64encode(gzip.compress(stream.read())).decode(),
+    }
+
+
+def transfer_encoded_file(file_path: str) -> dict[str, str]:
+    try:
+        with open(file_path, "rb") as f:
+            return transfer_encode(file_path, f)
+    except FileNotFoundError:
+        raise ArgumentTypeError(f"File not found: {file_path}")
+
+
+def transfer_encoded_extra_files(paths: list[pathlib.Path]) -> dict[str, str]:
+    def expanded(paths: list[pathlib.Path]) -> Generator[pathlib.Path, None, None]:
+        for path in paths:
+            if path.is_dir():
+                for root, _, file_names in os.walk(path):
+                    for file_name in file_names:
+                        yield pathlib.Path(root) / file_name
+            else:
+                yield path
+
+    buffer = io.BytesIO()
+
+    with ZipFile(buffer, "w") as zf:
+        for path in set(expanded(paths)):
+            zf.write(path.relative_to(CWD))
+
+    buffer.seek(0)
+    return transfer_encode("extra-files.zip", buffer)
+
+
 parser = configargparse.ArgumentParser(
     default_config_files=[
         "~/.locust.conf",
@@ -93,6 +142,7 @@ parser.add_argument(
     default="locustfile.py",
     help="The Python file that contains your test. Defaults to 'locustfile.py'.",
     env_var="LOCUST_LOCUSTFILE",
+    type=transfer_encoded_file,
 )
 parser.add_argument(
     "-u",
@@ -112,14 +162,14 @@ advanced.add_argument(
 )
 advanced.add_argument(
     "--requirements",
-    type=str,
+    type=transfer_encoded_file,
     help="Optional requirements.txt file that contains your external libraries.",
 )
 advanced.add_argument(
     "--login",
     action="store_true",
     default=False,
-    help=configargparse.SUPPRESS,
+    help="Launch an interactive session to authenticate your user.\nOnce completed your credentials will be stored and automatically refreshed for quite a long time.\nOnce those expires you will be prompted to perform another login.",
 )
 advanced.add_argument(
     "--non-interactive",
@@ -155,11 +205,15 @@ parser.add_argument(
     type=str,
     help="Set a profile to group the testruns together",
 )
+parser.add_argument(
+    "--extra-files",
+    nargs="*",
+    type=valid_extra_files_path,
+    help="A list of extra files or directories to upload. Space-separated, e.g. --extra-files testdata.csv *.py my-directory/",
+)
 
-options, locust_options = parser.parse_known_args()
-
-options: Namespace
-locust_options: list
+parsed_args: tuple[Namespace, list[str]] = parser.parse_known_args()
+options, locust_options = parsed_args
 
 logging.basicConfig(
     format="[LOCUST-CLOUD] %(levelname)s: %(message)s",
@@ -281,7 +335,7 @@ class ApiSession(requests.Session):
 
             if not all([username, password, region]):
                 print(
-                    "Running with --non-interaction requires that LOCUSTCLOUD_USERNAME, LOCUSTCLOUD_PASSWORD and LOCUSTCLOUD_REGION environment variables are set."
+                    "Running with --non-interactive requires that LOCUSTCLOUD_USERNAME, LOCUSTCLOUD_PASSWORD and LOCUSTCLOUD_REGION environment variables are set."
                 )
                 sys.exit(1)
 
@@ -348,8 +402,9 @@ class ApiSession(requests.Session):
             sys.exit(1)
 
         # TODO: Technically the /login endpoint can return a challenge for you
-        #       to change your password. Don't know how we should handle that
-        #       in the cli.
+        #       to change your password.
+        #       Now that we have a web based login flow we should force them to
+        #       do a locust-cloud --login if we get that.
 
         id_token = response.json()["cognito_client_id_token"]
         decoded = jwt.decode(id_token, options={"verify_signature": False})
@@ -587,23 +642,6 @@ def main() -> None:
         sys.exit()
 
     try:
-        try:
-            with open(options.locustfile, "rb") as f:
-                locustfile_data = base64.b64encode(gzip.compress(f.read())).decode()
-        except FileNotFoundError:
-            logger.error(f"File not found: {options.locustfile}")
-            sys.exit(1)
-
-        requirements_data = None
-
-        if options.requirements:
-            try:
-                with open(options.requirements, "rb") as f:
-                    requirements_data = base64.b64encode(gzip.compress(f.read())).decode()
-            except FileNotFoundError:
-                logger.error(f"File not found: {options.requirements}")
-                sys.exit(1)
-
         logger.info("Deploying load generators")
         locust_env_variables = [
             {"name": env_variable, "value": os.environ[env_variable]}
@@ -625,7 +663,7 @@ def main() -> None:
                 {"name": "LOCUSTCLOUD_PROFILE", "value": options.profile},
                 *locust_env_variables,
             ],
-            "locustfile": {"filename": options.locustfile, "data": locustfile_data},
+            "locustfile": options.locustfile,
             "user_count": options.users,
             "mock_server": options.mock_server,
         }
@@ -637,7 +675,10 @@ def main() -> None:
             payload["worker_count"] = options.workers
 
         if options.requirements:
-            payload["requirements"] = {"filename": options.requirements, "data": requirements_data}
+            payload["requirements"] = options.requirements
+
+        if options.extra_files:
+            payload["extra_files"] = transfer_encoded_extra_files(options.extra_files)
 
         try:
             response = session.post("/deploy", json=payload)
@@ -658,7 +699,7 @@ def main() -> None:
         session_id = response.json()["session_id"]
         logger.debug(f"Session ID is {session_id}")
 
-        logger.info("Waiting for pods to be ready...")
+        logger.info("Waiting for load generators to be ready...")
         websocket.connect(
             log_ws_url,
             auth=session_id,
@@ -708,8 +749,6 @@ def delete(session):
     except Exception as e:
         logger.error(f"Could not automatically tear down Locust Cloud: {e.__class__.__name__}:{e}")
 
-    logger.info("Done! ✨")  # FIXME: Should probably not say it's done since at this point it could still be running
-
 
 if __name__ == "__main__":
     main()

{locust_cloud-1.14.0.dist-info → locust_cloud-1.14.2.dist-info}/METADATA

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

locust_cloud-1.14.2.dist-info/RECORD

@@ -0,0 +1,5 @@
+locust_cloud/cloud.py,sha256=ro6oWsd6SEjBHDUWH7vrnvo0wY1Ljz8zeiJwhlaIF-g,26873
+locust_cloud-1.14.2.dist-info/METADATA,sha256=it51qVDFVV6aTB0b3bAuIvF9xfIVdx6Wk6UB31PLzhs,497
+locust_cloud-1.14.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+locust_cloud-1.14.2.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
+locust_cloud-1.14.2.dist-info/RECORD,,

locust_cloud-1.14.0.dist-info/RECORD

@@ -1,5 +0,0 @@
-locust_cloud/cloud.py,sha256=H_bATzG3Hyw8BRw0HmgLKTl5kX4y9KQuKCVTtzSAdjA,25466
-locust_cloud-1.14.0.dist-info/METADATA,sha256=mZOGdXwLoyzOO3TeABH2eKsnw6RTz2d5Av0POdyvaIc,497
-locust_cloud-1.14.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-locust_cloud-1.14.0.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
-locust_cloud-1.14.0.dist-info/RECORD,,