locust-cloud 1.20.3__py3-none-any.whl → 1.20.4__py3-none-any.whl

locust_cloud/__init__.py

@@ -1,26 +1,3 @@
-from locust import events
-from locust_cloud.args import cloud_parser
 from locust_cloud.cloud import main
 
 __all__ = ["main"]
-
-
-@events.init_command_line_parser.add_listener
-def _(parser):
-    cloud_group = parser.add_argument_group(
-        "Locust Cloud",
-        """Launches a distributed Locust run on locust.cloud infrastructure.
-
-Example: locust --cloud -f my_locustfile.py --users 1000 ...""",
-    )
-
-    # This arguments is defined here because only makes sense when
-    # running from locust core
-    cloud_group.add_argument(
-        "--cloud",
-        action="store_true",
-        help="Run Locust in cloud mode.",
-    )
-
-    for action in cloud_parser._actions:
-        cloud_group._add_action(action)

locust_cloud/apisession.py

@@ -3,7 +3,6 @@ import os
 import sys
 import time
 
-import jwt
 import requests
 from locust_cloud.common import VALID_REGIONS, __version__, get_api_url, read_cloud_config, write_cloud_config
 
@@ -42,7 +41,8 @@ class ApiSession(requests.Session):
 
             self.__refresh_token = response.json()["refresh_token"]
             id_token = response.json()["cognito_client_id_token"]
-
+            user_sub_id = response.json()["user_sub_id"]
+            id_token_expires = response.json()["id_token_expires"]
         else:
             config = read_cloud_config()
 
@@ -55,14 +55,14 @@ class ApiSession(requests.Session):
             self.__configure_for_region(config.region)
             self.__refresh_token = config.refresh_token
             id_token = config.id_token
+            user_sub_id = config.user_sub_id
+            id_token_expires = config.id_token_expires
 
         assert id_token
 
-        decoded = jwt.decode(id_token, options={"verify_signature": False})
-        self.__expiry_time = decoded["exp"] - 60  # Refresh 1 minute before expiry
+        self.__user_sub_id = user_sub_id
+        self.__id_token_expires = id_token_expires - 60  # Refresh 1 minute before expiry
         self.headers["Authorization"] = f"Bearer {id_token}"
-
-        self.__sub = decoded["sub"]
         self.headers["X-Client-Version"] = __version__
 
     def __configure_for_region(self, region: str) -> None:
@@ -73,12 +73,12 @@ class ApiSession(requests.Session):
         logger.debug(f"Lambda url: {self.api_url}")
 
     def __ensure_valid_authorization_header(self) -> None:
-        if self.__expiry_time > time.time():
+        if self.__id_token_expires > time.time():
             return
 
         response = requests.post(
             self.__login_url,
-            json={"user_sub_id": self.__sub, "refresh_token": self.__refresh_token},
+            json={"user_sub_id": self.__user_sub_id, "refresh_token": self.__refresh_token},
             headers={"X-Client-Version": __version__},
         )
 
@@ -92,13 +92,14 @@ class ApiSession(requests.Session):
         #       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})
-        self.__expiry_time = decoded["exp"] - 60  # Refresh 1 minute before expiry
+        id_token_expires = response.json()["id_token_expires"]
+        self.__id_token_expires = id_token_expires - 60  # Refresh 1 minute before expiry
         self.headers["Authorization"] = f"Bearer {id_token}"
 
         if not self.non_interactive:
             config = read_cloud_config()
             config.id_token = id_token
+            config.id_token_expires = id_token_expires
             write_cloud_config(config)
 
     def request(self, method, url, *args, **kwargs) -> requests.Response:

locust_cloud/args.py

@@ -115,6 +115,11 @@ cloud_parser.add_argument(
     action="store_true",
     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 expire you will be prompted to perform another login.",
 )
