From 526ca8a80c579a9dc5d802138d35179aaf9f2658 Mon Sep 17 00:00:00 2001 From: Eirini Koutsaniti Date: Thu, 17 Aug 2023 16:31:42 +0200 Subject: [PATCH] Add asynchronous client for pyfirecrest (#65) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ✨ Add typing * make upload/download more permissive * apply black * Add types in to documentation * Update types * fix formatting * Add undoc-members in types reference * First implementation of AsyncFirecrest client * Add httpx in the dependencies * Add some basic example for the async client * Set polling rate in example * Fix import * Fix external transfer * Add example for external transfers * Add async external objects in reference * Expose async external objects * Fix issue with mixed responses in case of errors * Get the access token before stalling requests * Add documentation for async client + refactor * Decouple async and sync clients * Split ExternalStorage objects to another file * Fix typo * Fix typo in filename * Refactor status tests * Add auth handler * Refactor compute tests * Refactor extra tests * Refactor reservation tests * Refactor storage tests * Refactor utilities tests * Duplicate compute testing for async version * Add compute tests * Add extra tests for async client * Add reservation async tests + small fixes * Add status tests for async client + small fixes * Add utilities tests + fix async whoami * Add storage tests for async version * Remove unused imports and fix annotation * Small fixes * Fix type errors * Small fix in docs * Remove old installation instructions --- async-example.md | 89 ++ docs/source/index.rst | 5 +- docs/source/reference_async.rst | 29 + .../{reference.rst => reference_basic.rst} | 4 +- .../{cli_reference.rst => reference_cli.rst} | 0 docs/source/reference_index.rst | 10 + docs/source/tutorial_async.rst | 124 ++ .../{tutorial.rst => tutorial_basic.rst} | 101 +- docs/source/tutorial_cli.rst | 40 + docs/source/tutorial_errors.rst | 30 + docs/source/tutorial_index.rst | 12 + docs/source/tutorial_logging.rst | 26 + firecrest/AsyncClient.py | 1397 +++++++++++++++++ firecrest/AsyncExternalStorage.py | 261 +++ firecrest/BasicClient.py | 231 +-- firecrest/ExternalStorage.py | 260 +++ firecrest/__init__.py | 12 +- pyproject.toml | 4 +- ..._authoriation.py => test_authorisation.py} | 56 +- tests/test_compute.py | 273 ++-- tests/test_compute_async.py | 405 +++++ tests/test_extras.py | 96 +- tests/test_extras_async.py | 167 ++ tests/test_reservation.py | 106 +- tests/test_reservation_async.py | 145 ++ tests/test_status.py | 188 +-- tests/test_status_async.py | 169 ++ tests/test_storage.py | 218 +-- tests/test_storage_async.py | 288 ++++ tests/test_utilities.py | 749 +++++---- tests/test_utilities_async.py | 653 ++++++++ 31 files changed, 5094 insertions(+), 1054 deletions(-) create mode 100644 async-example.md create mode 100644 docs/source/reference_async.rst rename docs/source/{reference.rst => reference_basic.rst} (96%) rename docs/source/{cli_reference.rst => reference_cli.rst} (100%) create mode 100644 docs/source/reference_index.rst create mode 100644 docs/source/tutorial_async.rst rename docs/source/{tutorial.rst => tutorial_basic.rst} (73%) create mode 100644 docs/source/tutorial_cli.rst create mode 100644 docs/source/tutorial_errors.rst create mode 100644 docs/source/tutorial_index.rst create mode 100644 docs/source/tutorial_logging.rst create mode 100644 firecrest/AsyncClient.py create mode 100644 firecrest/AsyncExternalStorage.py create mode 100644 firecrest/ExternalStorage.py rename tests/{test_authoriation.py => test_authorisation.py} (60%) create mode 100644 tests/test_compute_async.py create mode 100644 tests/test_extras_async.py create mode 100644 tests/test_reservation_async.py create mode 100644 tests/test_status_async.py create mode 100644 tests/test_storage_async.py create mode 100644 tests/test_utilities_async.py diff --git a/async-example.md b/async-example.md new file mode 100644 index 0000000..0c2ee92 --- /dev/null +++ b/async-example.md @@ -0,0 +1,89 @@ +# Examples for asyncio with pyfirecrest + +### Simple asynchronous workflow with the new client + +Here is an example of how to use the `AsyncFirecrest` client with asyncio. + +```python +import firecrest +import asyncio +import logging + + +# Setup variables before running the script +client_id = "" +client_secret = "" +token_uri = "" +firecrest_url = "" + +machine = "" +local_script_path = "" + +# Ignore this part, it is simply setup for logging +logger = logging.getLogger("simple_example") +logger.setLevel(logging.DEBUG) +ch = logging.StreamHandler() +ch.setLevel(logging.DEBUG) +formatter = logging.Formatter("%(asctime)s - %(message)s", datefmt="%H:%M:%S") +ch.setFormatter(formatter) +logger.addHandler(ch) + +async def workflow(client, i): + logger.info(f"{i}: Starting workflow") + job = await client.submit(machine, local_script_path) + logger.info(f"{i}: Submitted job with jobid: {job['jobid']}") + while True: + poll_res = await client.poll_active(machine, [job["jobid"]]) + if len(poll_res) < 1: + logger.info(f"{i}: Job {job['jobid']} is no longer active") + break + + logger.info(f"{i}: Job {job['jobid']} status: {poll_res[0]['state']}") + await asyncio.sleep(30) + + output = await client.view(machine, job["job_file_out"]) + logger.info(f"{i}: job output: {output}") + + +async def main(): + auth = firecrest.ClientCredentialsAuth(client_id, client_secret, token_uri) + client = firecrest.AsyncFirecrest(firecrest_url, authorization=auth) + + # Set up the desired polling rate for each microservice. The float number + # represents the number of seconds between consecutive requests in each + # microservice. Default is 5 seconds for now. + client.time_between_calls = { + "compute": 5, + "reservations": 5, + "status": 5, + "storage": 5, + "tasks": 5, + "utilities": 5, + } + + workflows = [workflow(client, i) for i in range(5)] + await asyncio.gather(*workflows) + + +asyncio.run(main()) + +``` + + +### External transfers with `AsyncFirecrest` + +The uploads and downloads work as before but you have to keep in mind which methods are coroutines. + +```python +# Download +down_obj = await client.external_download("cluster", "/remote/path/to/the/file") +status = await down_obj.status +print(status) +await down_obj.finish_download("my_local_file") + +# Upload +up_obj = await client.external_upload("cluster", "/path/to/local/file", "/remote/path/to/filesystem") +await up_obj.finish_upload() +status = await up_obj.status +print(status) +``` diff --git a/docs/source/index.rst b/docs/source/index.rst index e3782e3..e16f84f 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -17,9 +17,8 @@ You can also clone it from `Github `__ :caption: Contents: authorization - tutorial - reference - cli_reference + tutorial_index + reference_index Contact ======= diff --git a/docs/source/reference_async.rst b/docs/source/reference_async.rst new file mode 100644 index 0000000..79679f0 --- /dev/null +++ b/docs/source/reference_async.rst @@ -0,0 +1,29 @@ +Asynchronous FirecREST objects +============================== + +The library also provides an asynchronous API for the client: + +The ``AsyncFirecrest`` class +**************************** +.. autoclass:: firecrest.AsyncFirecrest + :members: + :undoc-members: + :show-inheritance: + + +The ``AsyncExternalDownload`` class +*********************************** +.. autoclass:: firecrest.AsyncExternalDownload + :inherited-members: + :members: + :undoc-members: + :show-inheritance: + + +The ``AsyncExternalUpload`` class +********************************* +.. autoclass:: firecrest.AsyncExternalUpload + :inherited-members: + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/reference.rst b/docs/source/reference_basic.rst similarity index 96% rename from docs/source/reference.rst rename to docs/source/reference_basic.rst index 78efa41..8fee278 100644 --- a/docs/source/reference.rst +++ b/docs/source/reference_basic.rst @@ -1,5 +1,5 @@ -Reference -========= +The basic client +================ The wrapper includes the ``Firecrest`` class, which is in practice a very basic client. Together with the authorisation class it takes care of the token and makes the appropriate calls for each action. diff --git a/docs/source/cli_reference.rst b/docs/source/reference_cli.rst similarity index 100% rename from docs/source/cli_reference.rst rename to docs/source/reference_cli.rst diff --git a/docs/source/reference_index.rst b/docs/source/reference_index.rst new file mode 100644 index 0000000..f81cd1b --- /dev/null +++ b/docs/source/reference_index.rst @@ -0,0 +1,10 @@ +Reference +========= + +.. toctree:: + :maxdepth: 2 + :caption: Contents: + + reference_basic + reference_async + reference_cli diff --git a/docs/source/tutorial_async.rst b/docs/source/tutorial_async.rst new file mode 100644 index 0000000..41557cf --- /dev/null +++ b/docs/source/tutorial_async.rst @@ -0,0 +1,124 @@ +How to use the asynchronous API [experimental] +============================================== + +In this tutorial, we will explore the asynchronous API of the pyFirecREST library. +Asynchronous programming is a powerful technique that allows you to write more efficient and responsive code by handling concurrent tasks without blocking the main execution flow. +This capability is particularly valuable when dealing with time-consuming operations such as network requests, I/O operations, or interactions with external services. + +In order to take advantage of the asynchronous client you may need to make many changes in your existing code, so the effort is worth it when you develop a code from the start or if you need to make a large number of requests. +You could submit hundreds or thousands of jobs, set a reasonable rate and pyFirecREST will handle it in the background without going over the request rate limit or overflowing the system. + +If you are already familiar with the synchronous version of pyFirecREST, you will find it quite straightforward to adapt to the asynchronous paradigm. + +We will be going through an example that will use the `asyncio library `__. +First you will need to create an ``AsyncFirecrest`` object, instead of the simple ``Firecrest`` object. + +.. code-block:: Python + + client = fc.AsyncFirecrest( + firecrest_url=, + authorization=MyAuthorizationClass() + ) + +As you can see in the reference, the methods of ``AsyncFirecrest`` have the same name as the ones from the simple client, with the same arguments and types, but you will need to use the async/await syntax when you call them. + +Here is an example of the calls we saw in the previous section: + +.. code-block:: Python + + # Getting all the systems + systems = await client.all_systems() + print(systems) + + # Getting the files of a directory + files = await client.list_files("cluster", "/home/test_user") + print(files) + + # Submit a job + job = await client.submit("cluster", "script.sh") + print(job) + + +The uploads and downloads work as before but you have to keep in mind which methods are coroutines. + +.. code-block:: Python + + # Download + down_obj = await client.external_download("cluster", "/remote/path/to/the/file") + status = await down_obj.status + print(status) + await down_obj.finish_download("my_local_file") + + # Upload + up_obj = await client.external_upload("cluster", "/path/to/local/file", "/remote/path/to/filesystem") + await up_obj.finish_upload() + status = await up_obj.status + print(status) + + +Here is a more complete example for how you could use the asynchronous client: + + +.. code-block:: Python + + import firecrest + import asyncio + import logging + + + # Setup variables before running the script + client_id = "" + client_secret = "" + token_uri = "" + firecrest_url = "" + + machine = "" + local_script_path = "" + + # This is simply setup for logging, you can ignore it + logger = logging.getLogger("simple_example") + logger.setLevel(logging.DEBUG) + ch = logging.StreamHandler() + ch.setLevel(logging.DEBUG) + formatter = logging.Formatter("%(asctime)s - %(message)s", datefmt="%H:%M:%S") + ch.setFormatter(formatter) + logger.addHandler(ch) + + async def workflow(client, i): + logger.info(f"{i}: Starting workflow") + job = await client.submit(machine, local_script_path) + logger.info(f"{i}: Submitted job with jobid: {job['jobid']}") + while True: + poll_res = await client.poll_active(machine, [job["jobid"]]) + if len(poll_res) < 1: + logger.info(f"{i}: Job {job['jobid']} is no longer active") + break + + logger.info(f"{i}: Job {job['jobid']} status: {poll_res[0]['state']}") + await asyncio.sleep(30) + + output = await client.view(machine, job["job_file_out"]) + logger.info(f"{i}: job output: {output}") + + + async def main(): + auth = firecrest.ClientCredentialsAuth(client_id, client_secret, token_uri) + client = firecrest.AsyncFirecrest(firecrest_url, authorization=auth) + + # Set up the desired polling rate for each microservice. The float number + # represents the number of seconds between consecutive requests in each + # microservice. + client.time_between_calls = { + "compute": 5, + "reservations": 5, + "status": 5, + "storage": 5, + "tasks": 5, + "utilities": 5, + } + + workflows = [workflow(client, i) for i in range(5)] + await asyncio.gather(*workflows) + + + asyncio.run(main()) diff --git a/docs/source/tutorial.rst b/docs/source/tutorial_basic.rst similarity index 73% rename from docs/source/tutorial.rst rename to docs/source/tutorial_basic.rst index 89248de..324c71e 100644 --- a/docs/source/tutorial.rst +++ b/docs/source/tutorial_basic.rst @@ -1,5 +1,5 @@ -Tutorial -======== +Simple tutorial +=============== Your starting point to use pyFirecREST will be the creation of a FirecREST object. This is simply a mini client that, in cooperation with the authorization object, will take care of the necessary requests that need to be made and handle the responses. @@ -9,7 +9,7 @@ For this tutorial we will assume the simplest kind of authorization class, where .. code-block:: Python - import firecrest as f7t + import firecrest as fc class MyAuthorizationClass: def __init__(self): @@ -19,7 +19,7 @@ For this tutorial we will assume the simplest kind of authorization class, where return # Setup the client with the appropriate URL and the authorization class - client = f7t.Firecrest(firecrest_url=, authorization=MyAuthorizationClass()) + client = fc.Firecrest(firecrest_url=, authorization=MyAuthorizationClass()) Simple blocking requests @@ -216,96 +216,3 @@ The simplest way to do the uploading through pyFirecREST is as follows: But, as before, you can get the necessary components for the upload from the ``object_storage_data`` property. You can get the link, as well as all the necessary arguments for the request to Object Storage and the full command you could perform manually from the terminal. - -Enable logging in your python code ----------------------------------- - -The simplest way to enable logging in your code would be to add this in the beginning of your file: - -.. code-block:: Python - - import logging - - logging.basicConfig( - level=logging.INFO, - format="%(levelname)s:%(name)s:%(message)s", - ) - -pyFirecREST has all of it's messages in `INFO` level. If you want to avoid messages from other packages, you can do the following: - -.. code-block:: Python - - import logging - - logging.basicConfig( - level=logging.WARNING, - format="%(levelname)s:%(name)s:%(message)s", - ) - logging.getLogger("firecrest").setLevel(logging.INFO) - - -Handling of errors ------------------- - -The methods of the Firecrest, ExternalUpload and ExternalDownload objects can raise exceptions in case something goes wrong. -When the error comes from the response of some request pyFirecREST will raise ``FirecrestException``. -In these cases you can manually examine all the responses from the requests in order to get more information, when the message is not informative enough. -These responses are from the requests package of python and you can get all types of useful information from it, like the status code, the json response, the headers and more. -Here is an example of the code that will handle those failures. - -.. code-block:: Python - - try: - parameters = client.parameters() - print(f"Firecrest parameters: {parameters}") - except fc.FirecrestException as e: - # You can just print the exception to get more information about the type of error, - # for example an invalid or expired token. - print(e) - # Or you can manually examine the responses. - print(e.responses[-1]) - except Exception as e: - # You might also get regular exceptions in some cases. For example when you are - # trying to upload a file that doesn't exist in your local filesystem. - pass - -CLI support ------------ - -After version 1.3.0, pyFirecREST comes together with a CLI but for now it can only be used with the `f7t.ClientCredentialsAuth` authentication class. - -You will need to set the environment variables ``FIRECREST_CLIENT_ID``, ``FIRECREST_CLIENT_SECRET`` and ``AUTH_TOKEN_URL`` to set up the Client Credentials client, as well as ``FIRECREST_URL`` with the URL for the FirecREST instance you are using. - -After that you can explore the capabilities of the CLI with the `--help` option: - -.. code-block:: bash - - firecrest --help - firecrest ls --help - firecrest submit --help - firecrest upload --help - firecrest download --help - firecrest submit-template --help - -Some basic examples: - -.. code-block:: bash - - # Get the available systems - firecrest systems - - # Get the parameters of different microservices of FirecREST - firecrest parameters - - # List files of directory - firecrest ls cluster1 /home - - # Submit a job - firecrest submit cluster script.sh - - # Upload a "small" file (you can check the maximum size in `UTILITIES_MAX_FILE_SIZE` from the `parameters` command) - firecrest upload --type=direct cluster local_file.txt /path/to/cluster/fs - - # Upload a "large" file - firecrest upload --type=external cluster local_file.txt /path/to/cluster/fs - # You will have to finish the upload with a second command that will be given in the output diff --git a/docs/source/tutorial_cli.rst b/docs/source/tutorial_cli.rst new file mode 100644 index 0000000..4e6c785 --- /dev/null +++ b/docs/source/tutorial_cli.rst @@ -0,0 +1,40 @@ +How to use the CLI +================== + +After version 1.3.0, pyFirecREST comes together with a CLI but for now it can only be used with the ``ClientCredentialsAuth`` authentication class. + +You will need to set the environment variables ``FIRECREST_CLIENT_ID``, ``FIRECREST_CLIENT_SECRET`` and ``AUTH_TOKEN_URL`` to set up the Client Credentials client, as well as ``FIRECREST_URL`` with the URL for the FirecREST instance you are using. + +After that you can explore the capabilities of the CLI with the `--help` option: + +.. code-block:: bash + + firecrest --help + firecrest ls --help + firecrest submit --help + firecrest upload --help + firecrest download --help + firecrest submit-template --help + +Some basic examples: + +.. code-block:: bash + + # Get the available systems + firecrest systems + + # Get the parameters of different microservices of FirecREST + firecrest parameters + + # List files of directory + firecrest ls cluster1 /home + + # Submit a job + firecrest submit cluster script.sh + + # Upload a "small" file (you can check the maximum size in `UTILITIES_MAX_FILE_SIZE` from the `parameters` command) + firecrest upload --type=direct cluster local_file.txt /path/to/cluster/fs + + # Upload a "large" file + firecrest upload --type=external cluster local_file.txt /path/to/cluster/fs + # You will have to finish the upload with a second command that will be given in the output diff --git a/docs/source/tutorial_errors.rst b/docs/source/tutorial_errors.rst new file mode 100644 index 0000000..71bb348 --- /dev/null +++ b/docs/source/tutorial_errors.rst @@ -0,0 +1,30 @@ +How to catch and debug errors +============================= + +The methods of the ``Firecrest``, ``ExternalUpload`` and ``ExternalDownload`` objects will raise exceptions in case something goes wrong. +The same happens for their asynchronous counterparts. +When the error comes from the response of some request pyFirecREST will raise ``FirecrestException``. +In these cases you can manually examine all the responses from the requests in order to get more information, when the message is not informative enough. +These responses are from the requests package of python and you can get all types of useful information from it, like the status code, the json response, the headers and more. +Here is an example of the code that will handle those failures. + +.. code-block:: Python + + import firecrest as fc + + + try: + parameters = client.parameters() + print(f"Firecrest parameters: {parameters}") + except fc.FirecrestException as e: + # You can just print the exception to get more information about the type of error, + # for example an invalid or expired token. + print(e) + # Or you can manually examine the responses. + print(e.responses[-1]) + print(e.responses[-1].status_code) + print(e.responses[-1].body) + except Exception as e: + # You might also get regular exceptions in some cases. For example when you are + # trying to upload a file that doesn't exist in your local filesystem. + print(f"A different exception was encountered: {e}") diff --git a/docs/source/tutorial_index.rst b/docs/source/tutorial_index.rst new file mode 100644 index 0000000..cfa7af3 --- /dev/null +++ b/docs/source/tutorial_index.rst @@ -0,0 +1,12 @@ +Tutorials +========= + +.. toctree:: + :maxdepth: 2 + :caption: Contents: + + tutorial_basic + tutorial_logging + tutorial_errors + tutorial_cli + tutorial_async diff --git a/docs/source/tutorial_logging.rst b/docs/source/tutorial_logging.rst new file mode 100644 index 0000000..0b0d635 --- /dev/null +++ b/docs/source/tutorial_logging.rst @@ -0,0 +1,26 @@ + +Enable logging in your python code +================================== + +The simplest way to enable logging in your code would be to add this in the beginning of your file: + +.. code-block:: Python + + import logging + + logging.basicConfig( + level=logging.INFO, + format="%(levelname)s:%(name)s:%(message)s", + ) + +pyFirecREST has all of it's messages in `INFO` level. If you want to avoid messages from other packages, you can do the following: + +.. code-block:: Python + + import logging + + logging.basicConfig( + level=logging.WARNING, + format="%(levelname)s:%(name)s:%(message)s", + ) + logging.getLogger("firecrest").setLevel(logging.INFO) diff --git a/firecrest/AsyncClient.py b/firecrest/AsyncClient.py new file mode 100644 index 0000000..2fb5b87 --- /dev/null +++ b/firecrest/AsyncClient.py @@ -0,0 +1,1397 @@ +# +# Copyright (c) 2019-2023, ETH Zurich. All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +# +from __future__ import annotations + +import asyncio +import httpx +from io import BytesIO +import jwt +import logging +import pathlib +import requests +import sys +import time + +from contextlib import nullcontext +from typing import ( + Any, + ContextManager, + Optional, + overload, + Sequence, + Tuple, + List, +) +from requests.compat import json # type: ignore +from packaging.version import Version, parse + +import firecrest.FirecrestException as fe +import firecrest.types as t +from firecrest.AsyncExternalStorage import AsyncExternalUpload, AsyncExternalDownload + + +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +logger = logging.getLogger(__name__) + +# This function is temporarily here +def handle_response(response): + print("\nResponse status code:") + print(response.status_code) + print("\nResponse headers:") + print(json.dumps(dict(response.headers), indent=4)) + print("\nResponse json:") + try: + print(json.dumps(response.json(), indent=4)) + except json.JSONDecodeError: + print("-") + + +class ComputeTask: + """Helper object for blocking methods that require multiple requests + """ + + def __init__( + self, + client: AsyncFirecrest, + task_id: str, + previous_responses: Optional[List[requests.Response]] = None, + ) -> None: + self._responses = [] if previous_responses is None else previous_responses + self._client = client + self._task_id = task_id + + async def poll_task(self, final_status): + logger.info(f"Polling task {self._task_id} until status is {final_status}") + resp = await self._client._task_safe(self._task_id, self._responses) + while resp["status"] < final_status: + # The rate limit is handled by the async client so no need to + # add a `sleep` here + resp = await self._client._task_safe(self._task_id, self._responses) + + logger.info(f'Status of {self._task_id} is {resp["status"]}') + return resp["data"] + + +class AsyncFirecrest: + """ + This is the basic class you instantiate to access the FirecREST API v1. + Necessary parameters are the firecrest URL and an authorization object. + + :param firecrest_url: FirecREST's URL + :param authorization: the authorization object. This object is responsible of handling the credentials and the only requirement for it is that it has a method get_access_token() that returns a valid access token. + :param verify: either a boolean, in which case it controls whether requests will verify the server’s TLS certificate, or a string, in which case it must be a path to a CA bundle to use + :param sa_role: this corresponds to the `F7T_AUTH_ROLE` configuration parameter of the site. If you don't know how FirecREST is setup it's better to leave the default. + """ + + TOO_MANY_REQUESTS_CODE = 429 + + def _retry_requests(func): + async def wrapper(*args, **kwargs): + client = args[0] + num_retries = 0 + resp = await func(*args, **kwargs) + while True: + if resp.status_code != client.TOO_MANY_REQUESTS_CODE: + break + elif ( + client.num_retries_rate_limit is not None + and num_retries >= client.num_retries_rate_limit + ): + logger.debug( + f"Rate limit is reached and the request has " + f"been retried already {num_retries} times" + ) + break + else: + reset = resp.headers.get( + "Retry-After", + default=resp.headers.get("RateLimit-Reset", default=10), + ) + reset = int(reset) + microservice = kwargs['endpoint'].split("/")[1] + client = args[0] + logger.info( + f"Rate limit in `{microservice}` is reached, next " + f"request will be possible in {reset} sec" + ) + client._next_request_ts[microservice] = ( + time.time() + reset + ) + resp = await func(*args, **kwargs) + num_retries += 1 + + return resp + + return wrapper + + def __init__( + self, + firecrest_url: str, + authorization: Any, + verify: Optional[str | bool] = None, + sa_role: str = "firecrest-sa", + ) -> None: + self._firecrest_url = firecrest_url + self._authorization = authorization + # This should be used only for blocking operations that require multiple requests, + # not for external upload/download + self._current_method_requests: List[requests.Response] = [] + self._verify = verify # TODO: not supported in httpx + self._sa_role = sa_role + #: This attribute will be passed to all the requests that will be made. + #: How many seconds to wait for the server to send data before giving up. + #: After that time a `requests.exceptions.Timeout` error will be raised. + #: + #: It can be a float or a tuple. More details here: https://www.python-httpx.org/advanced/#fine-tuning-the-configuration. + self.timeout: Any = None + # type is Any because of some incompatibility between https and requests library + + #: Number of retries in case the rate limit is reached. When it is set to `None`, the + #: client will keep trying until it gets a different status code than 429. + self.num_retries_rate_limit: Optional[int] = None + self._api_version: Version = parse("1.13.1") + self._session = httpx.AsyncClient() + + #: Seconds between requests in each microservice + self.time_between_calls: dict[str, float] = { # TODO more detailed docs + "compute": 5, + "reservations": 5, + "status": 5, + "storage": 5, + "tasks": 5, + "utilities": 5, + } + self._next_request_ts: dict[str, float] = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + self._locks = { + "compute": asyncio.Lock(), + "reservations": asyncio.Lock(), + "status": asyncio.Lock(), + "storage": asyncio.Lock(), + "tasks": asyncio.Lock(), + "utilities": asyncio.Lock(), + } + + def set_api_version(self, api_version: str) -> None: + """Set the version of the api of firecrest. By default it will be assumed that you are + using version 1.13.1 or compatible. The version is parsed by the `packaging` library. + """ + self._api_version = parse(api_version) + + @_retry_requests # type: ignore + async def _get_request(self, endpoint, additional_headers=None, params=None) -> httpx.Response: + microservice = endpoint.split("/")[1] + url = f"{self._firecrest_url}{endpoint}" + async with self._locks[microservice]: + await self._stall_request(microservice) + headers = { + "Authorization": f"Bearer {self._authorization.get_access_token()}" + } + if additional_headers: + headers.update(additional_headers) + + logger.info(f"Making GET request to {endpoint}") + resp = await self._session.get( + url=url, headers=headers, params=params, timeout=self.timeout + ) + self._next_request_ts[microservice] = ( + time.time() + self.time_between_calls[microservice] + ) + + return resp + + @_retry_requests # type: ignore + async def _post_request( + self, endpoint, additional_headers=None, data=None, files=None + ) -> httpx.Response: + microservice = endpoint.split("/")[1] + url = f"{self._firecrest_url}{endpoint}" + async with self._locks[microservice]: + await self._stall_request(microservice) + headers = { + "Authorization": f"Bearer {self._authorization.get_access_token()}" + } + if additional_headers: + headers.update(additional_headers) + + logger.info(f"Making POST request to {endpoint}") + resp = await self._session.post( + url=url, headers=headers, data=data, files=files, timeout=self.timeout + ) + self._next_request_ts[microservice] = ( + time.time() + self.time_between_calls[microservice] + ) + + return resp + + @_retry_requests # type: ignore + async def _put_request(self, endpoint, additional_headers=None, data=None) -> httpx.Response: + microservice = endpoint.split("/")[1] + url = f"{self._firecrest_url}{endpoint}" + async with self._locks[microservice]: + await self._stall_request(microservice) + headers = { + "Authorization": f"Bearer {self._authorization.get_access_token()}" + } + if additional_headers: + headers.update(additional_headers) + + logger.info(f"Making PUT request to {endpoint}") + resp = await self._session.put( + url=url, headers=headers, data=data, timeout=self.timeout + ) + self._next_request_ts[microservice] = ( + time.time() + self.time_between_calls[microservice] + ) + + return resp + + @_retry_requests # type: ignore + async def _delete_request( + self, endpoint, additional_headers=None, data=None + ) -> requests.Response: + microservice = endpoint.split("/")[1] + url = f"{self._firecrest_url}{endpoint}" + async with self._locks[microservice]: + await self._stall_request(microservice) + headers = { + "Authorization": f"Bearer {self._authorization.get_access_token()}" + } + if additional_headers: + headers.update(additional_headers) + + logger.info(f"Making DELETE request to {endpoint}") + # TODO: httpx doesn't support data in delete so we will have to + # keep using the `requests` package for this + resp = requests.delete( + url=url, headers=headers, data=data, timeout=self.timeout + ) + self._next_request_ts[microservice] = ( + time.time() + self.time_between_calls[microservice] + ) + + return resp + + async def _stall_request(self, microservice: str) -> None: + if self._next_request_ts[microservice] is not None: + while time.time() <= self._next_request_ts[microservice]: + logger.debug( + f"`{microservice}` microservice has received too many requests. " + f"Going to sleep for " + f"~{self._next_request_ts[microservice] - time.time()} sec" + ) + await asyncio.sleep(self._next_request_ts[microservice] - time.time()) + + @overload + def _json_response( + self, + responses: List[requests.Response], + expected_status_code: int, + allow_none_result: Literal[False] = ..., + ) -> dict: + ... + + @overload + def _json_response( + self, + responses: List[requests.Response], + expected_status_code: int, + allow_none_result: Literal[True], + ) -> Optional[dict]: + ... + + def _json_response( + self, + responses: List[requests.Response], + expected_status_code: int, + allow_none_result: bool = False, + ): + # Will examine only the last response + response = responses[-1] + status_code = response.status_code + # handle_response(response) + exc: fe.FirecrestException + for h in fe.ERROR_HEADERS: + if h in response.headers: + logger.critical(f"Header '{h}' is included in the response") + exc = fe.HeaderException(responses) + logger.critical(exc) + raise exc + + if status_code == 401: + logger.critical(f"Status of the response is 401") + exc = fe.UnauthorizedException(responses) + logger.critical(exc) + raise exc + elif status_code == 404: + logger.critical(f"Status of the response is 404") + exc = fe.NotFound(responses) + logger.critical(exc) + raise exc + elif status_code >= 400: + logger.critical(f"Status of the response is {status_code}") + exc = fe.FirecrestException(responses) + logger.critical(exc) + raise exc + elif status_code != expected_status_code: + logger.critical( + f"Unexpected status of last request {status_code}, it should have been {expected_status_code}" + ) + exc = fe.UnexpectedStatusException(responses, expected_status_code) + logger.critical(exc) + raise exc + + try: + ret = response.json() + except json.decoder.JSONDecodeError: + if allow_none_result: + ret = None + else: + exc = fe.NoJSONException(responses) + logger.critical(exc) + raise exc + + return ret + + async def _tasks( + self, + task_ids: Optional[List[str]] = None, + responses: Optional[List[requests.Response]] = None, + ) -> dict[str, t.Task]: + """Return a dictionary of FirecREST tasks and their last update. + When `task_ids` is an empty list or contains more than one element the + `/tasks` endpoint will be called. Otherwise `/tasks/{taskid}`. + When the `/tasks` is called the method will not give an error for invalid IDs, + but `/tasks/{taskid}` will raise an exception. + + :param task_ids: list of task IDs. When empty all tasks are returned. + :param responses: list of responses that are associated with these tasks (only relevant for error) + :calls: GET `/tasks` or `/tasks/{taskid}` + """ + task_ids = [] if task_ids is None else task_ids + responses = [] if responses is None else responses + endpoint = "/tasks/" + if len(task_ids) == 1: + endpoint += task_ids[0] + + resp = await self._get_request(endpoint=endpoint) + responses.append(resp) + taskinfo = self._json_response(responses, 200) + if len(task_ids) == 0: + return taskinfo["tasks"] + elif len(task_ids) == 1: + return {task_ids[0]: taskinfo["task"]} + else: + return {k: v for k, v in taskinfo["tasks"].items() if k in task_ids} + + async def _task_safe( + self, task_id: str, responses: Optional[List[requests.Response]] = None + ) -> t.Task: + if responses is None: + responses = self._current_method_requests + + task = (await self._tasks([task_id], responses))[task_id] + status = int(task["status"]) + exc: fe.FirecrestException + if status == 115: + logger.critical("Task has error status code 115") + exc = fe.StorageUploadException(responses) + logger.critical(exc) + raise exc + + if status == 118: + logger.critical("Task has error status code 118") + exc = fe.StorageDownloadException(responses) + logger.critical(exc) + raise exc + + if status >= 400: + logger.critical(f"Task has error status code {status}") + exc = fe.FirecrestException(responses) + logger.critical(exc) + raise exc + + return task + + async def _invalidate( + self, task_id: str, responses: Optional[List[requests.Response]] = None + ): + responses = [] if responses is None else responses + resp = await self._post_request( + endpoint="/storage/xfer-external/invalidate", + additional_headers={"X-Task-Id": task_id}, + ) + responses.append(resp) + return self._json_response(responses, 201, allow_none_result=True) + + async def _poll_tasks(self, task_id: str, final_status, sleep_time): + logger.info(f"Polling task {task_id} until status is {final_status}") + resp = await self._task_safe(task_id) + while resp["status"] < final_status: + t = next(sleep_time) + logger.info( + f'Status of {task_id} is {resp["status"]}, sleeping for {t} sec' + ) + await asyncio.sleep(t) + resp = await self._task_safe(task_id) + + logger.info(f'Status of {task_id} is {resp["status"]}') + return resp["data"] + + # Status + async def all_services(self) -> List[t.Service]: + """Returns a list containing all available micro services with a name, description, and status. + + :calls: GET `/status/services` + """ + resp = await self._get_request(endpoint="/status/services") + return self._json_response([resp], 200)["out"] + + async def service(self, service_name: str) -> t.Service: + """Returns information about a micro service. + Returns the name, description, and status. + + :param service_name: the service name + :calls: GET `/status/services/{service_name}` + """ + resp = await self._get_request(endpoint=f"/status/services/{service_name}") + return self._json_response([resp], 200) # type: ignore + + async def all_systems(self) -> List[t.System]: + """Returns a list containing all available systems and response status. + + :calls: GET `/status/systems` + """ + resp = await self._get_request(endpoint="/status/systems") + return self._json_response([resp], 200)["out"] + + async def system(self, system_name: str) -> t.System: + """Returns information about a system. + Returns the name, description, and status. + + :param system_name: the system name + :calls: GET `/status/systems/{system_name}` + """ + resp = await self._get_request(endpoint=f"/status/systems/{system_name}") + return self._json_response([resp], 200)["out"] + + async def parameters(self) -> t.Parameters: + """Returns configuration parameters of the FirecREST deployment that is associated with the client. + + :calls: GET `/status/parameters` + """ + resp = await self._get_request(endpoint="/status/parameters") + return self._json_response([resp], 200)["out"] + + # Utilities + async def list_files( + self, machine: str, target_path: str, show_hidden: bool = False + ) -> List[t.LsFile]: + """Returns a list of files in a directory. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :param show_hidden: show hidden files + :calls: GET `/utilities/ls` + """ + params: dict[str, Any] = {"targetPath": f"{target_path}"} + if show_hidden is True: + params["showhidden"] = show_hidden + + resp = await self._get_request( + endpoint="/utilities/ls", + additional_headers={"X-Machine-Name": machine}, + params=params, + ) + return self._json_response([resp], 200)["output"] + + async def mkdir( + self, machine: str, target_path: str, p: Optional[bool] = None + ) -> None: + """Creates a new directory. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :param p: no error if existing, make parent directories as needed + :calls: POST `/utilities/mkdir` + """ + data: dict[str, Any] = {"targetPath": target_path} + if p: + data["p"] = p + + resp = await self._post_request( + endpoint="/utilities/mkdir", + additional_headers={"X-Machine-Name": machine}, + data=data, + ) + self._json_response([resp], 201) + + async def mv(self, machine: str, source_path: str, target_path: str) -> None: + """Rename/move a file, directory, or symlink at the `source_path` to the `target_path` on `machine`'s filesystem. + + :param machine: the machine name where the filesystem belongs to + :param source_path: the absolute source path + :param target_path: the absolute target path + :calls: PUT `/utilities/rename` + """ + resp = await self._put_request( + endpoint="/utilities/rename", + additional_headers={"X-Machine-Name": machine}, + data={"targetPath": target_path, "sourcePath": source_path}, + ) + self._json_response([resp], 200) + + async def chmod(self, machine: str, target_path: str, mode: str) -> None: + """Changes the file mod bits of a given file according to the specified mode. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :param mode: same as numeric mode of linux chmod tool + :calls: PUT `/utilities/chmod` + """ + resp = await self._put_request( + endpoint="/utilities/chmod", + additional_headers={"X-Machine-Name": machine}, + data={"targetPath": target_path, "mode": mode}, + ) + self._json_response([resp], 200) + + async def chown( + self, + machine: str, + target_path: str, + owner: Optional[str] = None, + group: Optional[str] = None, + ) -> None: + """Changes the user and/or group ownership of a given file. + If only owner or group information is passed, only that information will be updated. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :param owner: owner ID for target + :param group: group ID for target + :calls: PUT `/utilities/chown` + """ + if owner is None and group is None: + return + + data = {"targetPath": target_path} + if owner: + data["owner"] = owner + + if group: + data["group"] = group + + resp = await self._put_request( + endpoint="/utilities/chown", + additional_headers={"X-Machine-Name": machine}, + data=data, + ) + self._json_response([resp], 200) + + async def copy(self, machine: str, source_path: str, target_path: str) -> None: + """Copies file from `source_path` to `target_path`. + + :param machine: the machine name where the filesystem belongs to + :param source_path: the absolute source path + :param target_path: the absolute target path + :calls: POST `/utilities/copy` + """ + resp = await self._post_request( + endpoint="/utilities/copy", + additional_headers={"X-Machine-Name": machine}, + data={"targetPath": target_path, "sourcePath": source_path}, + ) + self._json_response([resp], 201) + + async def file_type(self, machine: str, target_path: str) -> str: + """Uses the `file` linux application to determine the type of a file. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :calls: GET `/utilities/file` + """ + resp = await self._get_request( + endpoint="/utilities/file", + additional_headers={"X-Machine-Name": machine}, + params={"targetPath": target_path}, + ) + return self._json_response([resp], 200)["output"] + + async def stat( + self, machine: str, target_path: str, dereference: bool = False + ) -> t.StatFile: + """Uses the stat linux application to determine the status of a file on the machine's filesystem. + The result follows: https://docs.python.org/3/library/os.html#os.stat_result. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :param dereference: follow link (default False) + :calls: GET `/utilities/stat` + """ + params: dict[str, Any] = {"targetPath": target_path} + if dereference: + params["dereference"] = dereference + + resp = await self._get_request( + endpoint="/utilities/stat", + additional_headers={"X-Machine-Name": machine}, + params=params, + ) + return self._json_response([resp], 200)["output"] + + async def symlink(self, machine: str, target_path: str, link_path: str) -> None: + """Creates a symbolic link. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute path that the symlink will point to + :param link_path: the absolute path to the new symlink + :calls: POST `/utilities/symlink` + """ + resp = await self._post_request( + endpoint="/utilities/symlink", + additional_headers={"X-Machine-Name": machine}, + data={"targetPath": target_path, "linkPath": link_path}, + ) + self._json_response([resp], 201) + + async def simple_download( + self, machine: str, source_path: str, target_path: str | pathlib.Path | BytesIO + ) -> None: + """Blocking call to download a small file. + The maximun size of file that is allowed can be found from the parameters() call. + + :param machine: the machine name where the filesystem belongs to + :param source_path: the absolute source path + :param target_path: the target path in the local filesystem or binary stream + :calls: GET `/utilities/download` + """ + resp = await self._get_request( + endpoint="/utilities/download", + additional_headers={"X-Machine-Name": machine}, + params={"sourcePath": source_path}, + ) + self._json_response([resp], 200, allow_none_result=True) + context: ContextManager[BytesIO] = ( + open(target_path, "wb") # type: ignore + if isinstance(target_path, str) or isinstance(target_path, pathlib.Path) + else nullcontext(target_path) + ) + with context as f: + f.write(resp.content) + + async def simple_upload( + self, + machine: str, + source_path: str | pathlib.Path | BytesIO, + target_path: str, + filename: Optional[str] = None, + ) -> None: + """Blocking call to upload a small file. + The file that will be uploaded will have the same name as the source_path. + The maximum size of file that is allowed can be found from the parameters() call. + + :param machine: the machine name where the filesystem belongs to + :param source_path: the source path of the file or binary stream + :param target_path: the absolute target path of the directory where the file will be uploaded + :param filename: naming target file to filename (default is same as the local one) + :calls: POST `/utilities/upload` + """ + context: ContextManager[BytesIO] = ( + open(source_path, "rb") # type: ignore + if isinstance(source_path, str) or isinstance(source_path, pathlib.Path) + else nullcontext(source_path) + ) + with context as f: + # Set filename + if filename is not None: + f = (filename, f) # type: ignore + + resp = await self._post_request( + endpoint="/utilities/upload", + additional_headers={"X-Machine-Name": machine}, + data={"targetPath": target_path}, + files={"file": f}, + ) + + self._json_response([resp], 201) + + async def simple_delete(self, machine: str, target_path: str) -> None: + """Blocking call to delete a small file. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :calls: DELETE `/utilities/rm` + """ + resp = await self._delete_request( + endpoint="/utilities/rm", + additional_headers={"X-Machine-Name": machine}, + data={"targetPath": target_path}, + ) + self._json_response([resp], 204, allow_none_result=True) + + async def checksum(self, machine: str, target_path: str) -> str: + """Calculate the SHA256 (256-bit) checksum of a specified file. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :calls: GET `/utilities/checksum` + """ + resp = await self._get_request( + endpoint="/utilities/checksum", + additional_headers={"X-Machine-Name": machine}, + params={"targetPath": target_path}, + ) + return self._json_response([resp], 200)["output"] + + async def head( + self, + machine: str, + target_path: str, + bytes: Optional[int] = None, + lines: Optional[int] = None, + ) -> str: + """Display the beginning of a specified file. + By default 10 lines will be returned. + Bytes and lines cannot be specified simultaneously. + The final result will be smaller than `UTILITIES_MAX_FILE_SIZE` bytes. + This variable is available in the parameters command. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :param lines: the number of lines to be displayed + :param bytes: the number of bytes to be displayed + :calls: GET `/utilities/head` + """ + resp = await self._get_request( + endpoint="/utilities/head", + additional_headers={"X-Machine-Name": machine}, + params={"targetPath": target_path, "lines": lines, "bytes": bytes}, + ) + return self._json_response([resp], 200)["output"] + + async def tail( + self, + machine: str, + target_path: str, + bytes: Optional[int] = None, + lines: Optional[int] = None, + ) -> str: + """Display the last part of a specified file. + By default 10 lines will be returned. + Bytes and lines cannot be specified simultaneously. + The final result will be smaller than `UTILITIES_MAX_FILE_SIZE` bytes. + This variable is available in the parameters command. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :param lines: the number of lines to be displayed + :param bytes: the number of bytes to be displayed + :calls: GET `/utilities/head` + """ + resp = await self._get_request( + endpoint="/utilities/tail", + additional_headers={"X-Machine-Name": machine}, + params={"targetPath": target_path, "lines": lines, "bytes": bytes}, + ) + return self._json_response([resp], 200)["output"] + + async def view(self, machine: str, target_path: str) -> str: + """View the content of a specified file. + The final result will be smaller than `UTILITIES_MAX_FILE_SIZE` bytes. + This variable is available in the parameters command. + + :param machine: the machine name where the filesystem belongs to + :param target_path: the absolute target path + :calls: GET `/utilities/checksum` + """ + resp = await self._get_request( + endpoint="/utilities/view", + additional_headers={"X-Machine-Name": machine}, + params={"targetPath": target_path}, + ) + return self._json_response([resp], 200)["output"] + + async def whoami(self, machine=None) -> Optional[str]: + """Returns the username that FirecREST will be using to perform the other calls. + In the case the machine name is passed in the arguments, a call is made to the respective endpoint and the command whoami is run on the machine. + Otherwise, the library decodes the token and will return `None` if the token is not valid. + + :calls: GET `/utilities/whoami` + """ + if machine: + resp = await self._get_request( + endpoint="/utilities/whoami", + additional_headers={"X-Machine-Name": machine}, + ) + return self._json_response([resp], 200)["output"] + + try: + decoded = jwt.decode( + self._authorization.get_access_token(), + options={"verify_signature": False}, + ) + try: + if self._sa_role in decoded["realm_access"]["roles"]: + clientId = decoded["clientId"] + username = decoded["resource_access"][clientId]["roles"][0] + return username + + return decoded["preferred_username"] + except KeyError: + return decoded["preferred_username"] + + except Exception: + # Invalid token, cannot retrieve username + return None + + # Compute + async def _submit_request(self, machine: str, job_script, local_file, account=None): + if local_file: + with open(job_script, "rb") as f: + if account: + data = {"account": account} + else: + data = None + + resp = await self._post_request( + endpoint="/compute/jobs/upload", + additional_headers={"X-Machine-Name": machine}, + files={"file": f}, + data=data, + ) + else: + data = {"targetPath": job_script} + if account: + data["account"] = account + + resp = await self._post_request( + endpoint="/compute/jobs/path", + additional_headers={"X-Machine-Name": machine}, + data=data, + ) + + return self._json_response([resp], 201) + + async def _squeue_request(self, machine: str, jobs=None): + jobs = [] if jobs is None else jobs + params = {} + if jobs: + params = {"jobs": ",".join([str(j) for j in jobs])} + + resp = await self._get_request( + endpoint="/compute/jobs", + additional_headers={"X-Machine-Name": machine}, + params=params, + ) + return self._json_response([resp], 200) + + async def _acct_request( + self, machine: str, jobs=None, start_time=None, end_time=None + ): + jobs = [] if jobs is None else jobs + params = {} + if jobs: + params["jobs"] = ",".join(jobs) + + if start_time: + params["starttime"] = start_time + + if end_time: + params["endtime"] = end_time + + resp = await self._get_request( + endpoint="/compute/acct", + additional_headers={"X-Machine-Name": machine}, + params=params, + ) + return self._json_response([resp], 200) + + async def submit( + self, + machine: str, + job_script: str, + local_file: Optional[bool] = True, + account: Optional[str] = None, + ) -> t.JobSubmit: + """Submits a batch script to SLURM on the target system + + :param machine: the machine name where the scheduler belongs to + :param job_script: the path of the script (if it's local it can be relative path, if it is on the machine it has to be the absolute path) + :param local_file: batch file can be local (default) or on the machine's filesystem + :param account: submit the job with this project account + :calls: POST `/compute/jobs/upload` or POST `/compute/jobs/path` + + GET `/tasks/{taskid}` + """ + if local_file: + with open(job_script, "rb") as f: + if account: + data = {"account": account} + else: + data = None + + resp = await self._post_request( + endpoint="/compute/jobs/upload", + additional_headers={"X-Machine-Name": machine}, + files={"file": f}, + data=data, + ) + else: + data = {"targetPath": job_script} + if account: + data["account"] = account + + resp = await self._post_request( + endpoint="/compute/jobs/path", + additional_headers={"X-Machine-Name": machine}, + data=data, + ) + + json_response = self._json_response([resp], 201) + logger.info(f"Job submission task: {json_response['task_id']}") + t = ComputeTask(self, json_response["task_id"], [resp]) + return await t.poll_task("200") + + async def poll( + self, + machine: str, + jobs: Optional[Sequence[str | int]] = None, + start_time: Optional[str] = None, + end_time: Optional[str] = None, + ) -> List[t.JobAcct]: + """Retrieves information about submitted jobs. + This call uses the `sacct` command. + + :param machine: the machine name where the scheduler belongs to + :param jobs: list of the IDs of the jobs + :param start_time: Start time (and/or date) of job's query. Allowed formats are HH:MM[:SS] [AM|PM] MMDD[YY] or MM/DD[/YY] or MM.DD[.YY] MM/DD[/YY]-HH:MM[:SS] YYYY-MM-DD[THH:MM[:SS]] + :param end_time: End time (and/or date) of job's query. Allowed formats are HH:MM[:SS] [AM|PM] MMDD[YY] or MM/DD[/YY] or MM.DD[.YY] MM/DD[/YY]-HH:MM[:SS] YYYY-MM-DD[THH:MM[:SS]] + :calls: GET `/compute/acct` + + GET `/tasks/{taskid}` + """ + jobids = [str(j) for j in jobs] if jobs else [] + params = {} + if jobids: + params["jobs"] = ",".join(jobids) + + if start_time: + params["starttime"] = start_time + + if end_time: + params["endtime"] = end_time + + resp = await self._get_request( + endpoint="/compute/acct", + additional_headers={"X-Machine-Name": machine}, + params=params, + ) + json_response = self._json_response([resp], 200) + logger.info(f"Job polling task: {json_response['task_id']}") + t = ComputeTask(self, json_response["task_id"], [resp]) + res = await t.poll_task("200") + # When there is no job in the sacct output firecrest will return an empty dictionary instead of list + if isinstance(res, dict): + return list(res.values()) + else: + return res + + async def poll_active( + self, machine: str, jobs: Optional[Sequence[str | int]] = None + ) -> List[t.JobQueue]: + """Retrieves information about active jobs. + This call uses the `squeue -u ` command. + + :param machine: the machine name where the scheduler belongs to + :param jobs: list of the IDs of the jobs + :calls: GET `/compute/jobs` + + GET `/tasks/{taskid}` + """ + jobs = jobs if jobs else [] + jobids = [str(j) for j in jobs] + params = {} + if jobs: + params = {"jobs": ",".join([str(j) for j in jobids])} + + resp = await self._get_request( + endpoint="/compute/jobs", + additional_headers={"X-Machine-Name": machine}, + params=params, + ) + json_response = self._json_response([resp], 200) + logger.info(f"Job active polling task: {json_response['task_id']}") + t = ComputeTask(self, json_response["task_id"], [resp]) + dict_result = await t.poll_task("200") + return list(dict_result.values()) + + async def cancel(self, machine: str, job_id: str | int) -> str: + """Cancels running job. + This call uses the `scancel` command. + + :param machine: the machine name where the scheduler belongs to + :param job_id: the ID of the job + :calls: DELETE `/compute/jobs/{job_id}` + + GET `/tasks/{taskid}` + """ + resp = await self._delete_request( + endpoint=f"/compute/jobs/{job_id}", + additional_headers={"X-Machine-Name": machine}, + ) + json_response = self._json_response([resp], 200) + logger.info(f"Job cancellation task: {json_response['task_id']}") + t = ComputeTask(self, json_response["task_id"], [resp]) + return await t.poll_task("200") + + # Storage + async def _internal_transfer( + self, + endpoint, + machine, + source_path, + target_path, + job_name, + time, + stage_out_job_id, + account, + ret_response, + ): + data = {"targetPath": target_path} + if source_path: + data["sourcePath"] = source_path + + if job_name: + data["jobname"] = job_name + + if time: + data["time"] = time + + if stage_out_job_id: + data["stageOutJobId"] = stage_out_job_id + + if account: + data["account"] = account + + resp = await self._post_request( + endpoint=endpoint, additional_headers={"X-Machine-Name": machine}, data=data + ) + ret_response.append(resp) + return self._json_response([resp], 201) + + async def submit_move_job( + self, + machine: str, + source_path: str, + target_path: str, + job_name: Optional[str] = None, + time: Optional[str] = None, + stage_out_job_id: Optional[str] = None, + account: Optional[str] = None, + ) -> t.JobSubmit: + """Move files between internal CSCS file systems. + Rename/Move source_path to target_path. + Possible to stage-out jobs providing the SLURM ID of a production job. + More info about internal transfer: https://user.cscs.ch/storage/data_transfer/internal_transfer/ + + :param machine: the machine name where the scheduler belongs to + :param source_path: the absolute source path + :param target_path: the absolute target path + :param job_name: job name + :param time: limit on the total run time of the job. Acceptable time formats 'minutes', 'minutes:seconds', 'hours:minutes:seconds', 'days-hours', 'days-hours:minutes' and 'days-hours:minutes:seconds'. + :param stage_out_job_id: transfer data after job with ID {stage_out_job_id} is completed + :param account: name of the bank account to be used in SLURM. If not set, system default is taken. + :calls: POST `/storage/xfer-internal/mv` + + GET `/tasks/{taskid}` + """ + resp: List[requests.Response] = [] + endpoint = "/storage/xfer-internal/mv" + json_response = await self._internal_transfer( + endpoint, + machine, + source_path, + target_path, + job_name, + time, + stage_out_job_id, + account, + resp, + ) + logger.info(f"Job submission task: {json_response['task_id']}") + t = ComputeTask(self, json_response["task_id"], resp) + return await t.poll_task("200") + + async def submit_copy_job( + self, + machine: str, + source_path: str, + target_path: str, + job_name: Optional[str] = None, + time: Optional[str] = None, + stage_out_job_id: Optional[str] = None, + account: Optional[str] = None, + ) -> t.JobSubmit: + """Copy files between internal CSCS file systems. + Copy source_path to target_path. + Possible to stage-out jobs providing the SLURM Id of a production job. + More info about internal transfer: https://user.cscs.ch/storage/data_transfer/internal_transfer/ + + :param machine: the machine name where the scheduler belongs to + :param source_path: the absolute source path + :param target_path: the absolute target path + :param job_name: job name + :param time: limit on the total run time of the job. Acceptable time formats 'minutes', 'minutes:seconds', 'hours:minutes:seconds', 'days-hours', 'days-hours:minutes' and 'days-hours:minutes:seconds'. + :param stage_out_job_id: transfer data after job with ID {stage_out_job_id} is completed + :param account: name of the bank account to be used in SLURM. If not set, system default is taken. + :calls: POST `/storage/xfer-internal/cp` + + GET `/tasks/{taskid}` + """ + resp: List[requests.Response] = [] + endpoint = "/storage/xfer-internal/cp" + json_response = await self._internal_transfer( + endpoint, + machine, + source_path, + target_path, + job_name, + time, + stage_out_job_id, + account, + resp, + ) + logger.info(f"Job submission task: {json_response['task_id']}") + t = ComputeTask(self, json_response["task_id"], resp) + return await t.poll_task("200") + + async def submit_rsync_job( + self, + machine: str, + source_path: str, + target_path: str, + job_name: Optional[str] = None, + time: Optional[str] = None, + stage_out_job_id: Optional[str] = None, + account: Optional[str] = None, + ) -> t.JobSubmit: + """Transfer files between internal CSCS file systems. + Transfer source_path to target_path. + Possible to stage-out jobs providing the SLURM Id of a production job. + More info about internal transfer: https://user.cscs.ch/storage/data_transfer/internal_transfer/ + + :param machine: the machine name where the scheduler belongs to + :param source_path: the absolute source path + :param target_path: the absolute target path + :param job_name: job name + :param time: limit on the total run time of the job. Acceptable time formats 'minutes', 'minutes:seconds', 'hours:minutes:seconds', 'days-hours', 'days-hours:minutes' and 'days-hours:minutes:seconds'. + :param stage_out_job_id: transfer data after job with ID {stage_out_job_id} is completed + :param account: name of the bank account to be used in SLURM. If not set, system default is taken. + :calls: POST `/storage/xfer-internal/rsync` + + GET `/tasks/{taskid}` + """ + resp: List[requests.Response] = [] + endpoint = "/storage/xfer-internal/rsync" + json_response = await self._internal_transfer( + endpoint, + machine, + source_path, + target_path, + job_name, + time, + stage_out_job_id, + account, + resp, + ) + logger.info(f"Job submission task: {json_response['task_id']}") + t = ComputeTask(self, json_response["task_id"], resp) + return await t.poll_task("200") + + async def submit_delete_job( + self, + machine: str, + target_path: str, + job_name: Optional[str] = None, + time: Optional[str] = None, + stage_out_job_id: Optional[str] = None, + account: Optional[str] = None, + ) -> t.JobSubmit: + """Remove files in internal CSCS file systems. + Remove file in target_path. + Possible to stage-out jobs providing the SLURM Id of a production job. + More info about internal transfer: https://user.cscs.ch/storage/data_transfer/internal_transfer/ + + :param machine: the machine name where the scheduler belongs to + :param target_path: the absolute target path + :param job_name: job name + :param time: limit on the total run time of the job. Acceptable time formats 'minutes', 'minutes:seconds', 'hours:minutes:seconds', 'days-hours', 'days-hours:minutes' and 'days-hours:minutes:seconds'. + :param stage_out_job_id: transfer data after job with ID {stage_out_job_id} is completed + :param account: name of the bank account to be used in SLURM. If not set, system default is taken. + :calls: POST `/storage/xfer-internal/rm` + + GET `/tasks/{taskid}` + """ + resp: List[requests.Response] = [] + endpoint = "/storage/xfer-internal/rm" + json_response = await self._internal_transfer( + endpoint, + machine, + None, + target_path, + job_name, + time, + stage_out_job_id, + account, + resp, + ) + logger.info(f"Job submission task: {json_response['task_id']}") + t = ComputeTask(self, json_response["task_id"], resp) + return await t.poll_task("200") + + async def external_upload( + self, machine: str, source_path: str, target_path: str + ) -> AsyncExternalUpload: + """Non blocking call for the upload of larger files. + + :param machine: the machine where the filesystem belongs to + :param source_path: the source path in the local filesystem + :param target_path: the target path in the machine's filesystem + :returns: an ExternalUpload object + """ + resp = await self._post_request( + endpoint="/storage/xfer-external/upload", + additional_headers={"X-Machine-Name": machine}, + data={"targetPath": target_path, "sourcePath": source_path}, + ) + json_response = self._json_response([resp], 201)["task_id"] + return AsyncExternalUpload(self, json_response, [resp]) + + async def external_download( + self, machine: str, source_path: str + ) -> AsyncExternalDownload: + """Non blocking call for the download of larger files. + + :param machine: the machine where the filesystem belongs to + :param source_path: the source path in the local filesystem + :returns: an ExternalDownload object + """ + resp = await self._post_request( + endpoint="/storage/xfer-external/download", + additional_headers={"X-Machine-Name": machine}, + data={"sourcePath": source_path}, + ) + return AsyncExternalDownload( + self, self._json_response([resp], 201)["task_id"], [resp] + ) + + # Reservation + async def all_reservations(self, machine: str) -> List[dict]: + """List all active reservations and their status + + :param machine: the machine name + :calls: GET `/reservations` + """ + resp = await self._get_request( + endpoint="/reservations", additional_headers={"X-Machine-Name": machine} + ) + return self._json_response([resp], 200)["success"] + + async def create_reservation( + self, + machine: str, + reservation: str, + account: str, + number_of_nodes: str, + node_type: str, + start_time: str, + end_time: str, + ) -> None: + """Creates a new reservation with {reservation} name for a given SLURM groupname + + :param machine: the machine name + :param reservation: the reservation name + :param account: the account in SLURM to which the reservation is made for + :param number_of_nodes: number of nodes needed for the reservation + :param node_type: type of node + :param start_time: start time for reservation (YYYY-MM-DDTHH:MM:SS) + :param end_time: end time for reservation (YYYY-MM-DDTHH:MM:SS) + :calls: POST `/reservations` + """ + data = { + "reservation": reservation, + "account": account, + "numberOfNodes": number_of_nodes, + "nodeType": node_type, + "starttime": start_time, + "endtime": end_time, + } + resp = await self._post_request( + endpoint="/reservations", + additional_headers={"X-Machine-Name": machine}, + data=data, + ) + self._json_response([resp], 201) + + async def update_reservation( + self, + machine: str, + reservation: str, + account: str, + number_of_nodes: str, + node_type: str, + start_time: str, + end_time: str, + ) -> None: + """Updates an already created reservation named {reservation} + + :param machine: the machine name + :param reservation: the reservation name + :param account: the account in SLURM to which the reservation is made for + :param number_of_nodes: number of nodes needed for the reservation + :param node_type: type of node + :param start_time: start time for reservation (YYYY-MM-DDTHH:MM:SS) + :param end_time: end time for reservation (YYYY-MM-DDTHH:MM:SS) + :calls: PUT `/reservations/{reservation}` + """ + data = { + "account": account, + "numberOfNodes": number_of_nodes, + "nodeType": node_type, + "starttime": start_time, + "endtime": end_time, + } + resp = await self._put_request( + endpoint=f"/reservations/{reservation}", + additional_headers={"X-Machine-Name": machine}, + data=data, + ) + self._json_response([resp], 200) + + async def delete_reservation(self, machine: str, reservation: str) -> None: + """Deletes an already created reservation named {reservation} + + :param machine: the machine name + :param reservation: the reservation name + :calls: DELETE `/reservations/{reservation}` + """ + resp = await self._delete_request( + endpoint=f"/reservations/{reservation}", + additional_headers={"X-Machine-Name": machine}, + ) + self._json_response([resp], 204, allow_none_result=True) diff --git a/firecrest/AsyncExternalStorage.py b/firecrest/AsyncExternalStorage.py new file mode 100644 index 0000000..4dc1533 --- /dev/null +++ b/firecrest/AsyncExternalStorage.py @@ -0,0 +1,261 @@ +# +# Copyright (c) 2019-2023, ETH Zurich. All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +# +from __future__ import annotations + +import asyncio +from io import BufferedWriter +import itertools +import logging +import pathlib +import requests +import shutil +import sys +from typing import ( + ContextManager, + Optional, + List, + TYPE_CHECKING, +) +import urllib.request +from packaging.version import Version + +if TYPE_CHECKING: + from firecrest.AsyncClient import AsyncFirecrest + +from contextlib import nullcontext +from requests.compat import json # type: ignore + +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +logger = logging.getLogger(__name__) + + +class AsyncExternalStorage: + """External storage object. + """ + + _final_states: set[str] + + def __init__( + self, + client: AsyncFirecrest, + task_id: str, + previous_responses: Optional[List[requests.Response]] = None, + ) -> None: + previous_responses = [] if previous_responses is None else previous_responses + self._client = client + self._task_id = task_id + self._in_progress = True + self._status: Optional[str] = None + self._data = None + self._object_storage_data = None + self._responses = previous_responses + + @property + def client(self) -> AsyncFirecrest: + """Returns the client that will be used to get information for the task. + """ + return self._client + + @property + def task_id(self) -> str: + """Returns the FirecREST task ID that is associated with this transfer. + """ + return self._task_id + + async def _update(self) -> None: + if self._status not in self._final_states: + task = await self._client._task_safe(self._task_id, self._responses) + self._status = task["status"] + self._data = task["data"] + logger.info(f"Task {self._task_id} has status {self._status}") + if not self._object_storage_data: + if self._status == "111": + self._object_storage_data = task["data"]["msg"] + elif self._status == "117": + self._object_storage_data = task["data"] + + @property + async def status(self) -> str: + """Returns status of the task that is associated with this transfer. + + :calls: GET `/tasks/{taskid}` + """ + await self._update() + return self._status # type: ignore + + @property + async def in_progress(self) -> bool: + """Returns `False` when the transfer has been completed (succesfully or with errors), otherwise `True`. + + :calls: GET `/tasks/{taskid}` + """ + await self._update() + return self._status not in self._final_states + + @property + async def data(self) -> Optional[dict]: + """Returns the task information from the latest response. + + :calls: GET `/tasks/{taskid}` + """ + await self._update() + return self._data + + @property + async def object_storage_data(self): + """Returns the necessary information for the external transfer. + The call is blocking and in cases of large file transfers it might take a long time. + + :calls: GET `/tasks/{taskid}` + :rtype: dictionary or string + """ + if not self._object_storage_data: + await self._update() + + while not self._object_storage_data: + # No need for extra sleeping here, since the async client handles + # the rate of requests anyway + await self._update() + + return self._object_storage_data + + +class AsyncExternalUpload(AsyncExternalStorage): + """ + This class handles the external upload from a file. + + Tracks the progress of the upload through the status of the associated task. + Final states: *114* and *115*. + + +--------+--------------------------------------------------------------------+ + | Status | Description | + +========+====================================================================+ + | 110 | Waiting for Form URL from Object Storage to be retrieved | + +--------+--------------------------------------------------------------------+ + | 111 | Form URL from Object Storage received | + +--------+--------------------------------------------------------------------+ + | 112 | Object Storage confirms that upload to Object Storage has finished | + +--------+--------------------------------------------------------------------+ + | 113 | Download from Object Storage to server has started | + +--------+--------------------------------------------------------------------+ + | 114 | Download from Object Storage to server has finished | + +--------+--------------------------------------------------------------------+ + | 115 | Download from Object Storage error | + +--------+--------------------------------------------------------------------+ + + :param client: FirecREST client associated with the transfer + :param task_id: FirecrREST task associated with the transfer + """ + + def __init__( + self, + client: AsyncFirecrest, + task_id: str, + previous_responses: Optional[List[requests.Response]] = None, + ) -> None: + previous_responses = [] if previous_responses is None else previous_responses + super().__init__(client, task_id, previous_responses) + self._final_states = {"114", "115"} + logger.info(f"Creating ExternalUpload object for task {task_id}") + + async def finish_upload(self) -> None: + """Finish the upload process. + This call will upload the file to the staging area. + Check with the method `status` or `in_progress` to see the status of the transfer. + The transfer from the staging area to the systems's filesystem can take several seconds to start to start. + """ + c = (await self.object_storage_data)["command"] # typer: ignore + # LOCAL FIX FOR MAC + # c = c.replace("192.168.220.19", "localhost") + logger.info(f"Uploading the file to the staging area with the command: {c}") + proc = await asyncio.create_subprocess_shell( + c, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + ) + + _, stderr = await proc.communicate() + if proc.returncode != 0: + exc = Exception( + f"Failed to finish upload with error: {stderr.decode('utf-8')}" + ) + logger.critical(exc) + raise exc + + +class AsyncExternalDownload(AsyncExternalStorage): + """ + This class handles the external download from a file. + + Tracks the progress of the download through the status of the associated task. + Final states: *117* and *118*. + + +--------+--------------------------------------------------------------------+ + | Status | Description | + +========+====================================================================+ + | 116 | Started upload from filesystem to Object Storage | + +--------+--------------------------------------------------------------------+ + | 117 | Upload from filesystem to Object Storage has finished successfully | + +--------+--------------------------------------------------------------------+ + | 118 | Upload from filesystem to Object Storage has finished with errors | + +--------+--------------------------------------------------------------------+ + + :param client: FirecREST client associated with the transfer + :param task_id: FirecrREST task associated with the transfer + """ + + def __init__( + self, + client: AsyncFirecrest, + task_id: str, + previous_responses: Optional[List[requests.Response]] = None, + ) -> None: + previous_responses = [] if previous_responses is None else previous_responses + super().__init__(client, task_id, previous_responses) + self._final_states = {"117", "118"} + logger.info(f"Creating ExternalDownload object for task {task_id}") + + async def invalidate_object_storage_link(self) -> None: + """Invalidate the temporary URL for downloading. + + :calls: POST `/storage/xfer-external/invalidate` + """ + await self._client._invalidate(self._task_id) + + @property + async def object_storage_link(self) -> str: + """Get the direct download url for the file. The response from the FirecREST api + changed after version 1.13.0, so make sure to set to older version, if you are + using an older deployment. + + :calls: GET `/tasks/{taskid}` + """ + if self._client._api_version > Version("1.13.0"): + return await self.object_storage_data["url"] + else: + return await self.object_storage_data + + async def finish_download( + self, target_path: str | pathlib.Path | BufferedWriter + ) -> None: + """Finish the download process. + + :param target_path: the local path to save the file + """ + url = await self.object_storage_data + logger.info(f"Downloading the file from {url} and saving to {target_path}") + # LOCAL FIX FOR MAC + # url = url.replace("192.168.220.19", "localhost") + context: ContextManager[BufferedWriter] = ( + open(target_path, "wb") # type: ignore + if isinstance(target_path, str) or isinstance(target_path, pathlib.Path) + else nullcontext(target_path) + ) + with urllib.request.urlopen(url) as response, context as out_file: + shutil.copyfileobj(response, out_file) diff --git a/firecrest/BasicClient.py b/firecrest/BasicClient.py index 2553f9f..7a930d3 100644 --- a/firecrest/BasicClient.py +++ b/firecrest/BasicClient.py @@ -11,12 +11,8 @@ import logging import pathlib import requests -import shlex -import shutil -import subprocess import sys import time -import urllib.request from contextlib import nullcontext from io import BufferedWriter, BytesIO @@ -26,6 +22,7 @@ import firecrest.FirecrestException as fe import firecrest.types as t +from firecrest.ExternalStorage import ExternalUpload, ExternalDownload if sys.version_info >= (3, 8): from typing import Literal @@ -47,232 +44,6 @@ def handle_response(response): print("-") -class ExternalStorage: - """External storage object. - """ - - _final_states: set[str] - - def __init__( - self, - client: Firecrest, - task_id: str, - previous_responses: Optional[List[requests.Response]] = None, - ) -> None: - previous_responses = [] if previous_responses is None else previous_responses - self._client = client - self._task_id = task_id - self._in_progress = True - self._status: Optional[str] = None - self._data = None - self._object_storage_data = None - self._sleep_time = itertools.cycle([1, 5, 10]) - self._responses = previous_responses - - @property - def client(self) -> Firecrest: - """Returns the client that will be used to get information for the task. - """ - return self._client - - @property - def task_id(self) -> str: - """Returns the FirecREST task ID that is associated with this transfer. - """ - return self._task_id - - def _update(self) -> None: - if self._status not in self._final_states: - task = self._client._task_safe(self._task_id, self._responses) - self._status = task["status"] - self._data = task["data"] - logger.info(f"Task {self._task_id} has status {self._status}") - if not self._object_storage_data: - if self._status == "111": - self._object_storage_data = task["data"]["msg"] - elif self._status == "117": - self._object_storage_data = task["data"] - - @property - def status(self) -> str: - """Returns status of the task that is associated with this transfer. - - :calls: GET `/tasks/{taskid}` - """ - self._update() - return self._status # type: ignore - - @property - def in_progress(self) -> bool: - """Returns `False` when the transfer has been completed (succesfully or with errors), otherwise `True`. - - :calls: GET `/tasks/{taskid}` - """ - self._update() - return self._status not in self._final_states - - @property - def data(self) -> Optional[dict]: - """Returns the task information from the latest response. - - :calls: GET `/tasks/{taskid}` - """ - self._update() - return self._data - - @property - def object_storage_data(self): - """Returns the necessary information for the external transfer. - The call is blocking and in cases of large file transfers it might take a long time. - - :calls: GET `/tasks/{taskid}` - :rtype: dictionary or string - """ - if not self._object_storage_data: - self._update() - - while not self._object_storage_data: - t = next(self._sleep_time) - logger.info(f"Sleeping for {t} sec") - time.sleep(t) - self._update() - - return self._object_storage_data - - -class ExternalUpload(ExternalStorage): - """ - This class handles the external upload from a file. - - Tracks the progress of the upload through the status of the associated task. - Final states: *114* and *115*. - - +--------+--------------------------------------------------------------------+ - | Status | Description | - +========+====================================================================+ - | 110 | Waiting for Form URL from Object Storage to be retrieved | - +--------+--------------------------------------------------------------------+ - | 111 | Form URL from Object Storage received | - +--------+--------------------------------------------------------------------+ - | 112 | Object Storage confirms that upload to Object Storage has finished | - +--------+--------------------------------------------------------------------+ - | 113 | Download from Object Storage to server has started | - +--------+--------------------------------------------------------------------+ - | 114 | Download from Object Storage to server has finished | - +--------+--------------------------------------------------------------------+ - | 115 | Download from Object Storage error | - +--------+--------------------------------------------------------------------+ - - :param client: FirecREST client associated with the transfer - :param task_id: FirecrREST task associated with the transfer - """ - - def __init__( - self, - client: Firecrest, - task_id: str, - previous_responses: Optional[List[requests.Response]] = None, - ) -> None: - previous_responses = [] if previous_responses is None else previous_responses - super().__init__(client, task_id, previous_responses) - self._final_states = {"114", "115"} - logger.info(f"Creating ExternalUpload object for task {task_id}") - - def finish_upload(self) -> None: - """Finish the upload process. - This call will upload the file to the staging area. - Check with the method `status` or `in_progress` to see the status of the transfer. - The transfer from the staging area to the systems's filesystem can take several seconds to start to start. - """ - c = self.object_storage_data["command"] # typer: ignore - # LOCAL FIX FOR MAC - # c = c.replace("192.168.220.19", "localhost") - logger.info(f"Uploading the file to the staging area with the command: {c}") - command = subprocess.run( - shlex.split(c), stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) - if command.returncode != 0: - exc = Exception( - f"Failed to finish upload with error: {command.stderr.decode('utf-8')}" - ) - logger.critical(exc) - raise exc - - -class ExternalDownload(ExternalStorage): - """ - This class handles the external download from a file. - - Tracks the progress of the download through the status of the associated task. - Final states: *117* and *118*. - - +--------+--------------------------------------------------------------------+ - | Status | Description | - +========+====================================================================+ - | 116 | Started upload from filesystem to Object Storage | - +--------+--------------------------------------------------------------------+ - | 117 | Upload from filesystem to Object Storage has finished successfully | - +--------+--------------------------------------------------------------------+ - | 118 | Upload from filesystem to Object Storage has finished with errors | - +--------+--------------------------------------------------------------------+ - - :param client: FirecREST client associated with the transfer - :param task_id: FirecrREST task associated with the transfer - """ - - def __init__( - self, - client: Firecrest, - task_id: str, - previous_responses: Optional[List[requests.Response]] = None, - ) -> None: - previous_responses = [] if previous_responses is None else previous_responses - super().__init__(client, task_id, previous_responses) - self._final_states = {"117", "118"} - logger.info(f"Creating ExternalDownload object for task {task_id}") - - def invalidate_object_storage_link(self) -> None: - """Invalidate the temporary URL for downloading. - - :calls: POST `/storage/xfer-external/invalidate` - """ - self._client._invalidate(self._task_id) - - @property - def object_storage_link(self) -> str: - """Get the direct download url for the file. The response from the FirecREST api - changed after version 1.13.0, so make sure to set to older version, if you are - using an older deployment. - - :calls: GET `/tasks/{taskid}` - """ - if self._client._api_version > Version("1.13.0"): - return self.object_storage_data["url"] - else: - return self.object_storage_data - - def finish_download(self, target_path: str | pathlib.Path | BufferedWriter) -> None: - """Finish the download process. The response from the FirecREST api changed after - version 1.13.0, so make sure to set to older version, if you are using an older - deployment. - - :param target_path: the local path to save the file - - :calls: GET `/tasks/{taskid}` - """ - url = self.object_storage_link - logger.info(f"Downloading the file from {url} and saving to {target_path}") - # LOCAL FIX FOR MAC - # url = url.replace("192.168.220.19", "localhost") - context: ContextManager[BufferedWriter] = ( - open(target_path, "wb") # type: ignore - if isinstance(target_path, str) or isinstance(target_path, pathlib.Path) - else nullcontext(target_path) - ) - with urllib.request.urlopen(url) as response, context as out_file: - shutil.copyfileobj(response, out_file) - - class Firecrest: """ This is the basic class you instantiate to access the FirecREST API v1. diff --git a/firecrest/ExternalStorage.py b/firecrest/ExternalStorage.py new file mode 100644 index 0000000..96471db --- /dev/null +++ b/firecrest/ExternalStorage.py @@ -0,0 +1,260 @@ +# +# Copyright (c) 2019-2023, ETH Zurich. All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +# +from __future__ import annotations + +from io import BufferedWriter +import itertools +import logging +import pathlib +import requests +import shlex +import shutil +import subprocess +import sys +import time +from typing import ContextManager, Optional, List, TYPE_CHECKING +import urllib.request +from packaging.version import Version + +if TYPE_CHECKING: + from firecrest.BasicClient import Firecrest + +from contextlib import nullcontext +from requests.compat import json # type: ignore + +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +logger = logging.getLogger(__name__) + + +class ExternalStorage: + """External storage object. + """ + + _final_states: set[str] + + def __init__( + self, + client: Firecrest, + task_id: str, + previous_responses: Optional[List[requests.Response]] = None, + ) -> None: + previous_responses = [] if previous_responses is None else previous_responses + self._client = client + self._task_id = task_id + self._in_progress = True + self._status: Optional[str] = None + self._data = None + self._object_storage_data = None + self._sleep_time = itertools.cycle([1, 5, 10]) + self._responses = previous_responses + + @property + def client(self) -> Firecrest: + """Returns the client that will be used to get information for the task. + """ + return self._client + + @property + def task_id(self) -> str: + """Returns the FirecREST task ID that is associated with this transfer. + """ + return self._task_id + + def _update(self) -> None: + if self._status not in self._final_states: + task = self._client._task_safe(self._task_id, self._responses) + self._status = task["status"] + self._data = task["data"] + logger.info(f"Task {self._task_id} has status {self._status}") + if not self._object_storage_data: + if self._status == "111": + self._object_storage_data = task["data"]["msg"] + elif self._status == "117": + self._object_storage_data = task["data"] + + @property + def status(self) -> str: + """Returns status of the task that is associated with this transfer. + + :calls: GET `/tasks/{taskid}` + """ + self._update() + return self._status # type: ignore + + @property + def in_progress(self) -> bool: + """Returns `False` when the transfer has been completed (succesfully or with errors), otherwise `True`. + + :calls: GET `/tasks/{taskid}` + """ + self._update() + return self._status not in self._final_states + + @property + def data(self) -> Optional[dict]: + """Returns the task information from the latest response. + + :calls: GET `/tasks/{taskid}` + """ + self._update() + return self._data + + @property + def object_storage_data(self): + """Returns the necessary information for the external transfer. + The call is blocking and in cases of large file transfers it might take a long time. + + :calls: GET `/tasks/{taskid}` + :rtype: dictionary or string + """ + if not self._object_storage_data: + self._update() + + while not self._object_storage_data: + t = next(self._sleep_time) + logger.info(f"Sleeping for {t} sec") + time.sleep(t) + self._update() + + return self._object_storage_data + + +class ExternalUpload(ExternalStorage): + """ + This class handles the external upload from a file. + + Tracks the progress of the upload through the status of the associated task. + Final states: *114* and *115*. + + +--------+--------------------------------------------------------------------+ + | Status | Description | + +========+====================================================================+ + | 110 | Waiting for Form URL from Object Storage to be retrieved | + +--------+--------------------------------------------------------------------+ + | 111 | Form URL from Object Storage received | + +--------+--------------------------------------------------------------------+ + | 112 | Object Storage confirms that upload to Object Storage has finished | + +--------+--------------------------------------------------------------------+ + | 113 | Download from Object Storage to server has started | + +--------+--------------------------------------------------------------------+ + | 114 | Download from Object Storage to server has finished | + +--------+--------------------------------------------------------------------+ + | 115 | Download from Object Storage error | + +--------+--------------------------------------------------------------------+ + + :param client: FirecREST client associated with the transfer + :param task_id: FirecrREST task associated with the transfer + """ + + def __init__( + self, + client: Firecrest, + task_id: str, + previous_responses: Optional[List[requests.Response]] = None, + ) -> None: + previous_responses = [] if previous_responses is None else previous_responses + super().__init__(client, task_id, previous_responses) + self._final_states = {"114", "115"} + logger.info(f"Creating ExternalUpload object for task {task_id}") + + def finish_upload(self) -> None: + """Finish the upload process. + This call will upload the file to the staging area. + Check with the method `status` or `in_progress` to see the status of the transfer. + The transfer from the staging area to the systems's filesystem can take several seconds to start to start. + """ + c = self.object_storage_data["command"] # typer: ignore + # LOCAL FIX FOR MAC + # c = c.replace("192.168.220.19", "localhost") + logger.info(f"Uploading the file to the staging area with the command: {c}") + command = subprocess.run( + shlex.split(c), stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + if command.returncode != 0: + exc = Exception( + f"Failed to finish upload with error: {command.stderr.decode('utf-8')}" + ) + logger.critical(exc) + raise exc + + +class ExternalDownload(ExternalStorage): + """ + This class handles the external download from a file. + + Tracks the progress of the download through the status of the associated task. + Final states: *117* and *118*. + + +--------+--------------------------------------------------------------------+ + | Status | Description | + +========+====================================================================+ + | 116 | Started upload from filesystem to Object Storage | + +--------+--------------------------------------------------------------------+ + | 117 | Upload from filesystem to Object Storage has finished successfully | + +--------+--------------------------------------------------------------------+ + | 118 | Upload from filesystem to Object Storage has finished with errors | + +--------+--------------------------------------------------------------------+ + + :param client: FirecREST client associated with the transfer + :param task_id: FirecrREST task associated with the transfer + """ + + def __init__( + self, + client: Firecrest, + task_id: str, + previous_responses: Optional[List[requests.Response]] = None, + ) -> None: + previous_responses = [] if previous_responses is None else previous_responses + super().__init__(client, task_id, previous_responses) + self._final_states = {"117", "118"} + logger.info(f"Creating ExternalDownload object for task {task_id}") + + def invalidate_object_storage_link(self) -> None: + """Invalidate the temporary URL for downloading. + + :calls: POST `/storage/xfer-external/invalidate` + """ + self._client._invalidate(self._task_id) + + @property + def object_storage_link(self) -> str: + """Get the direct download url for the file. The response from the FirecREST api + changed after version 1.13.0, so make sure to set to older version, if you are + using an older deployment. + + :calls: GET `/tasks/{taskid}` + """ + if self._client._api_version > Version("1.13.0"): + return self.object_storage_data["url"] + else: + return self.object_storage_data + + def finish_download(self, target_path: str | pathlib.Path | BufferedWriter) -> None: + """Finish the download process. The response from the FirecREST api changed after + version 1.13.0, so make sure to set to older version, if you are using an older + deployment. + + :param target_path: the local path to save the file + + :calls: GET `/tasks/{taskid}` + """ + url = self.object_storage_link + logger.info(f"Downloading the file from {url} and saving to {target_path}") + # LOCAL FIX FOR MAC + # url = url.replace("192.168.220.19", "localhost") + context: ContextManager[BufferedWriter] = ( + open(target_path, "wb") # type: ignore + if isinstance(target_path, str) or isinstance(target_path, pathlib.Path) + else nullcontext(target_path) + ) + with urllib.request.urlopen(url) as response, context as out_file: + shutil.copyfileobj(response, out_file) diff --git a/firecrest/__init__.py b/firecrest/__init__.py index 0361a8e..22102b9 100644 --- a/firecrest/__init__.py +++ b/firecrest/__init__.py @@ -19,11 +19,13 @@ ) sys.exit(1) -from firecrest.BasicClient import ( - Firecrest, - ExternalDownload, - ExternalUpload, - ExternalStorage, +from firecrest.BasicClient import Firecrest +from firecrest.AsyncClient import AsyncFirecrest +from firecrest.ExternalStorage import ExternalDownload, ExternalUpload, ExternalStorage +from firecrest.AsyncExternalStorage import ( + AsyncExternalDownload, + AsyncExternalUpload, + AsyncExternalStorage, ) from firecrest.Authorization import ClientCredentialsAuth from firecrest.FirecrestException import ( diff --git a/pyproject.toml b/pyproject.toml index a4bbc16..5e66a5f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "requests>=2.14.0", "PyJWT>=2.4.0", "typer[all]~=0.7.0", + "httpx>=0.24.0", "packaging>=23.1" ] @@ -41,11 +42,12 @@ firecrest = "firecrest.cli_script:main" [project.optional-dependencies] test = [ - "httpretty>=1.0.3", "pytest>=5.3", "flake8~=5.0", "mypy~=0.991", "types-requests~=2.28.11", + "pytest-httpserver~=1.0.6", + "pytest-asyncio>=0.21.1" ] docs = [ "sphinx>=4.0", diff --git a/tests/test_authoriation.py b/tests/test_authorisation.py similarity index 60% rename from tests/test_authoriation.py rename to tests/test_authorisation.py index 357e671..2c0dc1c 100644 --- a/tests/test_authoriation.py +++ b/tests/test_authorisation.py @@ -1,13 +1,13 @@ -import httpretty import json import pytest from context import firecrest +from werkzeug.wrappers import Response -def auth_callback(request, uri, response_headers): - client_id = request.parsed_body["client_id"][0] - client_secret = request.parsed_body["client_secret"][0] +def auth_handler(request): + client_id = request.form["client_id"] + client_secret = request.form["client_secret"] if client_id == "valid_id": if client_secret == "valid_secret": ret = { @@ -18,8 +18,8 @@ def auth_callback(request, uri, response_headers): "not-before-policy": 0, "scope": "profile firecrest email", } - return [200, response_headers, json.dumps(ret)] - if client_secret == "valid_secret_2": + ret_status = 200 + elif client_secret == "valid_secret_2": ret = { "access_token": "token_2", "expires_in": 15, @@ -28,42 +28,32 @@ def auth_callback(request, uri, response_headers): "not-before-policy": 0, "scope": "profile firecrest email", } - return [200, response_headers, json.dumps(ret)] + ret_status = 200 else: ret = { "error": "unauthorized_client", "error_description": "Invalid client secret", } - return [400, response_headers, json.dumps(ret)] + ret_status = 400 else: ret = { "error": "invalid_client", "error_description": "Invalid client credentials", } - return [400, response_headers, json.dumps(ret)] + ret_status = 400 + return Response(json.dumps(ret), status=ret_status, content_type="application/json") -@pytest.fixture(autouse=True) -def setup_callbacks(): - httpretty.enable(allow_net_connect=False, verbose=True) - httpretty.register_uri( - httpretty.POST, - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", - body=auth_callback, - ) - - yield +@pytest.fixture +def auth_server(httpserver): + httpserver.expect_request("/auth/token").respond_with_handler(auth_handler) + return httpserver - httpretty.disable() - httpretty.reset() - -def test_client_credentials_valid(): +def test_client_credentials_valid(auth_server): auth_obj = firecrest.ClientCredentialsAuth( - "valid_id", - "valid_secret", - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + "valid_id", "valid_secret", auth_server.url_for("/auth/token") ) assert auth_obj._min_token_validity == 10 assert auth_obj.get_access_token() == "VALID_TOKEN" @@ -74,7 +64,7 @@ def test_client_credentials_valid(): auth_obj = firecrest.ClientCredentialsAuth( "valid_id", "valid_secret", - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + auth_server.url_for("/auth/token"), min_token_validity=20, ) assert auth_obj.get_access_token() == "VALID_TOKEN" @@ -83,11 +73,9 @@ def test_client_credentials_valid(): assert auth_obj.get_access_token() == "token_2" -def test_client_credentials_invalid_id(): +def test_client_credentials_invalid_id(auth_server): auth_obj = firecrest.ClientCredentialsAuth( - "invalid_id", - "valid_secret", - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + "invalid_id", "valid_secret", auth_server.url_for("/auth/token") ) with pytest.raises(Exception) as exc_info: auth_obj.get_access_token() @@ -95,11 +83,9 @@ def test_client_credentials_invalid_id(): assert "Client credentials error" in str(exc_info.value) -def test_client_credentials_invalid_secret(): +def test_client_credentials_invalid_secret(auth_server): auth_obj = firecrest.ClientCredentialsAuth( - "valid_id", - "invalid_secret", - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + "valid_id", "invalid_secret", auth_server.url_for("/auth/token") ) with pytest.raises(Exception) as exc_info: auth_obj.get_access_token() diff --git a/tests/test_compute.py b/tests/test_compute.py index c9e713f..fcbf858 100644 --- a/tests/test_compute.py +++ b/tests/test_compute.py @@ -1,47 +1,48 @@ import common -import httpretty import json import pytest import re -import test_authoriation as auth +import test_authorisation as auth from context import firecrest from firecrest import __app_name__, __version__, cli from typer.testing import CliRunner +from werkzeug.wrappers import Response +from werkzeug.wrappers import Request runner = CliRunner() @pytest.fixture -def valid_client(): +def valid_client(fc_server): class ValidAuthorization: def get_access_token(self): return "VALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def valid_credentials(): +def valid_credentials(fc_server, auth_server): return [ - "--firecrest-url=http://firecrest.cscs.ch", + f"--firecrest-url={fc_server.url_for('/')}", "--client-id=valid_id", "--client-secret=valid_secret", - "--token-url=https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + f"--token-url={auth_server.url_for('/auth/token')}", ] @pytest.fixture -def invalid_client(): +def invalid_client(fc_server): class InvalidAuthorization: def get_access_token(self): return "INVALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=InvalidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() ) @@ -63,23 +64,30 @@ def non_slurm_script(tmp_path): return script -httpretty.enable(allow_net_connect=False, verbose=True) - - -def submit_path_callback(request, uri, response_headers): +def submit_path_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Failed to submit job", "error": "Machine does not exist"}', - ] - - target_path = request.parsed_body["targetPath"][0] - account = request.parsed_body.get("account", [None])[0] + return Response( + json.dumps( + { + "description": "Failed to submit job", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + target_path = request.form["targetPath"] + account = request.form.get("account", None) + extra_headers = None if target_path == "/path/to/workdir/script.sh": if account is None: ret = { @@ -103,29 +111,43 @@ def submit_path_callback(request, uri, response_headers): } status_code = 201 else: - response_headers["X-Invalid-Path"] = f"{target_path} is an invalid path." + extra_headers = {"X-Invalid-Path": f"{target_path} is an invalid path."} ret = {"description": "Failed to submit job"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def submit_upload_callback(request, uri, response_headers): +def submit_upload_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Failed to submit job", "error": "Machine does not exist"}', - ] - - # I couldn't find a better way to get the params from the request - if b'form-data; name="file"; filename="script.sh"' in request.body: - if b"#!/bin/bash -l\n" in request.body: - if b"proj" in request.body: + return Response( + json.dumps( + { + "description": "Failed to submit job", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + extra_headers = None + if request.files["file"].filename == "script.sh": + if b"#!/bin/bash -l\n" in request.files["file"].read(): + if request.form.get("account", None) == "proj": ret = { "success": "Task created", "task_id": "submit_upload_job_id_proj_account", @@ -146,26 +168,40 @@ def submit_upload_callback(request, uri, response_headers): } status_code = 201 else: - response_headers["X-Invalid-Path"] = f"path is an invalid path." + extra_headers = {"X-Invalid-Path": f"path is an invalid path."} ret = {"description": "Failed to submit job"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def queue_callback(request, uri, response_headers): +def queue_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{ "description": "Failed to retrieve jobs information", "error": "Machine does not exists"}', - ] - - jobs = request.querystring.get("jobs", [""])[0].split(",") + return Response( + json.dumps( + { + "description": "Failed to retrieve jobs information", + "error": "Machine does not exists", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + jobs = request.args.get("jobs", "").split(",") if jobs == [""]: ret = { "success": "Task created", @@ -194,22 +230,33 @@ def queue_callback(request, uri, response_headers): } status_code = 200 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), status=status_code, content_type="application/json" + ) -def sacct_callback(request, uri, response_headers): +def sacct_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Failed to retrieve account information", "error": "Machine does not exist"}', - ] - - jobs = request.querystring.get("jobs", [""])[0].split(",") + return Response( + json.dumps( + { + "description": "Failed to retrieve account information", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + jobs = request.args.get("jobs", "").split(",") if jobs == [""]: ret = { "success": "Task created", @@ -239,21 +286,33 @@ def sacct_callback(request, uri, response_headers): } status_code = 200 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), status=status_code, content_type="application/json" + ) -def cancel_callback(request, uri, response_headers): +def cancel_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Failed to delete job", "error": "Machine does not exist"}', - ] + return Response( + json.dumps( + { + "description": "Failed to delete job", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + uri = request.url jobid = uri.split("/")[-1] if jobid == "35360071": ret = { @@ -277,7 +336,9 @@ def cancel_callback(request, uri, response_headers): } status_code = 200 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), status=status_code, content_type="application/json" + ) # Global variables for tasks @@ -293,9 +354,13 @@ def cancel_callback(request, uri, response_headers): cancel_result = 1 -def tasks_callback(request, uri, response_headers): +def tasks_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) global submit_path_retry global submit_upload_retry @@ -303,6 +368,7 @@ def tasks_callback(request, uri, response_headers): global queue_retry global cancel_retry + uri = request.url taskid = uri.split("/")[-1] if taskid == "tasks": # TODO: return all tasks @@ -745,55 +811,44 @@ def tasks_callback(request, uri, response_headers): } status_code = 200 - return [status_code, response_headers, json.dumps(ret)] - + return Response( + json.dumps(ret), status=status_code, content_type="application/json" + ) -@pytest.fixture(autouse=True) -def setup_callbacks(): - httpretty.enable(allow_net_connect=False, verbose=True) - httpretty.register_uri( - httpretty.POST, - "http://firecrest.cscs.ch/compute/jobs/path", - body=submit_path_callback, +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request("/compute/jobs/path", method="POST").respond_with_handler( + submit_path_handler ) - httpretty.register_uri( - httpretty.POST, - "http://firecrest.cscs.ch/compute/jobs/upload", - body=submit_upload_callback, - ) + httpserver.expect_request( + "/compute/jobs/upload", method="POST" + ).respond_with_handler(submit_upload_handler) - httpretty.register_uri( - httpretty.GET, "http://firecrest.cscs.ch/compute/acct", body=sacct_callback + httpserver.expect_request("/compute/acct", method="GET").respond_with_handler( + sacct_handler ) - httpretty.register_uri( - httpretty.GET, "http://firecrest.cscs.ch/compute/jobs", body=queue_callback + httpserver.expect_request("/compute/jobs", method="GET").respond_with_handler( + queue_handler ) - httpretty.register_uri( - httpretty.DELETE, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/compute\/jobs.*"), - body=cancel_callback, - ) + httpserver.expect_request( + re.compile("^/compute/jobs.*"), method="DELETE" + ).respond_with_handler(cancel_handler) - httpretty.register_uri( - httpretty.GET, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/tasks.*"), - body=tasks_callback, - ) + httpserver.expect_request( + re.compile("^/tasks/.*"), method="GET" + ).respond_with_handler(tasks_handler) - httpretty.register_uri( - httpretty.POST, - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", - body=auth.auth_callback, - ) + return httpserver - yield - httpretty.disable() - httpretty.reset() +@pytest.fixture +def auth_server(httpserver): + httpserver.expect_request("/auth/token").respond_with_handler(auth.auth_handler) + return httpserver def test_submit_remote(valid_client): diff --git a/tests/test_compute_async.py b/tests/test_compute_async.py new file mode 100644 index 0000000..1c577a8 --- /dev/null +++ b/tests/test_compute_async.py @@ -0,0 +1,405 @@ +import pytest +import re +import test_compute as basic_compute + +from context import firecrest +from firecrest import __app_name__, __version__ + + +@pytest.fixture +def valid_client(fc_server): + class ValidAuthorization: + def get_access_token(self): + return "VALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def invalid_client(fc_server): + class InvalidAuthorization: + def get_access_token(self): + return "INVALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def slurm_script(tmp_path): + tmp_dir = tmp_path / "script_dir" + tmp_dir.mkdir() + script = tmp_dir / "script.sh" + script.write_text("#!/bin/bash -l\n# Slurm job script\n") + return script + + +@pytest.fixture +def non_slurm_script(tmp_path): + tmp_dir = tmp_path / "script_dir" + tmp_dir.mkdir() + script = tmp_dir / "script.sh" + script.write_text("non slurm script\n") + return script + + +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request("/compute/jobs/path", method="POST").respond_with_handler( + basic_compute.submit_path_handler + ) + + httpserver.expect_request( + "/compute/jobs/upload", method="POST" + ).respond_with_handler(basic_compute.submit_upload_handler) + + httpserver.expect_request("/compute/acct", method="GET").respond_with_handler( + basic_compute.sacct_handler + ) + + httpserver.expect_request("/compute/jobs", method="GET").respond_with_handler( + basic_compute.queue_handler + ) + + httpserver.expect_request( + re.compile("^/compute/jobs.*"), method="DELETE" + ).respond_with_handler(basic_compute.cancel_handler) + + httpserver.expect_request( + re.compile("^/tasks/.*"), method="GET" + ).respond_with_handler(basic_compute.tasks_handler) + + return httpserver + + +@pytest.mark.asyncio +async def test_submit_remote(valid_client): + global submit_path_retry + submit_path_retry = 0 + assert await valid_client.submit( + machine="cluster1", job_script="/path/to/workdir/script.sh", local_file=False + ) == { + "job_data_err": "", + "job_data_out": "", + "job_file": "/path/to/workdir/script.sh", + "job_file_err": "/path/to/workdir/slurm-35335405.out", + "job_file_out": "/path/to/workdir/slurm-35335405.out", + "jobid": 35335405, + "result": "Job submitted", + } + submit_path_retry = 0 + assert await valid_client.submit( + machine="cluster1", + job_script="/path/to/workdir/script.sh", + local_file=False, + account="proj", + ) == { + "job_data_err": "", + "job_data_out": "", + "job_file": "/path/to/workdir/script.sh", + "job_file_err": "/path/to/workdir/slurm-35335405.out", + "job_file_out": "/path/to/workdir/slurm-35335405.out", + "jobid": 35335406, + "result": "Job submitted", + } + + +@pytest.mark.asyncio +async def test_submit_local(valid_client, slurm_script): + # Test submission for local script + global submit_upload_retry + submit_upload_retry = 0 + assert await valid_client.submit( + machine="cluster1", job_script=slurm_script, local_file=True + ) == { + "job_data_err": "", + "job_data_out": "", + "job_file": "/path/to/firecrest/submit_upload_job_id_default_account/script.sh", + "job_file_err": "/path/to/firecrest/submit_upload_job_id_default_account/slurm-35342667.out", + "job_file_out": "/path/to/firecrest/submit_upload_job_id_default_account/slurm-35342667.out", + "jobid": 35342667, + "result": "Job submitted", + } + submit_upload_retry = 0 + assert await valid_client.submit( + machine="cluster1", job_script=slurm_script, local_file=True, account="proj" + ) == { + "job_data_err": "", + "job_data_out": "", + "job_file": "/path/to/firecrest/submit_upload_job_id_proj_account/script.sh", + "job_file_err": "/path/to/firecrest/submit_upload_job_id_proj_account/slurm-35342667.out", + "job_file_out": "/path/to/firecrest/submit_upload_job_id_proj_account/slurm-35342667.out", + "jobid": 35342668, + "result": "Job submitted", + } + + +@pytest.mark.asyncio +async def test_submit_invalid_arguments(valid_client, non_slurm_script): + with pytest.raises(firecrest.HeaderException): + await valid_client.submit( + machine="cluster1", + job_script="/path/to/non/existent/file", + local_file=False, + ) + + global submit_path_retry + submit_path_retry = 0 + with pytest.raises(firecrest.FirecrestException): + await valid_client.submit( + machine="cluster1", + job_script="/path/to/non/slurm/file.sh", + local_file=False, + ) + + global submit_upload_retry + submit_upload_retry = 0 + + with pytest.raises(firecrest.FirecrestException): + await valid_client.submit( + machine="cluster1", job_script=non_slurm_script, local_file=True + ) + + +@pytest.mark.asyncio +async def test_submit_invalid_machine(valid_client, slurm_script): + with pytest.raises(firecrest.HeaderException): + await valid_client.submit( + machine="cluster2", job_script="/path/to/file", local_file=False + ) + + with pytest.raises(firecrest.HeaderException): + await valid_client.submit( + machine="cluster2", job_script=slurm_script, local_file=True + ) + + +@pytest.mark.asyncio +async def test_submit_invalid_client(invalid_client, slurm_script): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.submit( + machine="cluster1", job_script="/path/to/file", local_file=False + ) + + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.submit( + machine="cluster1", job_script=slurm_script, local_file=True + ) + + +@pytest.mark.asyncio +async def test_poll(valid_client): + global acct_retry + acct_retry = 0 + assert await valid_client.poll( + machine="cluster1", + jobs=[352, 2, "334"], + start_time="starttime", + end_time="endtime", + ) == [ + { + "jobid": "352", + "name": "firecrest_job_test", + "nodelist": "nid0[6227-6229]", + "nodes": "3", + "partition": "normal", + "start_time": "2021-11-29T16:31:07", + "state": "COMPLETED", + "time": "00:48:00", + "time_left": "2021-11-29T16:31:47", + "user": "username", + }, + { + "jobid": "334", + "name": "firecrest_job_test2", + "nodelist": "nid02401", + "nodes": "1", + "partition": "normal", + "start_time": "2021-11-29T16:31:07", + "state": "COMPLETED", + "time": "00:17:12", + "time_left": "2021-11-29T16:31:50", + "user": "username", + }, + ] + acct_retry = 0 + assert await valid_client.poll(machine="cluster1", jobs=[]) == [ + { + "jobid": "352", + "name": "firecrest_job_test", + "nodelist": "nid0[6227-6229]", + "nodes": "3", + "partition": "normal", + "start_time": "2021-11-29T16:31:07", + "state": "COMPLETED", + "time": "00:48:00", + "time_left": "2021-11-29T16:31:47", + "user": "username", + } + ] + assert await valid_client.poll(machine="cluster1", jobs=["empty"]) == [] + + +@pytest.mark.asyncio +async def test_poll_invalid_arguments(valid_client): + global acct_retry + acct_retry = 0 + + with pytest.raises(firecrest.FirecrestException): + await valid_client.poll(machine="cluster1", jobs=["l"]) + + +@pytest.mark.asyncio +async def test_poll_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.poll(machine="cluster2", jobs=[]) + + +@pytest.mark.asyncio +async def test_poll_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.poll(machine="cluster1", jobs=[]) + + +@pytest.mark.asyncio +async def test_poll_active(valid_client): + global queue_retry + queue_retry = 0 + assert await valid_client.poll_active(machine="cluster1", jobs=[352, 2, "334"]) == [ + { + "job_data_err": "", + "job_data_out": "", + "job_file": "(null)", + "job_file_err": "stderr-file-not-found", + "job_file_out": "stdout-file-not-found", + "jobid": "352", + "name": "interactive", + "nodelist": "nid02357", + "nodes": "1", + "partition": "debug", + "start_time": "6:38", + "state": "RUNNING", + "time": "2022-03-10T10:11:34", + "time_left": "23:22", + "user": "username", + } + ] + queue_retry = 0 + assert await valid_client.poll_active(machine="cluster1", jobs=[]) == [ + { + "job_data_err": "", + "job_data_out": "", + "job_file": "(null)", + "job_file_err": "stderr-file-not-found", + "job_file_out": "stdout-file-not-found", + "jobid": "352", + "name": "interactive", + "nodelist": "nid02357", + "nodes": "1", + "partition": "debug", + "start_time": "6:38", + "state": "RUNNING", + "time": "2022-03-10T10:11:34", + "time_left": "23:22", + "user": "username", + }, + { + "job_data_err": "", + "job_data_out": "", + "job_file": "(null)", + "job_file_err": "stderr-file-not-found", + "job_file_out": "stdout-file-not-found", + "jobid": "356", + "name": "interactive", + "nodelist": "nid02351", + "nodes": "1", + "partition": "debug", + "start_time": "6:38", + "state": "RUNNING", + "time": "2022-03-10T10:11:34", + "time_left": "23:22", + "user": "username", + }, + ] + + +@pytest.mark.asyncio +async def test_poll_active_invalid_arguments(valid_client): + global queue_retry + queue_retry = 0 + + with pytest.raises(firecrest.FirecrestException): + await valid_client.poll_active(machine="cluster1", jobs=["l"]) + + queue_retry = 0 + with pytest.raises(firecrest.FirecrestException): + # We assume that jobid is too old and is rejected by squeue + await valid_client.poll_active(machine="cluster1", jobs=["4"]) + + +@pytest.mark.asyncio +async def test_poll_active_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.poll_active(machine="cluster2", jobs=[]) + + +@pytest.mark.asyncio +async def test_poll_active_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.poll_active(machine="cluster1", jobs=[]) + + +@pytest.mark.asyncio +async def test_cancel(valid_client): + global cancel_retry + cancel_retry = 0 + # Make sure this doesn't raise an error + await valid_client.cancel(machine="cluster1", job_id=35360071) + + +@pytest.mark.asyncio +async def test_cancel_invalid_arguments(valid_client): + global cancel_retry + cancel_retry = 0 + with pytest.raises(firecrest.FirecrestException): + await valid_client.cancel(machine="cluster1", job_id="k") + + cancel_retry = 0 + with pytest.raises(firecrest.FirecrestException): + # Jobid 35360072 is from a different user + await valid_client.cancel(machine="cluster1", job_id="35360072") + + +@pytest.mark.asyncio +async def test_cancel_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.cancel(machine="cluster2", job_id=35360071) + + +@pytest.mark.asyncio +async def test_cancel_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.cancel(machine="cluster1", job_id=35360071) diff --git a/tests/test_extras.py b/tests/test_extras.py index 334f182..fdf5574 100644 --- a/tests/test_extras.py +++ b/tests/test_extras.py @@ -1,12 +1,13 @@ import common -import httpretty import json import pytest import re -import test_authoriation as auth +import test_authorisation as auth from context import firecrest from typer.testing import CliRunner +from werkzeug.wrappers import Response +from werkzeug.wrappers import Request from firecrest import __app_name__, __version__, cli @@ -21,7 +22,7 @@ def test_cli_version(): @pytest.fixture -def client1(): +def client1(fc_server): class ValidAuthorization: def get_access_token(self): # This token was created in https://jwt.io/ with payload: @@ -44,12 +45,12 @@ def get_access_token(self): return "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJyZWFsbV9hY2Nlc3MiOnsicm9sZXMiOlsiZmlyZWNyZXN0LXNhIl19LCJyZXNvdXJjZV9hY2Nlc3MiOnsiYm9iLWNsaWVudCI6eyJyb2xlcyI6WyJib2IiXX19LCJjbGllbnRJZCI6ImJvYi1jbGllbnQiLCJwcmVmZXJyZWRfdXNlcm5hbWUiOiJzZXJ2aWNlLWFjY291bnQtYm9iLWNsaWVudCJ9.XfCXDclEBh7faQrOF2piYdnb7c3AUiCxDesTkNSwpSY" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def client2(): +def client2(fc_server): class ValidAuthorization: def get_access_token(self): # This token was created in https://jwt.io/ with payload: @@ -64,12 +65,12 @@ def get_access_token(self): return "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJyZWFsbV9hY2Nlc3MiOnsicm9sZXMiOlsib3RoZXItcm9sZSJdfSwicHJlZmVycmVkX3VzZXJuYW1lIjoiYWxpY2UifQ.dpo1_F9jkV-RpNGqTaCNLbM-JPMnstDg7mQjzbwDp5g" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def client3(): +def client3(fc_server): class ValidAuthorization: def get_access_token(self): # This token was created in https://jwt.io/ with payload: @@ -79,45 +80,49 @@ def get_access_token(self): return "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJwcmVmZXJyZWRfdXNlcm5hbWUiOiJldmUifQ.SGVPDrJdy8b5jRpxcw9ILLsf8M2ljAYWxiN0A1b_1SE" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def valid_client(): +def valid_client(fc_server): class ValidAuthorization: def get_access_token(self): return "VALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def valid_credentials(): +def valid_credentials(fc_server, auth_server): return [ - "--firecrest-url=http://firecrest.cscs.ch", + f"--firecrest-url={fc_server.url_for('/')}", "--client-id=valid_id", "--client-secret=valid_secret", - "--token-url=https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + f"--token-url={auth_server.url_for('/auth/token')}", ] @pytest.fixture -def invalid_client(): - class ValidAuthorization: +def invalid_client(fc_server): + class InvalidAuthorization: def get_access_token(self): - return "INVALID TOKEN" + return "INVALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() ) -def tasks_callback(request, uri, response_headers): +def tasks_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) ret = { "tasks": { @@ -162,38 +167,35 @@ def tasks_callback(request, uri, response_headers): }, } } - if uri == "http://firecrest.cscs.ch/tasks/": - return [200, response_headers, json.dumps(ret)] - - task_id = uri.split("/")[-1] - if task_id in {"taskid_1", "taskid_2", "taskid_3"}: - ret = {"task": ret["tasks"][task_id]} - return [200, response_headers, json.dumps(ret)] - else: - ret = {"error": f"Task {task_id} does not exist"} - return [404, response_headers, json.dumps(ret)] - - -@pytest.fixture(autouse=True) -def setup_callbacks(): - httpretty.enable(allow_net_connect=False, verbose=True) - - httpretty.register_uri( - httpretty.GET, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/tasks.*"), - body=tasks_callback, + status_code = 200 + uri = request.url + if not uri.endswith("/tasks/"): + task_id = uri.split("/")[-1] + if task_id in {"taskid_1", "taskid_2", "taskid_3"}: + ret = {"task": ret["tasks"][task_id]} + status_code = 200 + else: + ret = {"error": f"Task {task_id} does not exist"} + status_code = 404 + + return Response( + json.dumps(ret), status=status_code, content_type="application/json" ) - httpretty.register_uri( - httpretty.POST, - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", - body=auth.auth_callback, - ) - yield +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request( + re.compile("^/tasks/.*"), method="GET" + ).respond_with_handler(tasks_handler) + + return httpserver + - httpretty.disable() - httpretty.reset() +@pytest.fixture +def auth_server(httpserver): + httpserver.expect_request("/auth/token").respond_with_handler(auth.auth_handler) + return httpserver def test_whoami(client1): diff --git a/tests/test_extras_async.py b/tests/test_extras_async.py new file mode 100644 index 0000000..5ccde1d --- /dev/null +++ b/tests/test_extras_async.py @@ -0,0 +1,167 @@ +import pytest +import re +import test_extras as basic_extras + +from context import firecrest + +from firecrest import __app_name__, __version__ + + +@pytest.fixture +def valid_client(fc_server): + class ValidAuthorization: + def get_access_token(self): + return "VALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def invalid_client(fc_server): + class InvalidAuthorization: + def get_access_token(self): + return "INVALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request( + re.compile("^/tasks/.*"), method="GET" + ).respond_with_handler(basic_extras.tasks_handler) + + return httpserver + + +@pytest.mark.asyncio +async def test_all_tasks(valid_client): + assert await valid_client._tasks() == { + "taskid_1": { + "created_at": "2022-08-16T07:18:54", + "data": "data", + "description": "description", + "hash_id": "taskid_1", + "last_modify": "2022-08-16T07:18:54", + "service": "storage", + "status": "114", + "task_id": "taskid_1", + "task_url": "TASK_IP/tasks/taskid_1", + "updated_at": "2022-08-16T07:18:54", + "user": "username", + }, + "taskid_2": { + "created_at": "2022-08-16T07:18:54", + "data": "data", + "description": "description", + "hash_id": "taskid_2", + "last_modify": "2022-08-16T07:18:54", + "service": "storage", + "status": "112", + "task_id": "taskid_2", + "task_url": "TASK_IP/tasks/taskid_2", + "updated_at": "2022-08-16T07:18:54", + "user": "username", + }, + "taskid_3": { + "created_at": "2022-08-16T07:18:54", + "data": "data", + "description": "description", + "hash_id": "taskid_3", + "last_modify": "2022-08-16T07:18:54", + "service": "storage", + "status": "111", + "task_id": "taskid_3", + "task_url": "TASK_IP/tasks/taskid_3", + "updated_at": "2022-08-16T07:18:54", + "user": "username", + }, + } + + +@pytest.mark.asyncio +async def test_subset_tasks(valid_client): + # "taskid_4" is not a valid id but it will be silently ignored + assert await valid_client._tasks(["taskid_1", "taskid_3", "taskid_4"]) == { + "taskid_1": { + "created_at": "2022-08-16T07:18:54", + "data": "data", + "description": "description", + "hash_id": "taskid_1", + "last_modify": "2022-08-16T07:18:54", + "service": "storage", + "status": "114", + "task_id": "taskid_1", + "task_url": "TASK_IP/tasks/taskid_1", + "updated_at": "2022-08-16T07:18:54", + "user": "username", + }, + "taskid_3": { + "created_at": "2022-08-16T07:18:54", + "data": "data", + "description": "description", + "hash_id": "taskid_3", + "last_modify": "2022-08-16T07:18:54", + "service": "storage", + "status": "111", + "task_id": "taskid_3", + "task_url": "TASK_IP/tasks/taskid_3", + "updated_at": "2022-08-16T07:18:54", + "user": "username", + }, + } + + +@pytest.mark.asyncio +async def test_one_task(valid_client): + assert await valid_client._tasks(["taskid_2"]) == { + "taskid_2": { + "created_at": "2022-08-16T07:18:54", + "data": "data", + "description": "description", + "hash_id": "taskid_2", + "last_modify": "2022-08-16T07:18:54", + "service": "storage", + "status": "112", + "task_id": "taskid_2", + "task_url": "TASK_IP/tasks/taskid_2", + "updated_at": "2022-08-16T07:18:54", + "user": "username", + } + } + + +@pytest.mark.asyncio +async def test_invalid_task(valid_client): + with pytest.raises(firecrest.FirecrestException): + await valid_client._tasks(["invalid_id"]) + + +@pytest.mark.asyncio +async def test_tasks_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client._tasks() diff --git a/tests/test_reservation.py b/tests/test_reservation.py index 865b455..a9abc38 100644 --- a/tests/test_reservation.py +++ b/tests/test_reservation.py @@ -1,119 +1,123 @@ -import httpretty import json import pytest import re -import test_authoriation as auth +import test_authorisation as auth from context import firecrest from firecrest import __app_name__, __version__, cli from typer.testing import CliRunner +from werkzeug.wrappers import Response +from werkzeug.wrappers import Request runner = CliRunner() @pytest.fixture -def valid_client(): +def valid_client(fc_server): class ValidAuthorization: def get_access_token(self): return "VALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def valid_credentials(): +def valid_credentials(fc_server, auth_server): return [ - "--firecrest-url=http://firecrest.cscs.ch", + f"--firecrest-url={fc_server.url_for('/')}", "--client-id=valid_id", "--client-secret=valid_secret", - "--token-url=https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + f"--token-url={auth_server.url_for('/auth/token')}", ] @pytest.fixture -def invalid_client(): +def invalid_client(fc_server): class InvalidAuthorization: def get_access_token(self): return "INVALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=InvalidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() ) -def all_reservations_callback(request, uri, response_headers): +def all_reservations_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) ret = {"success": []} - return [200, response_headers, json.dumps(ret)] + return Response(json.dumps(ret), status=200, content_type="application/json") -def create_reservation_callback(request, uri, response_headers): +def create_reservation_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) ret = {} - return [201, response_headers, json.dumps(ret)] + return Response(json.dumps(ret), status=201, content_type="application/json") -def update_reservation_callback(request, uri, response_headers): +def update_reservation_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) ret = {} - return [200, response_headers, json.dumps(ret)] + return Response(json.dumps(ret), status=200, content_type="application/json") -def delete_reservation_callback(request, uri, response_headers): +def delete_reservation_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) ret = {} - return [204, response_headers, json.dumps(ret)] - + return Response(json.dumps(ret), status=204, content_type="application/json") -@pytest.fixture(autouse=True) -def setup_callbacks(): - httpretty.enable(allow_net_connect=False, verbose=True) - httpretty.register_uri( - httpretty.GET, - "http://firecrest.cscs.ch/reservations", - body=all_reservations_callback, +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request("/reservations", method="GET").respond_with_handler( + all_reservations_handler ) - httpretty.register_uri( - httpretty.POST, - "http://firecrest.cscs.ch/reservations", - body=create_reservation_callback, + httpserver.expect_request("/reservations", method="POST").respond_with_handler( + create_reservation_handler ) - httpretty.register_uri( - httpretty.PUT, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/reservations\/.*"), - body=update_reservation_callback, - ) + httpserver.expect_request( + re.compile("^/reservations/.*"), method="PUT" + ).respond_with_handler(update_reservation_handler) - httpretty.register_uri( - httpretty.DELETE, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/reservations\/.*"), - body=delete_reservation_callback, - ) + httpserver.expect_request( + re.compile("^/reservations/.*"), method="DELETE" + ).respond_with_handler(delete_reservation_handler) - httpretty.register_uri( - httpretty.POST, - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", - body=auth.auth_callback, - ) + return httpserver - yield - httpretty.disable() - httpretty.reset() +@pytest.fixture +def auth_server(httpserver): + httpserver.expect_request("/auth/token").respond_with_handler(auth.auth_handler) + return httpserver def test_all_reservations(valid_client): diff --git a/tests/test_reservation_async.py b/tests/test_reservation_async.py new file mode 100644 index 0000000..205e44a --- /dev/null +++ b/tests/test_reservation_async.py @@ -0,0 +1,145 @@ +import pytest +import re +import test_reservation as basic_reservation + +from context import firecrest +from firecrest import __app_name__, __version__ + + +@pytest.fixture +def valid_client(fc_server): + class ValidAuthorization: + def get_access_token(self): + return "VALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def invalid_client(fc_server): + class InvalidAuthorization: + def get_access_token(self): + return "INVALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request("/reservations", method="GET").respond_with_handler( + basic_reservation.all_reservations_handler + ) + + httpserver.expect_request("/reservations", method="POST").respond_with_handler( + basic_reservation.create_reservation_handler + ) + + httpserver.expect_request( + re.compile("^/reservations/.*"), method="PUT" + ).respond_with_handler(basic_reservation.update_reservation_handler) + + httpserver.expect_request( + re.compile("^/reservations/.*"), method="DELETE" + ).respond_with_handler(basic_reservation.delete_reservation_handler) + + return httpserver + + +@pytest.mark.asyncio +async def test_all_reservations(valid_client): + assert await valid_client.all_reservations("cluster1") == [] + + +@pytest.mark.asyncio +async def test_all_reservations_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.all_reservations("cluster1") + + +@pytest.mark.asyncio +async def test_create_reservation(valid_client): + await valid_client.create_reservation( + "cluster1", + "reservation", + "account", + "number_of_nodes", + "node_type", + "start_time", + "end_time", + ) + + +@pytest.mark.asyncio +async def test_create_reservation_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.create_reservation( + "cluster1", + "reservation", + "account", + "number_of_nodes", + "node_type", + "start_time", + "end_time", + ) + + +@pytest.mark.asyncio +async def test_update_reservation(valid_client): + await valid_client.update_reservation( + "cluster1", + "reservation", + "account", + "number_of_nodes", + "node_type", + "start_time", + "end_time", + ) + + +@pytest.mark.asyncio +async def test_update_reservation_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.update_reservation( + "cluster1", + "reservation", + "account", + "number_of_nodes", + "node_type", + "start_time", + "end_time", + ) + + +@pytest.mark.asyncio +async def test_delete_reservation(valid_client): + await valid_client.delete_reservation("cluster1", "reservation") + + +@pytest.mark.asyncio +async def test_delete_reservation_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.delete_reservation("cluster1", "reservation") diff --git a/tests/test_status.py b/tests/test_status.py index 8fcc732..09e8d17 100644 --- a/tests/test_status.py +++ b/tests/test_status.py @@ -1,53 +1,81 @@ import common -import httpretty import json import pytest import re -import test_authoriation as auth +import test_authorisation as auth from context import firecrest from firecrest import __app_name__, __version__, cli from typer.testing import CliRunner +from werkzeug.wrappers import Response +from werkzeug.wrappers import Request runner = CliRunner() @pytest.fixture -def valid_client(): +def valid_client(fc_server): class ValidAuthorization: def get_access_token(self): return "VALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def valid_credentials(): +def valid_credentials(fc_server, auth_server): return [ - "--firecrest-url=http://firecrest.cscs.ch", + f"--firecrest-url={fc_server.url_for('/')}", "--client-id=valid_id", "--client-secret=valid_secret", - "--token-url=https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + f"--token-url={auth_server.url_for('/auth/token')}", ] @pytest.fixture -def invalid_client(): +def invalid_client(fc_server): class InvalidAuthorization: def get_access_token(self): return "INVALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=InvalidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() ) -def services_callback(request, uri, response_headers): +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request( + re.compile("^/status/services.*"), method="GET" + ).respond_with_handler(services_handler) + + httpserver.expect_request( + re.compile("^/status/systems.*"), method="GET" + ).respond_with_handler(systems_handler) + + httpserver.expect_request("/status/parameters", method="GET").respond_with_handler( + parameters_handler + ) + + return httpserver + + +@pytest.fixture +def auth_server(httpserver): + httpserver.expect_request("/auth/token").respond_with_handler(auth.auth_handler) + return httpserver + + +def services_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) ret = { "description": "List of services with status and description.", @@ -64,24 +92,28 @@ def services_callback(request, uri, response_headers): }, ], } - if uri == "http://firecrest.cscs.ch/status/services": - return [200, response_headers, json.dumps(ret)] - - service = uri.split("/")[-1] - if service == "utilities": - ret = ret["out"][0] - return [200, response_headers, json.dumps(ret)] - elif service == "compute": - ret = ret["out"][1] - return [200, response_headers, json.dumps(ret)] - else: - ret = {"description": "Service does not exists"} - return [404, response_headers, json.dumps(ret)] - - -def systems_callback(request, uri, response_headers): + ret_status = 200 + uri = request.url + if not uri.endswith("/status/services"): + service = uri.split("/")[-1] + if service == "utilities": + ret = ret["out"][0] + elif service == "compute": + ret = ret["out"][1] + else: + ret = {"description": "Service does not exists"} + ret_status = 404 + + return Response(json.dumps(ret), status=ret_status, content_type="application/json") + + +def systems_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) ret = { "description": "List of systems with status and description.", @@ -98,38 +130,44 @@ def systems_callback(request, uri, response_headers): }, ], } - if uri == "http://firecrest.cscs.ch/status/systems": - return [200, response_headers, json.dumps(ret)] - - service = uri.split("/")[-1] - if service == "cluster1": - ret = { - "description": "System information", - "out": { - "description": "System ready", - "status": "available", - "system": "cluster1", - }, - } - return [200, response_headers, json.dumps(ret)] - elif service == "cluster2": - ret = { - "description": "System information", - "out": { - "description": "System ready", - "status": "available", - "system": "cluster2", - }, - } - return [200, response_headers, json.dumps(ret)] - else: - ret = {"description": "System does not exists."} - return [404, response_headers, json.dumps(ret)] + ret_status = 200 + uri = request.url + if not uri.endswith("/status/systems"): + system = uri.split("/")[-1] + if system == "cluster1": + ret = { + "description": "System information", + "out": { + "description": "System ready", + "status": "available", + "system": "cluster1", + }, + } + ret_status = 200 + elif system == "cluster2": + ret = { + "description": "System information", + "out": { + "description": "System ready", + "status": "available", + "system": "cluster2", + }, + } + ret_status = 200 + else: + ret = {"description": "System does not exists."} + ret_status = 400 + + return Response(json.dumps(ret), status=ret_status, content_type="application/json") -def parameters_callback(request, uri, response_headers): +def parameters_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) ret = { "description": "Firecrest's parameters", @@ -154,41 +192,7 @@ def parameters_callback(request, uri, response_headers): ], }, } - return [200, response_headers, json.dumps(ret)] - - -@pytest.fixture(autouse=True) -def setup_callbacks(): - httpretty.enable(allow_net_connect=False, verbose=True) - - httpretty.register_uri( - httpretty.GET, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/status\/services.*"), - body=services_callback, - ) - - httpretty.register_uri( - httpretty.GET, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/status\/systems.*"), - body=systems_callback, - ) - - httpretty.register_uri( - httpretty.GET, - "http://firecrest.cscs.ch/status/parameters", - body=parameters_callback, - ) - - httpretty.register_uri( - httpretty.POST, - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", - body=auth.auth_callback, - ) - - yield - - httpretty.disable() - httpretty.reset() + return Response(json.dumps(ret), status=200, content_type="application/json") def test_all_services(valid_client): diff --git a/tests/test_status_async.py b/tests/test_status_async.py new file mode 100644 index 0000000..32f9658 --- /dev/null +++ b/tests/test_status_async.py @@ -0,0 +1,169 @@ +import pytest +import re +import test_status as basic_status + +from context import firecrest +from firecrest import __app_name__, __version__ + + +@pytest.fixture +def valid_client(fc_server): + class ValidAuthorization: + def get_access_token(self): + return "VALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def invalid_client(fc_server): + class InvalidAuthorization: + def get_access_token(self): + return "INVALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request( + re.compile("^/status/services.*"), method="GET" + ).respond_with_handler(basic_status.services_handler) + + httpserver.expect_request( + re.compile("^/status/systems.*"), method="GET" + ).respond_with_handler(basic_status.systems_handler) + + httpserver.expect_request("/status/parameters", method="GET").respond_with_handler( + basic_status.parameters_handler + ) + + return httpserver + + +@pytest.mark.asyncio +async def test_all_services(valid_client): + assert await valid_client.all_services() == [ + { + "description": "server up & flask running", + "service": "utilities", + "status": "available", + }, + { + "description": "server up & flask running", + "service": "compute", + "status": "available", + }, + ] + + +@pytest.mark.asyncio +async def test_all_services_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.all_services() + + +@pytest.mark.asyncio +async def test_service(valid_client): + assert await valid_client.service("utilities") == { + "description": "server up & flask running", + "service": "utilities", + "status": "available", + } + + +@pytest.mark.asyncio +async def test_invalid_service(valid_client): + with pytest.raises(firecrest.FirecrestException): + await valid_client.service("invalid_service") + + +@pytest.mark.asyncio +async def test_service_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.service("utilities") + + +@pytest.mark.asyncio +async def test_all_systems(valid_client): + assert await valid_client.all_systems() == [ + {"description": "System ready", "status": "available", "system": "cluster1"}, + {"description": "System ready", "status": "available", "system": "cluster2"}, + ] + + +@pytest.mark.asyncio +async def test_all_systems_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.all_systems() + + +@pytest.mark.asyncio +async def test_system(valid_client): + assert await valid_client.system("cluster1") == { + "description": "System ready", + "status": "available", + "system": "cluster1", + } + + +@pytest.mark.asyncio +async def test_invalid_system(valid_client): + with pytest.raises(firecrest.FirecrestException): + await valid_client.system("invalid_system") + + +@pytest.mark.asyncio +async def test_system_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.system("cluster1") + + +@pytest.mark.asyncio +async def test_parameters(valid_client): + assert await valid_client.parameters() == { + "storage": [ + {"name": "OBJECT_STORAGE", "unit": "", "value": "swift"}, + {"name": "STORAGE_TEMPURL_EXP_TIME", "unit": "seconds", "value": "2592000"}, + {"name": "STORAGE_MAX_FILE_SIZE", "unit": "MB", "value": "512000"}, + { + "name": "FILESYSTEMS", + "unit": "", + "value": [{"mounted": ["/fs1"], "system": "cluster1"}], + }, + ], + "utilities": [ + {"name": "UTILITIES_MAX_FILE_SIZE", "unit": "MB", "value": "5"}, + {"name": "UTILITIES_TIMEOUT", "unit": "seconds", "value": "5"}, + ], + } + + +@pytest.mark.asyncio +async def test_parameters_invalid(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.parameters() diff --git a/tests/test_storage.py b/tests/test_storage.py index 016541d..227c3da 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -1,63 +1,72 @@ import common -import httpretty import json import pytest import re -import test_authoriation as auth +import test_authorisation as auth from context import firecrest from firecrest import __app_name__, __version__, cli -from firecrest.BasicClient import ExternalUpload, ExternalDownload from typer.testing import CliRunner +from werkzeug.wrappers import Response +from werkzeug.wrappers import Request runner = CliRunner() @pytest.fixture -def valid_client(): +def valid_client(fc_server): class ValidAuthorization: def get_access_token(self): return "VALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def valid_credentials(): +def valid_credentials(fc_server, auth_server): return [ - "--firecrest-url=http://firecrest.cscs.ch", + f"--firecrest-url={fc_server.url_for('/')}", "--client-id=valid_id", "--client-secret=valid_secret", - "--token-url=https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + f"--token-url={auth_server.url_for('/auth/token')}", ] @pytest.fixture -def invalid_client(): +def invalid_client(fc_server): class InvalidAuthorization: def get_access_token(self): return "INVALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=InvalidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() ) -def internal_transfer_callback(request, uri, response_headers): +def internal_transfer_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Failed to submit job", "error": "Machine does not exist"}', - ] + return Response( + json.dumps( + { + "description": "Failed to submit job", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) ret = { "success": "Task created", @@ -65,31 +74,42 @@ def internal_transfer_callback(request, uri, response_headers): "task_url": "TASK_IP/tasks/internal_transfer_id", } status_code = 201 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), status=status_code, content_type="application/json" + ) -def external_download_callback(request, uri, response_headers): +def external_download_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] - - # TODO Machine is ignored at this point - # if request.headers["X-Machine-Name"] != "cluster1": - # response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - # return [ - # 400, - # response_headers, - # '{"description": "Failed to submit job", "error": "Machine does not exist"}', - # ] - - # I couldn't find a better way to get the params from the request - if b"sourcePath=%2Fpath%2Fto%2Fremote%2Fsourcelegacy" in request.body: + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) + + if request.headers["X-Machine-Name"] != "cluster1": + return Response( + json.dumps( + { + "description": "Failed to download file", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + extra_headers = None + source_path = request.form.get("sourcePath") + if source_path == "/path/to/remote/sourcelegacy": ret = { "success": "Task created", "task_id": "external_download_id_legacy", "task_url": "TASK_IP/tasks/external_download_id_legacy", } status_code = 201 - elif b"sourcePath=%2Fpath%2Fto%2Fremote%2Fsource" in request.body: + elif source_path == "/path/to/remote/source": ret = { "success": "Task created", "task_id": "external_download_id", @@ -97,30 +117,45 @@ def external_download_callback(request, uri, response_headers): } status_code = 201 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "sourcePath error"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def external_upload_callback(request, uri, response_headers): +def external_upload_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] - - # TODO Machine is ignored at this point - # if request.headers["X-Machine-Name"] != "cluster1": - # response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - # return [ - # 400, - # response_headers, - # '{"description": "Failed to submit job", "error": "Machine does not exist"}', - # ] - - # I couldn't find a better way to get the params from the request + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) + + if request.headers["X-Machine-Name"] != "cluster1": + return Response( + json.dumps( + { + "description": "Failed to upload file", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + source_path = request.form.get("sourcePath") + target_path = request.form.get("targetPath") + extra_headers = None if ( - b"sourcePath=%2Fpath%2Fto%2Flocal%2Fsource" in request.body - and b"targetPath=%2Fpath%2Fto%2Fremote%2Fdestination" in request.body + source_path == "/path/to/local/source" + and target_path == "/path/to/remote/destination" ): ret = { "success": "Task created", @@ -129,11 +164,16 @@ def external_upload_callback(request, uri, response_headers): } status_code = 201 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "sourcePath error"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) # Global variables for tasks @@ -145,14 +185,19 @@ def external_upload_callback(request, uri, response_headers): external_download_result = 0 -def storage_tasks_callback(request, uri, response_headers): +def storage_tasks_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) global internal_transfer_retry global external_download_retry global external_upload_retry + uri = request.url taskid = uri.split("/")[-1] if taskid == "tasks": # TODO: return all tasks @@ -251,7 +296,7 @@ def storage_tasks_callback(request, uri, response_headers): "data": { "source": "/path/to/remote/source", "system_name": "machine", - "url": "https://object_storage_link.com" + "url": "https://object_storage_link.com", }, "description": "Started upload from filesystem to Object Storage", "hash_id": taskid, @@ -354,47 +399,36 @@ def storage_tasks_callback(request, uri, response_headers): } status_code = 200 - return [status_code, response_headers, json.dumps(ret)] - + return Response( + json.dumps(ret), status=status_code, content_type="application/json" + ) -@pytest.fixture(autouse=True) -def setup_callbacks(): - httpretty.enable(allow_net_connect=False, verbose=True) - httpretty.register_uri( - httpretty.POST, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/storage\/xfer-internal.*"), - body=internal_transfer_callback, - ) +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request( + re.compile("^/storage/xfer-internal.*"), method="POST" + ).respond_with_handler(internal_transfer_handler) - httpretty.register_uri( - httpretty.GET, - re.compile(r"http:\/\/firecrest\.cscs\.ch\/tasks.*"), - body=storage_tasks_callback, - ) + httpserver.expect_request( + re.compile("^/tasks/.*"), method="GET" + ).respond_with_handler(storage_tasks_handler) - httpretty.register_uri( - httpretty.POST, - "http://firecrest.cscs.ch/storage/xfer-external/download", - body=external_download_callback, - ) + httpserver.expect_request( + "/storage/xfer-external/download", method="POST" + ).respond_with_handler(external_download_handler) - httpretty.register_uri( - httpretty.POST, - "http://firecrest.cscs.ch/storage/xfer-external/upload", - body=external_upload_callback, - ) + httpserver.expect_request( + "/storage/xfer-external/upload", method="POST" + ).respond_with_handler(external_upload_handler) - httpretty.register_uri( - httpretty.POST, - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", - body=auth.auth_callback, - ) + return httpserver - yield - httpretty.disable() - httpretty.reset() +@pytest.fixture +def auth_server(httpserver): + httpserver.expect_request("/auth/token").respond_with_handler(auth.auth_handler) + return httpserver def test_internal_transfer(valid_client): @@ -583,7 +617,7 @@ def test_external_download(valid_client): external_download_retry = 0 valid_client.set_api_version("1.14.0") obj = valid_client.external_download("cluster1", "/path/to/remote/source") - assert isinstance(obj, ExternalDownload) + assert isinstance(obj, firecrest.ExternalDownload) assert obj._task_id == "external_download_id" assert obj.client == valid_client @@ -593,7 +627,7 @@ def test_external_download_legacy(valid_client): external_download_retry = 0 valid_client.set_api_version("1.13.0") obj = valid_client.external_download("cluster1", "/path/to/remote/sourcelegacy") - assert isinstance(obj, ExternalDownload) + assert isinstance(obj, firecrest.ExternalDownload) assert obj._task_id == "external_download_id_legacy" assert obj.client == valid_client @@ -646,7 +680,7 @@ def test_external_upload(valid_client): obj = valid_client.external_upload( "cluster1", "/path/to/local/source", "/path/to/remote/destination" ) - assert isinstance(obj, ExternalUpload) + assert isinstance(obj, firecrest.ExternalUpload) assert obj._task_id == "external_upload_id" assert obj.client == valid_client diff --git a/tests/test_storage_async.py b/tests/test_storage_async.py new file mode 100644 index 0000000..223925a --- /dev/null +++ b/tests/test_storage_async.py @@ -0,0 +1,288 @@ +import pytest +import re +import test_storage as basic_storage + +from context import firecrest + +from firecrest import __app_name__, __version__ + + +@pytest.fixture +def valid_client(fc_server): + class ValidAuthorization: + def get_access_token(self): + return "VALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def invalid_client(fc_server): + class InvalidAuthorization: + def get_access_token(self): + return "INVALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request( + re.compile("^/storage/xfer-internal.*"), method="POST" + ).respond_with_handler(basic_storage.internal_transfer_handler) + + httpserver.expect_request( + re.compile("^/tasks/.*"), method="GET" + ).respond_with_handler(basic_storage.storage_tasks_handler) + + httpserver.expect_request( + "/storage/xfer-external/download", method="POST" + ).respond_with_handler(basic_storage.external_download_handler) + + httpserver.expect_request( + "/storage/xfer-external/upload", method="POST" + ).respond_with_handler(basic_storage.external_upload_handler) + + return httpserver + + +@pytest.mark.asyncio +async def test_internal_transfer(valid_client): + global internal_transfer_retry + + # mv job + internal_transfer_retry = 0 + assert await valid_client.submit_move_job( + machine="cluster1", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) == { + "job_data_err": "", + "job_data_out": "", + "job_file": "/path/to/firecrest/internal_transfer_id/sbatch-job.sh", + "job_file_err": "/path/to/firecrest/internal_transfer_id/job-35363861.err", + "job_file_out": "/path/to/firecrest/internal_transfer_id/job-35363861.out", + "jobid": 35363861, + "result": "Job submitted", + } + + # cp job + internal_transfer_retry = 0 + assert await valid_client.submit_copy_job( + machine="cluster1", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) == { + "job_data_err": "", + "job_data_out": "", + "job_file": "/path/to/firecrest/internal_transfer_id/sbatch-job.sh", + "job_file_err": "/path/to/firecrest/internal_transfer_id/job-35363861.err", + "job_file_out": "/path/to/firecrest/internal_transfer_id/job-35363861.out", + "jobid": 35363861, + "result": "Job submitted", + } + + # rsync job + internal_transfer_retry = 0 + assert await valid_client.submit_rsync_job( + machine="cluster1", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) == { + "job_data_err": "", + "job_data_out": "", + "job_file": "/path/to/firecrest/internal_transfer_id/sbatch-job.sh", + "job_file_err": "/path/to/firecrest/internal_transfer_id/job-35363861.err", + "job_file_out": "/path/to/firecrest/internal_transfer_id/job-35363861.out", + "jobid": 35363861, + "result": "Job submitted", + } + + # rm job + internal_transfer_retry = 0 + assert await valid_client.submit_delete_job( + machine="cluster1", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) == { + "job_data_err": "", + "job_data_out": "", + "job_file": "/path/to/firecrest/internal_transfer_id/sbatch-job.sh", + "job_file_err": "/path/to/firecrest/internal_transfer_id/job-35363861.err", + "job_file_out": "/path/to/firecrest/internal_transfer_id/job-35363861.out", + "jobid": 35363861, + "result": "Job submitted", + } + + +@pytest.mark.asyncio +async def test_internal_transfer_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + # mv job + await valid_client.submit_move_job( + machine="cluster2", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) + + with pytest.raises(firecrest.HeaderException): + # cp job + await valid_client.submit_copy_job( + machine="cluster2", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) + + with pytest.raises(firecrest.HeaderException): + # rsync job + await valid_client.submit_rsync_job( + machine="cluster2", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) + + with pytest.raises(firecrest.HeaderException): + # rm job + await valid_client.submit_delete_job( + machine="cluster2", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) + + +@pytest.mark.asyncio +async def test_internal_transfer_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + # mv job + await invalid_client.submit_move_job( + machine="cluster1", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) + + with pytest.raises(firecrest.UnauthorizedException): + # cp job + await invalid_client.submit_copy_job( + machine="cluster1", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) + + with pytest.raises(firecrest.UnauthorizedException): + # rsync job + await invalid_client.submit_rsync_job( + machine="cluster1", + source_path="/path/to/source", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) + + with pytest.raises(firecrest.UnauthorizedException): + # rm job + await invalid_client.submit_delete_job( + machine="cluster1", + target_path="/path/to/destination", + job_name="mv-job", + time="2", + stage_out_job_id="35363851", + account="project", + ) + + +@pytest.mark.asyncio +async def test_external_download(valid_client): + global external_download_retry + external_download_retry = 0 + valid_client.set_api_version("1.14.0") + obj = await valid_client.external_download("cluster1", "/path/to/remote/source") + assert isinstance(obj, firecrest.AsyncExternalDownload) + assert obj._task_id == "external_download_id" + assert obj.client == valid_client + + +@pytest.mark.asyncio +async def test_external_download_legacy(valid_client): + global external_download_retry + external_download_retry = 0 + valid_client.set_api_version("1.13.0") + obj = await valid_client.external_download("cluster1", "/path/to/remote/sourcelegacy") + assert isinstance(obj, firecrest.AsyncExternalDownload) + assert obj._task_id == "external_download_id_legacy" + assert obj.client == valid_client + + +@pytest.mark.asyncio +async def test_external_upload(valid_client): + global external_upload_retry + external_upload_retry = 0 + obj = await valid_client.external_upload( + "cluster1", "/path/to/local/source", "/path/to/remote/destination" + ) + assert isinstance(obj, firecrest.AsyncExternalUpload) + assert obj._task_id == "external_upload_id" + assert obj.client == valid_client diff --git a/tests/test_utilities.py b/tests/test_utilities.py index 9a577f7..2554d62 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -1,64 +1,73 @@ import common -import httpretty import json import pytest -import re -import test_authoriation as auth +import test_authorisation as auth from context import firecrest from firecrest import __app_name__, __version__, cli from typer.testing import CliRunner +from werkzeug.wrappers import Response +from werkzeug.wrappers import Request runner = CliRunner() @pytest.fixture -def valid_client(): +def valid_client(fc_server): class ValidAuthorization: def get_access_token(self): return "VALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=ValidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() ) @pytest.fixture -def valid_credentials(): +def valid_credentials(fc_server, auth_server): return [ - "--firecrest-url=http://firecrest.cscs.ch", + f"--firecrest-url={fc_server.url_for('/')}", "--client-id=valid_id", "--client-secret=valid_secret", - "--token-url=https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", + f"--token-url={auth_server.url_for('/auth/token')}", ] @pytest.fixture -def invalid_client(): +def invalid_client(fc_server): class InvalidAuthorization: def get_access_token(self): return "INVALID_TOKEN" return firecrest.Firecrest( - firecrest_url="http://firecrest.cscs.ch", authorization=InvalidAuthorization() + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() ) -def ls_callback(request, uri, response_headers): +def ls_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on ls operation", "error": "Machine does not exist"}', - ] - - targetPath = request.querystring.get("targetPath", [None])[0] + return Response( + json.dumps( + { + "description": "Error on ls operation", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + targetPath = request.args.get("targetPath") if targetPath == "/path/to/valid/dir": ret = { "description": "List of contents", @@ -85,7 +94,7 @@ def ls_callback(request, uri, response_headers): }, ], } - showhidden = request.querystring.get("showhidden", [False])[0] + showhidden = request.args.get("showhidden", False) if showhidden: ret["output"].append( { @@ -100,57 +109,86 @@ def ls_callback(request, uri, response_headers): } ) - return [200, response_headers, json.dumps(ret)] + return Response(json.dumps(ret), status=200, content_type="application/json") if targetPath == "/path/to/invalid/dir": - response_headers["X-Invalid-Path"] = "path is an invalid path" - return [400, response_headers, '{"description": "Error on ls operation"}'] + extra_headers = {"X-Invalid-Path": "path is an invalid path"} + ret = {"description": "Error on ls operation"} + return Response( + json.dumps(ret), + status=400, + headers=extra_headers, + content_type="application/json", + ) -def mkdir_callback(request, uri, response_headers): +def mkdir_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on mkdir operation", "error": "Machine does not exist"}', - ] - - target_path = request.parsed_body["targetPath"][0] - p = request.parsed_body.get("p", [False])[0] + return Response( + json.dumps( + { + "description": "Error on mkdir operation", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + extra_headers = None + target_path = request.form["targetPath"] + p = request.form.get("p", False) if target_path == "path/to/valid/dir" or ( target_path == "path/to/valid/dir/with/p" and p ): ret = {"description": "Success to mkdir file or directory.", "output": ""} status_code = 201 else: - response_headers[ + extra_headers[ "X-Invalid-Path" ] = "sourcePath and/or targetPath are invalid paths" ret = {"description": "Error on mkdir operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def mv_callback(request, uri, response_headers): +def mv_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on rename operation", "error": "Machine does not exist"}', - ] - - source_path = request.parsed_body["sourcePath"][0] - target_path = request.parsed_body["targetPath"][0] + return Response( + json.dumps( + { + "description": "Error on rename operation", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + source_path = request.form["sourcePath"] + target_path = request.form["targetPath"] if ( source_path == "/path/to/valid/source" and target_path == "/path/to/valid/destination" @@ -162,24 +200,35 @@ def mv_callback(request, uri, response_headers): ret = {"description": "Error on rename operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), status=status_code, content_type="application/json" + ) -def chmod_callback(request, uri, response_headers): +def chmod_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on chmod operation", "error": "Machine does not exist"}', - ] - - target_path = request.parsed_body["targetPath"][0] - mode = request.parsed_body["mode"][0] + return Response( + json.dumps( + { + "description": "Error on chmod operation", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + extra_headers = None + target_path = request.form["targetPath"] + mode = request.form["mode"] if target_path == "/path/to/valid/file" and mode == "777": ret = { "description": "Success to chmod file or directory.", @@ -188,29 +237,43 @@ def chmod_callback(request, uri, response_headers): status_code = 200 else: # FIXME: FirecREST sets the X-Invalid-Path even when the problem is the mode argument - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on chmod operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def chown_callback(request, uri, response_headers): +def chown_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on chown operation", "error": "Machine does not exist"}', - ] - - target_path = request.parsed_body["targetPath"][0] - owner = request.parsed_body.get("owner", [""])[0] - group = request.parsed_body.get("group", [""])[0] + return Response( + json.dumps( + { + "description": "Error on chown operation", + "error": "Machine does not exist", + } + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + extra_headers = None + target_path = request.form["targetPath"] + owner = request.form.get("owner", "") + group = request.form.get("group", "") if target_path == "/path/to/file" and owner == "new_owner" and group == "new_group": ret = { "description": "Success to chown file or directory.", @@ -218,28 +281,39 @@ def chown_callback(request, uri, response_headers): } status_code = 200 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on chown operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def copy_callback(request, uri, response_headers): +def copy_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on copy operation", "error": "Machine does not exist"}', - ] - - source_path = request.parsed_body["sourcePath"][0] - target_path = request.parsed_body["targetPath"][0] + return Response( + json.dumps( + {"description": "Error on copy operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + extra_headers = None + source_path = request.form["sourcePath"] + target_path = request.form["targetPath"] if ( source_path == "/path/to/valid/source" and target_path == "/path/to/valid/destination" @@ -247,26 +321,38 @@ def copy_callback(request, uri, response_headers): ret = {"description": "Success to copy file or directory.", "output": ""} status_code = 201 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on copy operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def file_type_callback(request, uri, response_headers): +def file_type_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on file operation", "error": "Machine does not exist"}', - ] - - targetPath = request.querystring.get("targetPath", [None])[0] + return Response( + json.dumps( + {"description": "Error on file operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + extra_headers = None + targetPath = request.args.get("targetPath") if targetPath == "/path/to/empty/file": ret = {"description": "Success to file file or directory.", "output": "empty"} status_code = 200 @@ -277,27 +363,39 @@ def file_type_callback(request, uri, response_headers): } status_code = 200 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on file operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def stat_callback(request, uri, response_headers): +def stat_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on file operation", "error": "Machine does not exist"}', - ] - - targetPath = request.querystring.get("targetPath", [None])[0] - deref = request.querystring.get("dereference", [False])[0] + return Response( + json.dumps( + {"description": "Error on file operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + extra_headers = None + targetPath = request.args.get("targetPath") + deref = request.args.get("dereference", False) if targetPath == "/path/to/link": if deref: ret = { @@ -334,317 +432,388 @@ def stat_callback(request, uri, response_headers): } status_code = 200 else: - response_headers["X-Not-Found"] = "sourcePath not found" + extra_headers = {"X-Not-Found": "sourcePath not found"} ret = {"description": "Error on stat operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def symlink_callback(request, uri, response_headers): +def symlink_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on symlink operation", "error": "Machine does not exist"}', - ] - - target_path = request.parsed_body["targetPath"][0] - link_path = request.parsed_body["linkPath"][0] + return Response( + json.dumps( + {"description": "Error on symlink operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + extra_headers = None + target_path = request.form["targetPath"] + link_path = request.form["linkPath"] if target_path == "/path/to/file" and link_path == "/path/to/link": ret = {"description": "Success to link file or directory.", "output": ""} status_code = 201 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on symlink operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def simple_download_callback(request, uri, response_headers): +def simple_download_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on download operation", "error": "Machine does not exist"}', - ] - - source_path = request.querystring.get("sourcePath", [None])[0] + return Response( + json.dumps( + {"description": "Error on download operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + source_path = request.args.get("sourcePath") if source_path == "/path/to/remote/source": ret = "Hello!\n" status_code = 200 - return [status_code, response_headers, ret] + return Response(ret, status=status_code) else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on download operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def simple_upload_callback(request, uri, response_headers): +def simple_upload_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on download operation", "error": "Machine does not exist"}', - ] - - # I couldn't find a better way to get the params from the request - if ( - b'form-data; name="targetPath"\r\n\r\n/path/to/remote/destination' - in request.body - ): + return Response( + json.dumps( + {"description": "Error on download operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + if request.form["targetPath"] == "/path/to/remote/destination": + extra_headers = None ret = {"description": "File upload successful", "output": ""} status_code = 201 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on upload operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def simple_delete_callback(request, uri, response_headers): +def simple_delete_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on download operation", "error": "Machine does not exist"}', - ] - - target_path = request.parsed_body["targetPath"][0] + return Response( + json.dumps( + {"description": "Error on download operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + target_path = request.form["targetPath"] if target_path == "/path/to/file": + extra_headers = None ret = {"description": "File delete successful", "output": ""} status_code = 204 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on delete operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def checksum_callback(request, uri, response_headers): +def checksum_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on checksum operation", "error": "Machine does not exist"}', - ] - - target_path = request.querystring.get("targetPath", [None])[0] + return Response( + json.dumps( + {"description": "Error on checksum operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + target_path = request.args.get("targetPath") if target_path == "/path/to/file": + extra_headers = None ret = { "description": "Success to checksum file or directory.", "output": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", } status_code = 200 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on checksum operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def view_callback(request, uri, response_headers): +def view_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on head operation", "error": "Machine does not exist"}', - ] - - target_path = request.querystring.get("targetPath", [None])[0] + return Response( + json.dumps( + {"description": "Error on head operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + target_path = request.args.get("targetPath") if target_path == "/path/to/file": + extra_headers = None ret = {"description": "Success to head file or directory.", "output": "hello\n"} status_code = 200 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on head operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] - + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def head_tail_callback(request, uri, response_headers): - is_tail_req = "tail" in uri +def head_tail_handler(request: Request): + is_tail_req = "tail" in request.url if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on head operation", "error": "Machine does not exist"}', - ] - - target_path = request.querystring.get("targetPath", [None])[0] - lines = request.querystring.get("lines", [None])[0] - bytes = request.querystring.get("bytes", [None])[0] + return Response( + json.dumps( + {"description": "Error on head operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + target_path = request.args.get("targetPath") + lines = request.args.get("lines") + bytes = request.args.get("bytes") if target_path == "/path/to/file": + extra_headers = None if lines and int(lines) < 10: - result = int(lines)*"hello\n" + result = int(lines) * "hello\n" else: - result = 10*"hello\n" + result = 10 * "hello\n" if bytes: if is_tail_req: - result = result[-int(bytes):] + result = result[-int(bytes) :] else: - result = result[0:int(bytes)] + result = result[0 : int(bytes)] ret = {"description": "Success to head file.", "output": result} status_code = 200 else: - response_headers["X-Invalid-Path"] = "path is an invalid path" + extra_headers = {"X-Invalid-Path": "path is an invalid path"} ret = {"description": "Error on head operation"} status_code = 400 - return [status_code, response_headers, json.dumps(ret)] + return Response( + json.dumps(ret), + status=status_code, + headers=extra_headers, + content_type="application/json", + ) -def whoami_callback(request, uri, response_headers): +def whoami_handler(request: Request): if request.headers["Authorization"] != "Bearer VALID_TOKEN": - return [401, response_headers, '{"message": "Bad token; invalid JSON"}'] + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) if request.headers["X-Machine-Name"] != "cluster1": - response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist" - return [ - 400, - response_headers, - '{"description": "Error on whoami operation", "error": "Machine does not exist"}', - ] - - return [ - 200, - response_headers, - '{"description": "Success on whoami operation.", "output": "username"}' - ] + return Response( + json.dumps( + {"description": "Error on whoami operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + return Response( + json.dumps({"description": "Success on whoami operation.", "output": "username"}), + status=200, + content_type="application/json", + ) -@pytest.fixture(autouse=True) -def setup_callbacks(): - httpretty.enable(allow_net_connect=False, verbose=True) - httpretty.register_uri( - httpretty.GET, "http://firecrest.cscs.ch/utilities/ls", body=ls_callback +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request("/utilities/ls", method="GET").respond_with_handler( + ls_handler ) - httpretty.register_uri( - httpretty.POST, "http://firecrest.cscs.ch/utilities/mkdir", body=mkdir_callback + httpserver.expect_request("/utilities/mkdir", method="POST").respond_with_handler( + mkdir_handler ) - httpretty.register_uri( - httpretty.PUT, "http://firecrest.cscs.ch/utilities/rename", body=mv_callback + httpserver.expect_request("/utilities/rename", method="PUT").respond_with_handler( + mv_handler ) - httpretty.register_uri( - httpretty.PUT, "http://firecrest.cscs.ch/utilities/chmod", body=chmod_callback + httpserver.expect_request("/utilities/chmod", method="PUT").respond_with_handler( + chmod_handler ) - httpretty.register_uri( - httpretty.PUT, "http://firecrest.cscs.ch/utilities/chown", body=chown_callback + httpserver.expect_request("/utilities/chown", method="PUT").respond_with_handler( + chown_handler ) - httpretty.register_uri( - httpretty.POST, "http://firecrest.cscs.ch/utilities/copy", body=copy_callback + httpserver.expect_request("/utilities/copy", method="POST").respond_with_handler( + copy_handler ) - httpretty.register_uri( - httpretty.GET, - "http://firecrest.cscs.ch/utilities/file", - body=file_type_callback, + httpserver.expect_request("/utilities/file", method="GET").respond_with_handler( + file_type_handler ) - httpretty.register_uri( - httpretty.GET, "http://firecrest.cscs.ch/utilities/stat", body=stat_callback + httpserver.expect_request("/utilities/stat", method="GET").respond_with_handler( + stat_handler ) - httpretty.register_uri( - httpretty.POST, - "http://firecrest.cscs.ch/utilities/symlink", - body=symlink_callback, + httpserver.expect_request("/utilities/symlink", method="POST").respond_with_handler( + symlink_handler ) - httpretty.register_uri( - httpretty.GET, - "http://firecrest.cscs.ch/utilities/download", - body=simple_download_callback, + httpserver.expect_request("/utilities/download", method="GET").respond_with_handler( + simple_download_handler ) - httpretty.register_uri( - httpretty.POST, - "http://firecrest.cscs.ch/utilities/upload", - body=simple_upload_callback, + httpserver.expect_request("/utilities/upload", method="POST").respond_with_handler( + simple_upload_handler ) - httpretty.register_uri( - httpretty.DELETE, - "http://firecrest.cscs.ch/utilities/rm", - body=simple_delete_callback, + httpserver.expect_request("/utilities/rm", method="DELETE").respond_with_handler( + simple_delete_handler ) - httpretty.register_uri( - httpretty.GET, - "http://firecrest.cscs.ch/utilities/checksum", - body=checksum_callback, + httpserver.expect_request("/utilities/checksum", method="GET").respond_with_handler( + checksum_handler ) - httpretty.register_uri( - httpretty.GET, "http://firecrest.cscs.ch/utilities/view", body=view_callback + httpserver.expect_request("/utilities/view", method="GET").respond_with_handler( + view_handler ) - httpretty.register_uri( - httpretty.GET, "http://firecrest.cscs.ch/utilities/head", body=head_tail_callback + httpserver.expect_request("/utilities/head", method="GET").respond_with_handler( + head_tail_handler ) - httpretty.register_uri( - httpretty.GET, "http://firecrest.cscs.ch/utilities/tail", body=head_tail_callback + httpserver.expect_request("/utilities/tail", method="GET").respond_with_handler( + head_tail_handler ) - httpretty.register_uri( - httpretty.GET, "http://firecrest.cscs.ch/utilities/whoami", body=whoami_callback + httpserver.expect_request("/utilities/whoami", method="GET").respond_with_handler( + whoami_handler ) - httpretty.register_uri( - httpretty.POST, - "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token", - body=auth.auth_callback, - ) + return httpserver - yield - httpretty.disable() - httpretty.reset() +@pytest.fixture +def auth_server(httpserver): + httpserver.expect_request("/auth/token").respond_with_handler(auth.auth_handler) + return httpserver def test_list_files(valid_client): @@ -1230,8 +1399,8 @@ def test_view(valid_client): def test_head(valid_client): - assert valid_client.head("cluster1", "/path/to/file") == 10*"hello\n" - assert valid_client.head("cluster1", "/path/to/file", lines=2) == 2*"hello\n" + assert valid_client.head("cluster1", "/path/to/file") == 10 * "hello\n" + assert valid_client.head("cluster1", "/path/to/file", lines=2) == 2 * "hello\n" assert valid_client.head("cluster1", "/path/to/file", bytes=4) == "hell" @@ -1270,8 +1439,8 @@ def test_cli_head(valid_credentials): def test_tail(valid_client): - assert valid_client.tail("cluster1", "/path/to/file") == 10*"hello\n" - assert valid_client.tail("cluster1", "/path/to/file", lines=2) == 2*"hello\n" + assert valid_client.tail("cluster1", "/path/to/file") == 10 * "hello\n" + assert valid_client.tail("cluster1", "/path/to/file", lines=2) == 2 * "hello\n" assert valid_client.tail("cluster1", "/path/to/file", bytes=5) == "ello\n" @@ -1325,14 +1494,14 @@ def test_view_invalid_client(invalid_client): def test_whoami(valid_client): - assert valid_client.whoami('cluster1') == "username" + assert valid_client.whoami("cluster1") == "username" def test_whoami_invalid_machine(valid_client): with pytest.raises(firecrest.HeaderException): - valid_client.whoami('cluster2') + valid_client.whoami("cluster2") def test_whoami_invalid_client(invalid_client): with pytest.raises(firecrest.UnauthorizedException): - invalid_client.whoami('cluster1') + invalid_client.whoami("cluster1") diff --git a/tests/test_utilities_async.py b/tests/test_utilities_async.py new file mode 100644 index 0000000..6606d2f --- /dev/null +++ b/tests/test_utilities_async.py @@ -0,0 +1,653 @@ +import pytest +import test_utilities as basic_utilities + +from context import firecrest + +from firecrest import __app_name__, __version__ + + +@pytest.fixture +def valid_client(fc_server): + class ValidAuthorization: + def get_access_token(self): + return "VALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=ValidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def invalid_client(fc_server): + class InvalidAuthorization: + def get_access_token(self): + return "INVALID_TOKEN" + + client = firecrest.AsyncFirecrest( + firecrest_url=fc_server.url_for("/"), authorization=InvalidAuthorization() + ) + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + return client + + +@pytest.fixture +def fc_server(httpserver): + httpserver.expect_request("/utilities/ls", method="GET").respond_with_handler( + basic_utilities.ls_handler + ) + + httpserver.expect_request("/utilities/mkdir", method="POST").respond_with_handler( + basic_utilities.mkdir_handler + ) + + httpserver.expect_request("/utilities/rename", method="PUT").respond_with_handler( + basic_utilities.mv_handler + ) + + httpserver.expect_request("/utilities/chmod", method="PUT").respond_with_handler( + basic_utilities.chmod_handler + ) + + httpserver.expect_request("/utilities/chown", method="PUT").respond_with_handler( + basic_utilities.chown_handler + ) + + httpserver.expect_request("/utilities/copy", method="POST").respond_with_handler( + basic_utilities.copy_handler + ) + + httpserver.expect_request("/utilities/file", method="GET").respond_with_handler( + basic_utilities.file_type_handler + ) + + httpserver.expect_request("/utilities/stat", method="GET").respond_with_handler( + basic_utilities.stat_handler + ) + + httpserver.expect_request("/utilities/symlink", method="POST").respond_with_handler( + basic_utilities.symlink_handler + ) + + httpserver.expect_request("/utilities/download", method="GET").respond_with_handler( + basic_utilities.simple_download_handler + ) + + httpserver.expect_request("/utilities/upload", method="POST").respond_with_handler( + basic_utilities.simple_upload_handler + ) + + httpserver.expect_request("/utilities/rm", method="DELETE").respond_with_handler( + basic_utilities.simple_delete_handler + ) + + httpserver.expect_request("/utilities/checksum", method="GET").respond_with_handler( + basic_utilities.checksum_handler + ) + + httpserver.expect_request("/utilities/view", method="GET").respond_with_handler( + basic_utilities.view_handler + ) + + httpserver.expect_request("/utilities/head", method="GET").respond_with_handler( + basic_utilities.head_tail_handler + ) + + httpserver.expect_request("/utilities/tail", method="GET").respond_with_handler( + basic_utilities.head_tail_handler + ) + + httpserver.expect_request("/utilities/whoami", method="GET").respond_with_handler( + basic_utilities.whoami_handler + ) + + return httpserver + + +@pytest.mark.asyncio +async def test_list_files(valid_client): + assert await valid_client.list_files("cluster1", "/path/to/valid/dir") == [ + { + "group": "group", + "last_modified": "2021-08-10T15:26:52", + "link_target": "", + "name": "file.txt", + "permissions": "r-xr-xr-x.", + "size": "180", + "type": "-", + "user": "user", + }, + { + "group": "group", + "last_modified": "2021-10-07T09:17:01", + "link_target": "", + "name": "projectdir", + "permissions": "rwxr-xr-x", + "size": "4096", + "type": "d", + "user": "user", + }, + ] + + assert await valid_client.list_files( + "cluster1", "/path/to/valid/dir", show_hidden=True + ) == [ + { + "group": "group", + "last_modified": "2021-08-10T15:26:52", + "link_target": "", + "name": "file.txt", + "permissions": "r-xr-xr-x.", + "size": "180", + "type": "-", + "user": "user", + }, + { + "group": "group", + "last_modified": "2021-10-07T09:17:01", + "link_target": "", + "name": "projectdir", + "permissions": "rwxr-xr-x", + "size": "4096", + "type": "d", + "user": "user", + }, + { + "group": "group", + "last_modified": "2021-11-26T09:34:59", + "link_target": "", + "name": ".hiddenfile", + "permissions": "rwxrwxr-x", + "size": "4096", + "type": "-", + "user": "user", + }, + ] + + +@pytest.mark.asyncio +async def test_list_files_invalid_path(valid_client): + with pytest.raises(firecrest.FirecrestException): + await valid_client.list_files("cluster1", "/path/to/invalid/dir") + + +@pytest.mark.asyncio +async def test_list_files_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.list_files("cluster2", "/path/to/dir") + + +@pytest.mark.asyncio +async def test_list_files_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.list_files("cluster1", "/path/to/dir") + + +@pytest.mark.asyncio +async def test_mkdir(valid_client): + # Make sure these don't raise an error + await valid_client.mkdir("cluster1", "path/to/valid/dir") + await valid_client.mkdir("cluster1", "path/to/valid/dir/with/p", p=True) + + +@pytest.mark.asyncio +async def test_mkdir_invalid_path(valid_client): + with pytest.raises(firecrest.FirecrestException): + await valid_client.mkdir("cluster1", "/path/to/invalid/dir") + + +@pytest.mark.asyncio +async def test_mkdir_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.mkdir("cluster2", "path/to/dir") + + +@pytest.mark.asyncio +async def test_mkdir_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.mkdir("cluster1", "path/to/dir") + + +@pytest.mark.asyncio +async def test_mv(valid_client): + # Make sure this doesn't raise an error + await valid_client.mv("cluster1", "/path/to/valid/source", "/path/to/valid/destination") + + +@pytest.mark.asyncio +async def test_mv_invalid_paths(valid_client): + with pytest.raises(firecrest.FirecrestException): + await valid_client.mv( + "cluster1", "/path/to/invalid/source", "/path/to/valid/destination" + ) + + with pytest.raises(firecrest.FirecrestException): + await valid_client.mv( + "cluster1", "/path/to/valid/source", "/path/to/invalid/destination" + ) + + +@pytest.mark.asyncio +async def test_mv_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.mv("cluster2", "/path/to/source", "/path/to/destination") + + +@pytest.mark.asyncio +async def test_mv_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.mv("cluster1", "/path/to/source", "/path/to/destination") + + +@pytest.mark.asyncio +async def test_chmod(valid_client): + # Make sure this doesn't raise an error + await valid_client.chmod("cluster1", "/path/to/valid/file", "777") + + +@pytest.mark.asyncio +async def test_chmod_invalid_arguments(valid_client): + with pytest.raises(firecrest.FirecrestException): + await valid_client.chmod("cluster1", "/path/to/invalid/file", "777") + + with pytest.raises(firecrest.FirecrestException): + await valid_client.chmod("cluster1", "/path/to/valid/file", "random_string") + + +@pytest.mark.asyncio +async def test_chmod_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.chmod("cluster2", "/path/to/file", "700") + + +@pytest.mark.asyncio +async def test_chmod_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.chmod("cluster1", "/path/to/file", "600") + + +@pytest.mark.asyncio +async def test_chown(valid_client): + # Make sure this doesn't raise an error + await valid_client.chown( + "cluster1", "/path/to/file", owner="new_owner", group="new_group" + ) + # Call will immediately return if neither owner and nor group are set + await valid_client.chown("cluster", "path") + + +@pytest.mark.asyncio +async def test_chown_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.chown( + "cluster1", "/bad/path", owner="new_owner", group="new_group" + ) + + with pytest.raises(firecrest.HeaderException): + await valid_client.chown( + "cluster1", "/path/to/file", owner="bad_owner", group="new_group" + ) + + with pytest.raises(firecrest.HeaderException): + await valid_client.chown( + "cluster1", "/path/to/file", owner="new_owner", group="bad_group" + ) + + +@pytest.mark.asyncio +async def test_chown_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.chown( + "cluster2", "/path/to/file", owner="new_owner", group="new_group" + ) + + +@pytest.mark.asyncio +async def test_chown_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.chown( + "cluster1", "/path/to/file", owner="new_owner", group="new_group" + ) + + +@pytest.mark.asyncio +async def test_copy(valid_client): + # Make sure this doesn't raise an error + await valid_client.copy("cluster1", "/path/to/valid/source", "/path/to/valid/destination") + + +@pytest.mark.asyncio +async def test_copy_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.copy( + "cluster1", "/path/to/invalid/source", "/path/to/valid/destination" + ) + + with pytest.raises(firecrest.HeaderException): + await valid_client.copy( + "cluster1", "/path/to/valid/source", "/path/to/invalid/destination" + ) + + +@pytest.mark.asyncio +async def test_copy_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.copy("cluster2", "/path/to/source", "/path/to/destination") + + +@pytest.mark.asyncio +async def test_copy_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.copy("cluster1", "/path/to/source", "/path/to/destination") + + +@pytest.mark.asyncio +async def test_file_type(valid_client): + assert await valid_client.file_type("cluster1", "/path/to/empty/file") == "empty" + assert await valid_client.file_type("cluster1", "/path/to/directory") == "directory" + + +@pytest.mark.asyncio +async def test_file_type_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.file_type("cluster1", "/path/to/invalid/file") + + +@pytest.mark.asyncio +async def test_file_type_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.file_type("cluster2", "/path/to/file") + + +@pytest.mark.asyncio +async def test_file_type_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.file_type("cluster1", "/path/to/file") + + +@pytest.mark.asyncio +async def test_stat(valid_client): + assert await valid_client.stat("cluster1", "/path/to/link") == { + "atime": 1655197211, + "ctime": 1655197211, + "dev": 2418024346, + "gid": 1000, + "ino": 648577971375854279, + "mode": 777, + "mtime": 1655197211, + "nlink": 1, + "size": 8, + "uid": 25948, + } + assert await valid_client.stat("cluster1", "/path/to/link", dereference=False) == { + "atime": 1655197211, + "ctime": 1655197211, + "dev": 2418024346, + "gid": 1000, + "ino": 648577971375854279, + "mode": 777, + "mtime": 1655197211, + "nlink": 1, + "size": 8, + "uid": 25948, + } + assert await valid_client.stat("cluster1", "/path/to/link", dereference=True) == { + "atime": 1653660606, + "ctime": 1653660606, + "dev": 2418024346, + "gid": 1000, + "ino": 648577914584968738, + "mode": 644, + "mtime": 1653660606, + "nlink": 1, + "size": 0, + "uid": 25948, + } + + +@pytest.mark.asyncio +async def test_stat_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.stat("cluster1", "/path/to/invalid/file") + + +@pytest.mark.asyncio +async def test_stat_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.stat("cluster2", "/path/to/file") + + +@pytest.mark.asyncio +async def test_stat_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.stat("cluster1", "/path/to/file") + + +@pytest.mark.asyncio +async def test_symlink(valid_client): + # Make sure this doesn't raise an error + await valid_client.symlink("cluster1", "/path/to/file", "/path/to/link") + + +@pytest.mark.asyncio +async def test_symlink_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.symlink("cluster1", "/path/to/invalid/file", "/path/to/link") + + with pytest.raises(firecrest.HeaderException): + await valid_client.symlink("cluster1", "/path/to/file", "/path/to/invalid/link") + + +@pytest.mark.asyncio +async def test_symlink_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.symlink("cluster2", "/path/to/file", "/path/to/link") + + +@pytest.mark.asyncio +async def test_symlink_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.symlink("cluster1", "/path/to/file", "/path/to/link") + + +@pytest.mark.asyncio +async def test_simple_download(valid_client, tmp_path): + tmp_dir = tmp_path / "download_dir" + tmp_dir.mkdir() + local_file = tmp_dir / "hello.txt" + # Make sure this doesn't raise an error + await valid_client.simple_download("cluster1", "/path/to/remote/source", local_file) + + with open(local_file) as f: + assert f.read() == "Hello!\n" + + +@pytest.mark.asyncio +async def test_simple_download_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.simple_download( + "cluster1", "/path/to/invalid/file", "/path/to/local/destination" + ) + + +@pytest.mark.asyncio +async def test_simple_download_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.simple_download( + "cluster2", "/path/to/source", "/path/to/destination" + ) + + +@pytest.mark.asyncio +async def test_simple_download_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.simple_download( + "cluster1", "/path/to/source", "/path/to/destination" + ) + + +@pytest.mark.asyncio +async def test_simple_upload(valid_client, tmp_path): + tmp_dir = tmp_path / "upload_dir" + tmp_dir.mkdir() + local_file = tmp_dir / "hello.txt" + local_file.write_text("hi") + # Make sure this doesn't raise an error + await valid_client.simple_upload("cluster1", local_file, "/path/to/remote/destination") + + +@pytest.mark.asyncio +async def test_simple_upload_invalid_arguments(valid_client, tmp_path): + tmp_dir = tmp_path / "download_dir" + tmp_dir.mkdir() + local_file = tmp_dir / "hello_invalid.txt" + local_file.write_text("hi") + with pytest.raises(firecrest.HeaderException): + await valid_client.simple_upload( + "cluster1", local_file, "/path/to/invalid/destination" + ) + + +@pytest.mark.asyncio +async def test_simple_upload_invalid_machine(valid_client, tmp_path): + tmp_dir = tmp_path / "download_dir" + tmp_dir.mkdir() + local_file = tmp_dir / "hello_invalid.txt" + local_file.write_text("hi") + with pytest.raises(firecrest.HeaderException): + await valid_client.simple_upload( + "cluster2", local_file, "/path/to/remote/destination" + ) + + +@pytest.mark.asyncio +async def test_simple_upload_invalid_client(invalid_client, tmp_path): + tmp_dir = tmp_path / "download_dir" + tmp_dir.mkdir() + local_file = tmp_dir / "hello_invalid.txt" + local_file.write_text("hi") + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.simple_upload( + "cluster1", local_file, "/path/to/remote/destination" + ) + + +@pytest.mark.asyncio +async def test_simple_delete(valid_client): + # Make sure this doesn't raise an error + await valid_client.simple_delete("cluster1", "/path/to/file") + + +@pytest.mark.asyncio +async def test_simple_delete_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.simple_delete("cluster1", "/path/to/invalid/file") + + +@pytest.mark.asyncio +async def test_simple_delete_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.simple_delete("cluster2", "/path/to/file") + + +@pytest.mark.asyncio +async def test_simple_delete_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.simple_delete("cluster1", "/path/to/file") + + +@pytest.mark.asyncio +async def test_checksum(valid_client): + assert ( + await valid_client.checksum("cluster1", "/path/to/file") + == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + ) + + +@pytest.mark.asyncio +async def test_checksum_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.checksum("cluster1", "/path/to/invalid/file") + + +@pytest.mark.asyncio +async def test_checksum_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.checksum("cluster2", "/path/to/file") + + +@pytest.mark.asyncio +async def test_checksum_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.checksum("cluster1", "/path/to/file") + + +@pytest.mark.asyncio +async def test_view(valid_client): + assert await valid_client.view("cluster1", "/path/to/file") == "hello\n" + + +@pytest.mark.asyncio +async def test_head(valid_client): + assert await valid_client.head("cluster1", "/path/to/file") == 10 * "hello\n" + assert await valid_client.head("cluster1", "/path/to/file", lines=2) == 2 * "hello\n" + assert await valid_client.head("cluster1", "/path/to/file", bytes=4) == "hell" + + +@pytest.mark.asyncio +async def test_tail(valid_client): + assert await valid_client.tail("cluster1", "/path/to/file") == 10 * "hello\n" + assert await valid_client.tail("cluster1", "/path/to/file", lines=2) == 2 * "hello\n" + assert await valid_client.tail("cluster1", "/path/to/file", bytes=5) == "ello\n" + + +@pytest.mark.asyncio +async def test_view_invalid_arguments(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.view("cluster1", "/path/to/invalid/file") + + +@pytest.mark.asyncio +async def test_view_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.view("cluster2", "/path/to/file") + + +@pytest.mark.asyncio +async def test_view_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.view("cluster1", "/path/to/file") + + +@pytest.mark.asyncio +async def test_whoami(valid_client): + assert await valid_client.whoami("cluster1") == "username" + + +@pytest.mark.asyncio +async def test_whoami_invalid_machine(valid_client): + with pytest.raises(firecrest.HeaderException): + await valid_client.whoami("cluster2") + + +@pytest.mark.asyncio +async def test_whoami_invalid_client(invalid_client): + with pytest.raises(firecrest.UnauthorizedException): + await invalid_client.whoami("cluster1")