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

locust_cloud/cloud.py

@@ -1,21 +1,29 @@
 import base64
 import gzip
+import importlib.metadata
+import json
 import logging
 import os
+import pathlib
 import sys
 import threading
+import time
 import tomllib
 import urllib.parse
+import webbrowser
 from argparse import Namespace
 from collections import OrderedDict
+from dataclasses import dataclass
 from typing import IO, Any
 
 import configargparse
+import jwt
+import platformdirs
 import requests
 import socketio
 import socketio.exceptions
-from locust_cloud import __version__
-from locust_cloud.credential_manager import CredentialError, CredentialManager
+
+__version__ = importlib.metadata.version("locust-cloud")
 
 
 class LocustTomlConfigParser(configargparse.TomlConfigParser):
@@ -49,7 +57,7 @@ parser = configargparse.ArgumentParser(
         "cloud.conf",
     ],
     auto_env_var_prefix="LOCUSTCLOUD_",
-    formatter_class=configargparse.RawDescriptionHelpFormatter,
+    formatter_class=configargparse.RawTextHelpFormatter,
     config_file_parser_class=configargparse.CompositeConfigParser(
         [
             LocustTomlConfigParser(["tool.locust"]),
@@ -108,36 +116,16 @@ advanced.add_argument(
     help="Optional requirements.txt file that contains your external libraries.",
 )
 advanced.add_argument(
-    "--region",
-    type=str,
-    default=os.environ.get("AWS_DEFAULT_REGION"),
-    help="Sets the AWS region to use for the deployed cluster, e.g. us-east-1. It defaults to use AWS_DEFAULT_REGION env var, like AWS tools.",
-)
-parser.add_argument(
-    "--aws-access-key-id",
-    type=str,
-    help=configargparse.SUPPRESS,
-    env_var="AWS_ACCESS_KEY_ID",
-    default=None,
-)
-parser.add_argument(
-    "--aws-secret-access-key",
-    type=str,
-    help=configargparse.SUPPRESS,
-    env_var="AWS_SECRET_ACCESS_KEY",
-    default=None,
-)
-parser.add_argument(
-    "--username",
-    type=str,
+    "--login",
+    action="store_true",
+    default=False,
     help=configargparse.SUPPRESS,
-    default=os.getenv("LOCUST_CLOUD_USERNAME", None),  # backwards compatitibility for dmdb
 )
-parser.add_argument(
-    "--password",
-    type=str,
-    help=configargparse.SUPPRESS,
-    default=os.getenv("LOCUST_CLOUD_PASSWORD", None),  # backwards compatitibility for dmdb
+advanced.add_argument(
+    "--non-interactive",
+    action="store_true",
+    default=False,
+    help="This can be set when, for example, running in a CI/CD environment to ensure no interactive steps while executing.\nRequires that LOCUSTCLOUD_USERNAME, LOCUSTCLOUD_PASSWORD and LOCUSTCLOUD_REGION environment variables are set.",
 )
 parser.add_argument(
     "--workers",
@@ -153,7 +141,7 @@ parser.add_argument(
 parser.add_argument(
     "--image-tag",
     type=str,
-    default="latest",
+    default=None,
     help=configargparse.SUPPRESS,  # overrides the locust-cloud docker image tag. for internal use
 )
 parser.add_argument(
@@ -169,6 +157,7 @@ parser.add_argument(
 )
 
 options, locust_options = parser.parse_known_args()
+
 options: Namespace
 locust_options: list
 
@@ -178,58 +167,426 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 # Restore log level for other libs. Yes, this can be done more nicely
-logging.getLogger("botocore").setLevel(logging.INFO)
-logging.getLogger("boto3").setLevel(logging.INFO)
 logging.getLogger("requests").setLevel(logging.INFO)
 logging.getLogger("urllib3").setLevel(logging.INFO)
 
+cloud_conf_file = pathlib.Path(platformdirs.user_config_dir(appname="locust-cloud")) / "config"
+valid_regions = ["us-east-1", "eu-north-1"]
 
-api_url = os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", f"https://api.{options.region}.locust.cloud/1")
 
+def get_api_url(region):
+    return os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", f"https://api.{region}.locust.cloud/1")
 
-def main() -> None:
-    if options.version:
-        print(f"locust-cloud version {__version__}")
-        sys.exit(0)
 
-    if not options.region:
-        logger.error(
-            "Setting a region is required to use Locust Cloud. Please ensure the AWS_DEFAULT_REGION env variable or the --region flag is set."
-        )
+@dataclass
+class CloudConfig:
+    id_token: str | None = None
+    refresh_token: str | None = None
+    refresh_token_expires: int = 0
+    region: str | None = None
+
+
+def read_cloud_config() -> CloudConfig:
+    if cloud_conf_file.exists():
+        with open(cloud_conf_file) as f:
+            return CloudConfig(**json.load(f))
+
+    return CloudConfig()
+
+
+def write_cloud_config(config: CloudConfig) -> None:
+    cloud_conf_file.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(cloud_conf_file, "w") as f:
+        json.dump(config.__dict__, f)
+
+
+def web_login() -> None:
+    print("Enter the number for the region to authenticate against")
+    print()
+    for i, valid_region in enumerate(valid_regions, start=1):
+        print(f"  {i}. {valid_region}")
+    print()
+    choice = input("> ")
+    try:
+        region_index = int(choice) - 1
+        assert 0 <= region_index < len(valid_regions)
+    except (ValueError, AssertionError):
+        print(f"Not a valid choice: '{choice}'")
         sys.exit(1)
-    if options.region:
-        os.environ["AWS_DEFAULT_REGION"] = options.region
 
-    if not ((options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)):
-        logger.error(
-            "Authentication is required to use Locust Cloud. Please ensure the LOCUSTCLOUD_USERNAME and LOCUSTCLOUD_PASSWORD environment variables are set."
-        )
+    region = valid_regions[region_index]
+
+    try:
+        response = requests.post(f"{get_api_url(region)}/cli-auth")
+        response.raise_for_status()
+        response_data = response.json()
+        authentication_url = response_data["authentication_url"]
+        result_url = response_data["result_url"]
+    except Exception as e:
+        print("Something went wrong trying to authorize the locust-cloud CLI:", str(e))
         sys.exit(1)
+
+    message = f"""
+Attempting to automatically open the SSO authorization page in your default browser.
+If the browser does not open or you wish to use a different device to authorize this request, open the following URL:
+
+{authentication_url}
+    """.strip()
+    print()
+    print(message)
+
+    webbrowser.open_new_tab(authentication_url)
+
+    while True:  # Should there be some kind of timeout?
+        response = requests.get(result_url)
+
+        if not response.ok:
+            print("Oh no!")
+            print(response.text)
+            sys.exit(1)
+
+        data = response.json()
+
+        if data["state"] == "pending":
+            time.sleep(1)
+            continue
+        elif data["state"] == "failed":
+            print(f"\nFailed to authorize CLI: {data['reason']}")
+            sys.exit(1)
+        elif data["state"] == "authorized":
+            print("\nAuthorization succeded")
+            break
+        else:
+            print("\nGot unexpected response when authorizing CLI")
+            sys.exit(1)
+
+    config = CloudConfig(
+        id_token=data["id_token"],
+        refresh_token=data["refresh_token"],
+        refresh_token_expires=data["refresh_token_expires"],
+        region=region,
+    )
+    write_cloud_config(config)
+
+
+class ApiSession(requests.Session):
+    def __init__(self) -> None:
+        super().__init__()
+
+        if options.non_interactive:
+            username = os.getenv("LOCUSTCLOUD_USERNAME")
+            password = os.getenv("LOCUSTCLOUD_PASSWORD")
+            region = os.getenv("LOCUSTCLOUD_REGION")
+
+            if not all([username, password, region]):
+                print(
+                    "Running with --non-interaction requires that LOCUSTCLOUD_USERNAME, LOCUSTCLOUD_PASSWORD and LOCUSTCLOUD_REGION environment variables are set."
+                )
+                sys.exit(1)
+
+            if region not in valid_regions:
+                print("Environment variable LOCUSTCLOUD_REGION needs to be set to one of", ", ".join(valid_regions))
+                sys.exit(1)
+
+            self.__configure_for_region(region)
+            response = requests.post(
+                self.__login_url,
+                json={"username": username, "password": password},
+                headers={"X-Client-Version": __version__},
+            )
+            if not response.ok:
+                print(f"Authentication failed: {response.text}")
+                sys.exit(1)
+
+            self.__refresh_token = response.json()["refresh_token"]
+            id_token = response.json()["cognito_client_id_token"]
+
+        else:
+            config = read_cloud_config()
+
+            if config.refresh_token_expires < time.time() + 24 * 60 * 60:
+                message = "You need to authenticate before proceeding. Please run:\n    locust-cloud --login"
+                print(message)
+                sys.exit(1)
+
+            assert config.region
+            self.__configure_for_region(config.region)
+            self.__refresh_token = config.refresh_token
+            id_token = config.id_token
+
+        assert id_token
+
+        decoded = jwt.decode(id_token, options={"verify_signature": False})
+        self.__expiry_time = decoded["exp"] - 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:
+        self.__region = region
+        self.api_url = get_api_url(region)
+        self.__login_url = f"{self.api_url}/auth/login"
+
+        logger.debug(f"Lambda url: {self.api_url}")
+
+    def __ensure_valid_authorization_header(self) -> None:
+        if self.__expiry_time > time.time():
+            return
+
+        logger.info(f"Authenticating ({self.__region}, v{__version__})")
+
+        response = requests.post(
+            self.__login_url,
+            json={"user_sub_id": self.__sub, "refresh_token": self.__refresh_token},
+            headers={"X-Client-Version": __version__},
+        )
+
+        if not response.ok:
+            logger.error(f"Authentication failed: {response.text}")
+            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.
+
+        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
+        self.headers["Authorization"] = f"Bearer {id_token}"
+
+        if not options.non_interactive:
+            config = read_cloud_config()
+            config.id_token = id_token
+            write_cloud_config(config)
+
+    def request(self, method, url, *args, **kwargs) -> requests.Response:
+        self.__ensure_valid_authorization_header()
+        return super().request(method, f"{self.api_url}{url}", *args, **kwargs)
+
+
+class SessionMismatchError(Exception):
+    pass
+
+
+class WebsocketTimeout(Exception):
+    pass
+
+
+class Websocket:
+    def __init__(self) -> None:
+        """
+        This class was created to encapsulate all the logic involved in the websocket implementation.
+        The behaviour of the socketio client once a connection has been established
+        is to try to reconnect forever if the connection is lost.
+        The way this can be canceled is by setting the _reconnect_abort (threading.Event) on the client
+        in which case it will simply proceed with shutting down without giving any indication of an error.
+        This class handles timeouts for connection attempts as well as some logic around when the
+        socket can be shut down. See descriptions on the methods for further details.
+        """
+        self.__shutdown_allowed = threading.Event()
+        self.__timeout_on_disconnect = True
+        self.initial_connect_timeout = 120
+        self.reconnect_timeout = 10
+        self.wait_timeout = 0
+        self.exception: None | Exception = None
+
+        self.sio = socketio.Client(handle_sigint=False)
+        self.sio._reconnect_abort = threading.Event()
+        # The _reconnect_abort value on the socketio client will be populated with a newly created threading.Event if it's not already set.
+        # There is no way to set this by passing it in the constructor.
+        # This event is the only way to interupt the retry logic when the connection is attempted.
+
+        self.sio.on("connect", self.__on_connect)
+        self.sio.on("disconnect", self.__on_disconnect)
+        self.sio.on("connect_error", self.__on_connect_error)
+        self.sio.on("events", self.__on_events)
+
+        self.__processed_events: set[int] = set()
+
+    def __set_connection_timeout(self, timeout) -> None:
+        """
+        Start a threading.Timer that will set the threading.Event on the socketio client
+        that aborts any further attempts to reconnect, sets an exception on the websocket
+        that will be raised from the wait method and the threading.Event __shutdown_allowed
+        on the websocket that tells the wait method that it should stop blocking.
+        """
+
+        def _timeout():
+            logger.debug(f"Websocket connection timed out after {timeout} seconds")
+            self.sio._reconnect_abort.set()
+            self.exception = WebsocketTimeout("Timed out connecting to locust master")
+            self.__shutdown_allowed.set()
+
+        self.__connect_timeout_timer = threading.Timer(timeout, _timeout)
+        self.__connect_timeout_timer.daemon = True
+        logger.debug(f"Setting websocket connection timeout to {timeout} seconds")
+        self.__connect_timeout_timer.start()
+
+    def connect(self, url, *, auth) -> None:
+        """
+        Send along retry=True when initiating the socketio client connection
+        to make it use it's builtin logic for retrying failed connections that
+        is usually used for reconnections. This will retry forever.
+        When connecting start a timer to trigger disabling the retry logic and
+        raise a WebsocketTimeout exception.
+        """
+        ws_connection_info = urllib.parse.urlparse(url)
+        self.__set_connection_timeout(self.initial_connect_timeout)
+        try:
+            self.sio.connect(
+                f"{ws_connection_info.scheme}://{ws_connection_info.netloc}",
+                auth=auth,
+                retry=True,
+                **{"socketio_path": ws_connection_info.path} if ws_connection_info.path else {},
+            )
+        except socketio.exceptions.ConnectionError:
+            if self.exception:
+                raise self.exception
+
+            raise
+
+    def shutdown(self) -> None:
+        """
+        When shutting down the socketio client a disconnect event will fire.
+        Before doing so disable the behaviour of starting a threading.Timer
+        to handle timeouts on attempts to reconnect since no further such attempts
+        will be made.
+        If such a timer is already running, cancel it since the client is being shutdown.
+        """
+        self.__timeout_on_disconnect = False
+        if hasattr(self, "__connect_timeout_timer"):
+            self.__connect_timeout_timer.cancel()
+        self.sio.shutdown()
+
+    def wait(self, timeout=False) -> bool:
+        """
+        Block until the threading.Event __shutdown_allowed is set, with a timeout if indicated.
+        If an exception has been set on the websocket (from a connection timeout timer or the
+        __on_connect_error method), raise it.
+        """
+        timeout = self.wait_timeout if timeout else None
+        logger.debug(f"Waiting for shutdown for {str(timeout)+'s' if timeout else 'ever'}")
+        res = self.__shutdown_allowed.wait(timeout)
+        if self.exception:
+            raise self.exception
+        return res
+
+    def __on_connect(self) -> None:
+        """
+        This gets events whenever a connection is successfully established.
+        When this happens, cancel the running threading.Timer that would
+        abort reconnect attempts and raise a WebsocketTimeout exception.
+        The wait_timeout is originally set to zero when creating the websocket
+        but once a connection has been established this is raised to ensure
+        that the server is given the chance to send all the logs and an
+        official shutdown event.
+        """
+        self.__connect_timeout_timer.cancel()
+        self.wait_timeout = 90
+        logger.debug("Websocket connected")
+
+    def __on_disconnect(self) -> None:
+        """
+        This gets events whenever a connection is lost.
+        The socketio client will try to reconnect forever so,
+        unless the behaviour has been disabled, a threading.Timer
+        is started that will abort reconnect attempts and raise a
+        WebsocketTimeout exception.
+        """
+        if self.__timeout_on_disconnect:
+            self.__set_connection_timeout(self.reconnect_timeout)
+        logger.debug("Websocket disconnected")
+
+    def __on_events(self, data):
+        """
+        This gets events explicitly sent by the websocket server.
+        This will either be messages to print on stdout/stderr or
+        an indication that the CLI can shut down in which case the
+        threading.Event __shutdown_allowed gets set on the websocket
+        that tells the wait method that it should stop blocking.
+        """
+        shutdown = False
+        shutdown_message = ""
+
+        if data["id"] in self.__processed_events:
+            logger.debug(f"Got duplicate data on websocket, id {data['id']}")
+            return
+
+        self.__processed_events.add(data["id"])
+
+        for event in data["events"]:
+            type = event["type"]
+
+            if type == "shutdown":
+                shutdown = True
+                shutdown_message = event["message"]
+            elif type == "stdout":
+                sys.stdout.write(event["message"])
+            elif type == "stderr":
+                sys.stderr.write(event["message"])
+            else:
+                raise Exception("Unexpected event type")
+
+        if shutdown:
+            logger.debug("Got shutdown from locust master")
+            if shutdown_message:
+                print(shutdown_message)
+
+            self.__shutdown_allowed.set()
+
+    def __on_connect_error(self, data) -> None:
+        """
+        This gets events whenever there's an error during connection attempts.
+        The specific case that is handled below is triggered when the connection
+        is made with the auth parameter not matching the session ID on the server.
+        If this error occurs it's because the connection is attempted towards an
+        instance of locust not started by this CLI.
+
+        In that case:
+        Cancel the running threading.Timer that would abort reconnect attempts
+        and raise a WebsocketTimeout exception.
+        Set an exception on the websocket that will be raised from the wait method.
+        Cancel further reconnect attempts.
+        Set the threading.Event __shutdown_allowed on the websocket that tells the
+        wait method that it should stop blocking.
+        """
+        # Do nothing if it's not the specific case we know how to deal with
+        if not (isinstance(data, dict) and data.get("message") == "Session mismatch"):
+            return
+
+        self.__connect_timeout_timer.cancel()
+        self.exception = SessionMismatchError(
+            "The session from this run of locust-cloud did not match the one on the server"
+        )
+        self.sio._reconnect_abort.set()
+        self.__shutdown_allowed.set()
+
+
+def main() -> None:
+    if options.version:
+        print(f"locust-cloud version {__version__}")
+        sys.exit(0)
     if not options.locustfile:
         logger.error("A locustfile is required to run a test.")
         sys.exit(1)
 
-    try:
-        logger.info(f"Authenticating ({options.region}, v{__version__})")
-        logger.debug(f"Lambda url: {api_url}")
-        credential_manager = CredentialManager(
-            lambda_url=api_url,
-            access_key=options.aws_access_key_id,
-            secret_key=options.aws_secret_access_key,
-            username=options.username,
-            password=options.password,
-        )
+    if options.login:
+        try:
+            web_login()
+        except KeyboardInterrupt:
+            pass
+        sys.exit()
 
-        credentials = credential_manager.get_current_credentials()
-        cognito_client_id_token = credentials["cognito_client_id_token"]
-        aws_access_key_id = credentials.get("access_key")
-        aws_secret_access_key = credentials.get("secret_key")
-        aws_session_token = credentials.get("token", "")
+    session = ApiSession()
+    websocket = Websocket()
 
-        if options.delete:
-            delete(credential_manager)
-            return
+    if options.delete:
+        delete(session)
+        sys.exit()
 
+    try:
         try:
             with open(options.locustfile, "rb") as f:
                 locustfile_data = base64.b64encode(gzip.compress(f.read())).decode()
@@ -249,54 +606,46 @@ def main() -> None:
 
         logger.info("Deploying load generators")
         locust_env_variables = [
-            {"name": env_variable, "value": str(os.environ[env_variable])}
+            {"name": env_variable, "value": os.environ[env_variable]}
             for env_variable in os.environ
             if env_variable.startswith("LOCUST_")
-            and not env_variable
-            in [
+            and env_variable
+            not in [
                 "LOCUST_LOCUSTFILE",
                 "LOCUST_USERS",
                 "LOCUST_WEB_HOST_DISPLAY_NAME",
                 "LOCUST_SKIP_MONKEY_PATCH",
             ]
-            and os.environ[env_variable]
         ]
-        deploy_endpoint = f"{api_url}/deploy"
         payload = {
             "locust_args": [
                 {"name": "LOCUST_USERS", "value": str(options.users)},
                 {"name": "LOCUST_FLAGS", "value": " ".join(locust_options)},
-                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": api_url},
+                {"name": "LOCUSTCLOUD_DEPLOYER_URL", "value": session.api_url},
                 {"name": "LOCUSTCLOUD_PROFILE", "value": options.profile},
                 *locust_env_variables,
             ],
             "locustfile": {"filename": options.locustfile, "data": locustfile_data},
             "user_count": options.users,
-            "image_tag": options.image_tag,
             "mock_server": options.mock_server,
         }
+
+        if options.image_tag is not None:
+            payload["image_tag"] = options.image_tag
+
         if options.workers is not None:
             payload["worker_count"] = options.workers
+
         if options.requirements:
             payload["requirements"] = {"filename": options.requirements, "data": requirements_data}
-        headers = {
-            "Authorization": f"Bearer {cognito_client_id_token}",
-            "Content-Type": "application/json",
-            "AWS_ACCESS_KEY_ID": aws_access_key_id,
-            "AWS_SECRET_ACCESS_KEY": aws_secret_access_key,
-            "AWS_SESSION_TOKEN": aws_session_token,
-            "X-Client-Version": __version__,
-        }
+
         try:
-            # logger.info(payload) # might be useful when debugging sometimes
-            response = requests.post(deploy_endpoint, json=payload, headers=headers)
+            response = session.post("/deploy", json=payload)
         except requests.exceptions.RequestException as e:
             logger.error(f"Failed to deploy the load generators: {e}")
             sys.exit(1)
 
-        if response.status_code == 200:
-            log_ws_url = response.json()["log_ws_url"]
-        else:
+        if response.status_code != 200:
             try:
                 logger.error(f"{response.json()['Message']} (HTTP {response.status_code}/{response.reason})")
             except Exception:
@@ -304,98 +653,50 @@ def main() -> None:
                     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}")
-        sys.exit(1)
-    except KeyboardInterrupt:
-        logger.debug("Interrupted by user")
-        sys.exit(0)
-
-    logger.debug("Load generators deployed successfully!")
-    logger.info("Waiting for pods to be ready...")
-
-    shutdown_allowed = threading.Event()
-    shutdown_allowed.set()
-    reconnect_aborted = threading.Event()
-    connect_timeout = threading.Timer(2 * 60, reconnect_aborted.set)
-
-    try:
-        ws_connection_info = urllib.parse.urlparse(log_ws_url)
-        sio = socketio.Client(handle_sigint=False)
-
-        @sio.event
-        def connect():
-            shutdown_allowed.clear()
-            connect_timeout.cancel()
-            logger.debug("Websocket connection established, switching to Locust logs")
-
-        @sio.event
-        def disconnect():
-            logger.debug("Websocket disconnected")
 
-        @sio.event
-        def stderr(message):
-            sys.stderr.write(message)
+        log_ws_url = response.json()["log_ws_url"]
+        session_id = response.json()["session_id"]
+        logger.debug(f"Session ID is {session_id}")
 
-        @sio.event
-        def stdout(message):
-            sys.stdout.write(message)
-
-        @sio.event
-        def shutdown(message):
-            logger.debug("Got shutdown from locust master")
-            if message:
-                print(message)
-
-            shutdown_allowed.set()
-
-        # The _reconnect_abort value on the socketio client will be populated with a newly created threading.Event if it's not already set.
-        # There is no way to set this by passing it in the constructor.
-        # This event is the only way to interupt the retry logic when the connection is attempted.
-        sio._reconnect_abort = reconnect_aborted
-        connect_timeout.start()
-        sio.connect(
-            f"{ws_connection_info.scheme}://{ws_connection_info.netloc}",
-            socketio_path=ws_connection_info.path,
-            retry=True,
+        logger.info("Waiting for pods to be ready...")
+        websocket.connect(
+            log_ws_url,
+            auth=session_id,
         )
-        logger.debug("Waiting for shutdown")
-        shutdown_allowed.wait()
+        websocket.wait()
 
     except KeyboardInterrupt:
         logger.debug("Interrupted by user")
-        delete(credential_manager)
-        shutdown_allowed.wait(timeout=90)
+        delete(session)
+        try:
+            websocket.wait(timeout=True)
+        except (WebsocketTimeout, SessionMismatchError) as e:
+            logger.error(str(e))
+            sys.exit(1)
+    except WebsocketTimeout as e:
+        logger.error(str(e))
+        delete(session)
+        sys.exit(1)
+    except SessionMismatchError as e:
+        # In this case we do not trigger the teardown since the running instance is not ours
+        logger.error(str(e))
+        sys.exit(1)
     except Exception as e:
         logger.exception(e)
-        delete(credential_manager)
+        delete(session)
         sys.exit(1)
     else:
-        delete(credential_manager)
+        delete(session)
     finally:
-        sio.shutdown()
+        logger.debug("Shutting down websocket")
+        websocket.shutdown()
 
 
-def delete(credential_manager):
+def delete(session):
     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', '')}",
-            "X-Client-Version": __version__,
-        }
-
-        token = refreshed_credentials.get("token")
-        if token:
-            headers["AWS_SESSION_TOKEN"] = token
-
-        response = requests.delete(
-            f"{api_url}/teardown",
-            headers=headers,
+        response = session.delete(
+            "/teardown",
         )
 
         if response.status_code == 200:
@@ -407,7 +708,7 @@ def delete(credential_manager):
     except Exception as e:
         logger.error(f"Could not automatically tear down Locust Cloud: {e.__class__.__name__}:{e}")
 
-    logger.info("Done! ✨")
+    logger.info("Done! ✨")  # FIXME: Should probably not say it's done since at this point it could still be running
 
 
 if __name__ == "__main__":