+cloud_parser.add_argument(
+    "--logout",
+    action="store_true",
+    help="Removes the authentication credentials",
+)
 cloud_parser.add_argument(
     "--delete",
     action="store_true",
@@ -213,3 +218,23 @@ combined_cloud_parser.add_argument(
     choices=["DEBUG", "INFO", "WARNING", "ERROR"],
     default="INFO",
 )
+
+
+def add_locust_cloud_argparse(parser):
+    cloud_group = parser.add_argument_group(
+        "Locust Cloud",
+        """Launches a distributed Locust run on locust.cloud infrastructure.
+
+Example: locust --cloud -f my_locustfile.py --users 1000 ...""",
+    )
+
+    # This arguments is defined here because only makes sense when
+    # running from locust core
+    cloud_group.add_argument(
+        "--cloud",
+        action="store_true",
+        help="Run Locust in cloud mode.",
+    )
+
+    for action in cloud_parser._actions:
+        cloud_group._add_action(action)

locust_cloud/cloud.py

@@ -9,7 +9,7 @@ from locust_cloud.apisession import ApiSession
 from locust_cloud.args import combined_cloud_parser
 from locust_cloud.common import __version__
 from locust_cloud.input_events import input_listener
-from locust_cloud.web_login import web_login
+from locust_cloud.web_login import logout, web_login
 from locust_cloud.websocket import SessionMismatchError, Websocket, WebsocketTimeout
 
 logger = logging.getLogger(__name__)
@@ -39,7 +39,13 @@ def main():
     if options.login:
         try:
             web_login()
-        except KeyboardInterrupt:
+        except Exception:
+            pass
+        return
+    if options.logout:
+        try:
+            logout()
+        except Exception:
             pass
         return
 

locust_cloud/common.py

@@ -16,9 +16,11 @@ CLOUD_CONF_FILE = pathlib.Path(platformdirs.user_config_dir(appname="locust-clou
 @dataclass
 class CloudConfig:
     id_token: str | None = None
+    user_sub_id: str | None = None
     refresh_token: str | None = None
     refresh_token_expires: int = 0
     region: str | None = None
+    id_token_expires: int = 0
 
 
 def get_api_url(region):
@@ -38,3 +40,8 @@ def write_cloud_config(config: CloudConfig) -> None:
 
     with open(CLOUD_CONF_FILE, "w") as f:
         json.dump(config.__dict__, f)
+
+
+def delete_cloud_config() -> None:
+    if CLOUD_CONF_FILE.exists():
+        CLOUD_CONF_FILE.unlink()

locust_cloud/web_login.py

@@ -3,7 +3,7 @@ import time
 import webbrowser
 
 import requests
-from locust_cloud.common import VALID_REGIONS, CloudConfig, get_api_url, write_cloud_config
+from locust_cloud.common import VALID_REGIONS, CloudConfig, delete_cloud_config, get_api_url, write_cloud_config
 
 POLLING_FREQUENCY = 1
 
@@ -71,7 +71,13 @@ If the browser does not open or you wish to use a different device to authorize
     config = CloudConfig(
         id_token=data["id_token"],
         refresh_token=data["refresh_token"],
+        user_sub_id=data["user_sub_id"],
         refresh_token_expires=data["refresh_token_expires"],
+        id_token_expires=data["id_token_expires"],
         region=region,
     )
     write_cloud_config(config)
+
+
+def logout():
+    delete_cloud_config()

{locust_cloud-1.20.3.dist-info → locust_cloud-1.20.4.dist-info}/METADATA

@@ -1,16 +1,17 @@
 Metadata-Version: 2.4
 Name: locust-cloud
-Version: 1.20.3
+Version: 1.20.4
 Summary: Locust Cloud
 Project-URL: homepage, https://locust.cloud
 Project-URL: repository, https://github.com/locustcloud/locust-cloud
 License: MIT
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: locust>=2.36.1
+Requires-Dist: configargparse>=1.5.5
+Requires-Dist: gevent<25.0.0,>=24.10.1
 Requires-Dist: platformdirs<5.0.0,>=4.3.6
-Requires-Dist: pyjwt<3.0,>=2.0
 Requires-Dist: python-socketio[client]==5.13.0
+Requires-Dist: tomli>=1.1.0; python_version < '3.11'
 Description-Content-Type: text/markdown
 
 # Locust Cloud

locust_cloud-1.20.4.dist-info/RECORD

@@ -0,0 +1,18 @@
+locust_cloud/__init__.py,sha256=6z2hE5rUP9WJyYgr-7XC2GhIV-05m8XxjOsnb8ae1WY,56
+locust_cloud/apisession.py,sha256=gWqNADGL3H745fefm_tE1O1516OBfOhqZRlc-BgIudM,4224
+locust_cloud/args.py,sha256=warMJQvUg_liowJWHS0XF400rQQ0PUigPSutfwG9BLY,7661
+locust_cloud/cloud.py,sha256=LFO_Ykst0bOBYPqfxS6ll7abAvCCqyUpD3ql4nem9ys,6264
+locust_cloud/common.py,sha256=GVKkWcbbqd9n8oU-fHZRVZw3jGtuIVGSCLD2ZizeEo0,1160
+locust_cloud/input_events.py,sha256=MyxccgboHByICuK6VpQCCJhZQqTZAacNmkSpw-gxBEw,3420
+locust_cloud/web_login.py,sha256=Byu62DKyf8NSbPdP9I2Q3lJh_lZgxjn0OZcRgnMIlXg,2555
+locust_cloud/websocket.py,sha256=9Q7nTFuAwVhgW74DlJNcHTZXOQ1drsXi8hX9ciZhWlQ,8998
+locust_cloud/docs/.gitignore,sha256=ghNPcjYkjQXNS_eVmu2hQFhq6FIUliAD1O2CJhulS2o,10
+locust_cloud/docs/1-first-run.rst,sha256=dYYVxVb1RlWPNvhE98H1D_mzTJ1H2NWNn435nag3_jA,691
+locust_cloud/docs/2-examples.rst,sha256=RublnAK3jHRsEtbJAr1y2NEEoN0YAjkzP1np-7tJ3aE,9565
+locust_cloud/docs/locust-cloud.rst,sha256=NMfWGZdT0YPa0gPPuOcKplzBcqiEK7buPzXHOlzY5T0,98
+locust_cloud/docs/images/locust-cloud-screenshot.png,sha256=ag0IxBi-40VexC84MApol1GCgRCL2h-l8NQDTMaeTyE,477350
+locust_cloud-1.20.4.dist-info/METADATA,sha256=dEwbv0pE4L1vUKIfC9yoMFAGLuqafCiODj_HUW3-HlU,644
+locust_cloud-1.20.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+locust_cloud-1.20.4.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
+locust_cloud-1.20.4.dist-info/licenses/LICENSE,sha256=Ow6fY6ta4KIjdlWalmxGvRP8yLmetvkbkl-SdHMjPIs,1093
+locust_cloud-1.20.4.dist-info/RECORD,,

locust_cloud-1.20.3.dist-info/RECORD

@@ -1,18 +0,0 @@
-locust_cloud/__init__.py,sha256=9USe4yHuvpP5xrXm0dVsi08pu2l5TcjxQl6Pei8w3lk,713
-locust_cloud/apisession.py,sha256=kCr271_l0IeMGw0L563mOecqPJj4OD9h2J3vxCM5zYQ,4015
-locust_cloud/args.py,sha256=jBAKDNckEfP5v_2rKEMWWiWK_sjW6kX3sNgzgQkGk0Q,6974
-locust_cloud/cloud.py,sha256=esKhV4BXIkn2-1kPyUkW-o8k7GpevgaJz8CrA0S89VE,6149
-locust_cloud/common.py,sha256=cFrDVKpi9OEmH6giOuj9HoIUFSBArixNtNHzZIgDvPE,992
-locust_cloud/input_events.py,sha256=MyxccgboHByICuK6VpQCCJhZQqTZAacNmkSpw-gxBEw,3420
-locust_cloud/web_login.py,sha256=1j2AQoEM6XVSDtE1q0Ryrs4jFEx07r9IQfZCoFAQXJg,2400
-locust_cloud/websocket.py,sha256=9Q7nTFuAwVhgW74DlJNcHTZXOQ1drsXi8hX9ciZhWlQ,8998
-locust_cloud/docs/.gitignore,sha256=ghNPcjYkjQXNS_eVmu2hQFhq6FIUliAD1O2CJhulS2o,10
-locust_cloud/docs/1-first-run.rst,sha256=dYYVxVb1RlWPNvhE98H1D_mzTJ1H2NWNn435nag3_jA,691
-locust_cloud/docs/2-examples.rst,sha256=RublnAK3jHRsEtbJAr1y2NEEoN0YAjkzP1np-7tJ3aE,9565
-locust_cloud/docs/locust-cloud.rst,sha256=NMfWGZdT0YPa0gPPuOcKplzBcqiEK7buPzXHOlzY5T0,98
-locust_cloud/docs/images/locust-cloud-screenshot.png,sha256=ag0IxBi-40VexC84MApol1GCgRCL2h-l8NQDTMaeTyE,477350
-locust_cloud-1.20.3.dist-info/METADATA,sha256=ifRrqs3KTmgT6Rt230IXc_5Clb6xwRL7jJPV-lfQ5wk,576
-locust_cloud-1.20.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-locust_cloud-1.20.3.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
-locust_cloud-1.20.3.dist-info/licenses/LICENSE,sha256=Ow6fY6ta4KIjdlWalmxGvRP8yLmetvkbkl-SdHMjPIs,1093
-locust_cloud-1.20.3.dist-info/RECORD,,