windmill/integration_tests/test/wmill_integration_test_utils.py

import datetime
import time
import httpx
import json
import os



class WindmillClient:
    _url: str
    _token: str
    _workspace: str

    _client: httpx.Client

    def __init__(self):
        self._workspace = "integration-tests"
        self._url = "http://localhost:8000"
        self._token = self._login()

        self._client = self._init_client()
        print("New client for Windmill version {}".format(self.get_version()))
        self._create_workspace()
        self._set_license_key()

    def __del__(self):
        self._logout()
        self._client.close()

    def _login(self) -> str:
        with httpx.Client(base_url=self._url) as unauth_client:
            response = unauth_client.post(
                "/api/auth/login",
                json={
                    "email": "admin@windmill.dev",
                    "password": "changeme",
                },
            )
            if response.status_code // 100 != 2:
                raise Exception(response.content.decode())
            return response.content.decode()

    def _logout(self) -> None:
        response = self._client.post(
            "/api/auth/logout",
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())

    def _init_client(self):
        token = self._token
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}",
        }
        return httpx.Client(
            base_url=self._url,
            headers=headers,
            timeout=60.0,  # Go/Rust compilation can take 10+ seconds on first run
        )

    def _set_license_key(self):
        license_key = os.environ.get("LICENSE_KEY", "").strip()
        print(
            "Setting license key to {}...{}".format(license_key[:15], license_key[-15:])
        )
        response = self._client.post(
            "/api/settings/global/license_key",
            json={
                "value": license_key,
            },
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())

    def _create_workspace(self):
        print("Creating workspace for integration tests")
        exists = self._client.post(
            "/api/workspaces/exists",
            json={
                "id": self._workspace,
            },
        )
        if exists.status_code // 100 == 2 and exists.content.decode() == "true":
            print("Workspace already exists, not creating it")
            return
        response = self._client.post(
            "/api/workspaces/create",
            json={
                "id": self._workspace,
                "name": self._workspace
            },
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.content.decode()

    def set_npm_config_registry(self, registry_url: str):
        response = self._client.post(
            "/api/settings/global/npm_config_registry",
            json={
                "value": registry_url,
            },
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())

    def run_sync(self, path: str, args: dict, type: str = "p"):
        print(f"Running {path} with args {args}")
        response = self._client.post(
            f"/api/w/{self._workspace}/jobs/run_wait_result/{type}/{path}",
            json=args,
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.json()

    def create_script(self, path: str, content: str, language: str, tag: str = None):
        print(f"Creating script {path}")

        payload = {
            "path": path,
            "content": content,
            "description": "",
            "summary": "",
            "language": language,
        }

        if tag is not None:
            payload["tag"] = tag

        response = self._client.post(
            f"/api/w/{self._workspace}/scripts/create",
            json=payload,
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        script_hash = response.content.decode()
        print(f"Script hash for path {path} is {script_hash}")
        time_now = datetime.datetime.now(datetime.timezone.utc)
        while datetime.datetime.now(
            datetime.timezone.utc
        ) - time_now < datetime.timedelta(seconds=60):
            response = self._client.get(
                f"/api/w/{self._workspace}/scripts/deployment_status/h/{script_hash}"
            )
            if response.status_code // 100 != 2:
                raise Exception(response.content.decode())
            elif response.json()["lock"] != None:
                # deployment successful -> return
                return
            elif response.json()["lock_error_logs"] != None:
                raise Exception(response.json()["lock_error_logs"])
            print(f"Waiting for script {path} with hash {script_hash} to be deployed")
            time.sleep(1)
        raise Exception(f"Script deployment failed for {path}")

    def delete_script(self, path: str):
        print(f"Deleting script {path}")
        response = self._client.post(
            f"/api/w/{self._workspace}/scripts/delete/p/{path}",
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.content.decode()

    def create_flow(self, path: str, flow_value_json: str):
        print(f"Creating flow {path}")
        parsed_flow = json.loads(flow_value_json)
        if "path" not in parsed_flow:
            parsed_flow["path"] = path
        response = self._client.post(
            f"/api/w/{self._workspace}/flows/create",
            json=parsed_flow,
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.content.decode()

    def delete_flow(self, path: str):
        print(f"Deleting flow {path}")
        response = self._client.delete(
            f"/api/w/{self._workspace}/flows/delete/{path}",
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.content.decode()

    def create_schedule(
        self,
        path: str,
        runnable_path: str,
        type: str = "script",
        schedule: str = "*/5 * * * * *",
        args: dict = {},
    ):
        print(f"Creating schedule {path}")
        response = self._client.post(
            f"/api/w/{self._workspace}/schedules/create",
            json={
                "path": path,
                "schedule": schedule,
                "timezone": "Europe/Paris",
                "script_path": runnable_path,
                "is_flow": type == "flow",
                "args": args,
                "enabled": True,
            },
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.content.decode()

    def delete_schedule(self, path: str):
        print(f"Deleting schedule {path}")
        response = self._client.delete(
            f"/api/w/{self._workspace}/schedules/delete/{path}",
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.content.decode()

    def create_variable(
        self,
        path: str,
        value: str,
    ):
        print(f"Creating variable {path} with value '{value}'")
        response = self._client.post(
            f"/api/w/{self._workspace}/variables/create",
            json={
                "path": path,
                "value": value,
                "description": "",
                "is_secret": False,
            },
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.content.decode()

    def delete_variable(self, path: str):
        print(f"Deleting variable {path}")
        response = self._client.delete(
            f"/api/w/{self._workspace}/variables/delete/{path}",
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.content.decode()

    def get_latest_job_runs(self, path: str):
        response = self._client.get(
            f"/api/w/{self._workspace}/jobs/list?script_path_exact={path}"
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())
        return response.json()

    def get_version(self):
        response = self._client.get("/api/version")
        return response.content.decode()

    def get_global_custom_tags(self):
        """
        Get the current list of global custom tags.

        Returns:
            list: List of custom tags or empty list if not set or an error occurred.
        """
        try:
            response = self._client.get("/api/settings/global/custom_tags")
            if response.status_code // 100 == 2:
                tags = response.json()
                return tags if tags is not None else []
            else:
                print(f"Error retrieving global custom tags: Status {response.status_code}, Response: {response.content.decode()}")
                return []
        except Exception as e:
            print(f"Exception when retrieving global custom tags: {e}")
            return []

    def add_global_custom_tag(self, tag):
        """
        Add a tag to the global custom tags if it's not already present.

        Args:
            tag (str): The tag to add to global custom tags.

        Returns:
            bool: True if the tag was added or already exists, False if there was an error.
        """
        try:
            current_tags = self.get_global_custom_tags()

            if tag in current_tags:
                print(f"Tag '{tag}' already exists in global custom tags")
                return True

            new_tags = current_tags + [tag]
            print(f"Adding '{tag}' to global custom tags: {new_tags}")

            response = self._client.post(
                "/api/settings/global/custom_tags",
                json={
                    "value": new_tags,
                },
            )

            if response.status_code // 100 == 2:
                print(f"Successfully added '{tag}' to global custom tags")
                return True
            else:
                print(f"Error adding tag to global custom tags: Status {response.status_code}, Response: {response.content.decode()}")
                return False
        except Exception as e:
            print(f"Exception when adding global custom tag: {e}")
            return False

    def remove_global_custom_tag(self, tag):
        """
        Remove a tag from the global custom tags if it exists.

        Args:
            tag (str): The tag to remove from global custom tags.

        Returns:
            bool: True if the tag was removed or didn't exist, False if there was an error.
        """
        try:
            current_tags = self.get_global_custom_tags()

            if tag not in current_tags:
                print(f"Tag '{tag}' doesn't exist in global custom tags")
                return True

            new_tags = [t for t in current_tags if t != tag]
            print(f"Removing '{tag}' from global custom tags: {new_tags}")

            response = self._client.post(
                "/api/settings/global/custom_tags",
                json={
                    "value": new_tags,
                },
            )

            if response.status_code // 100 == 2:
                print(f"Successfully removed '{tag}' from global custom tags")
                return True
            else:
                print(f"Error removing tag from global custom tags: Status {response.status_code}, Response: {response.content.decode()}")
                return False
        except Exception as e:
            print(f"Exception when removing global custom tag: {e}")
            return False

    def get_workers_list(self, ping_since=60, page=0, per_page=100):
        """
        Get a list of workers currently connected to the Windmill server.

        Args:
            ping_since (int): Only include workers that have pinged in the last N seconds. Default is 60.
            page (int): Page number for pagination. Default is 0.
            per_page (int): Number of results per page. Default is 100.

        Returns:
            list: List of worker objects or empty list if no workers found or an error occurred.
        """
        try:
            params = {
                "page": page,
                "per_page": per_page,
                "ping_since": ping_since
            }

            response = self._client.get(
                "/api/workers/list",
                params=params
            )
            if response.status_code // 100 == 2:
                return response.json()
            else:
                print(f"Error retrieving workers list: Status {response.status_code}, Response: {response.content.decode()}")
                return []
        except Exception as e:
            print(f"Exception when retrieving workers list: {e}")
            return []

    def create_agent_token(self, worker_group="agent", tags=None, exp=None):
        """
        Create an agent JWT token using superadmin privilege.

        Args:
            worker_group (str): The worker group for the agent, defaults to "agent"
            tags (list): Tags for the agent, defaults to ["agent"]
            exp (int): Expiration timestamp, defaults to a timestamp about 1 year in the future

        Returns:
            str: The JWT token for the agent
        """
        if tags is None:
            tags = ["agent"]

        if exp is None:
            exp = int(time.time()) + 31536000  # 60*60*24*365 = 1 year

        print(f"Creating agent token for worker_group={worker_group}, tags={tags}")
        response = self._client.post(
            "/api/agent_workers/create_agent_token",
            json={
                "worker_group": worker_group,
                "tags": tags,
                "exp": exp
            },
        )
        if response.status_code // 100 != 2:
            raise Exception(response.content.decode())

        token = response.content.decode().strip('"')
        print(f"Created agent token: {token}")
        return token