From cf968164dbdf9b3dca156dca0a910726f8b4f5e5 Mon Sep 17 00:00:00 2001 From: "Mr. Bubbles" Date: Thu, 11 Jul 2024 07:35:39 +0200 Subject: [PATCH] Add methods to add packages and upload container (#48) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add add_package and upload_container methods * Bump version 1.3.2 → 1.4.0 --- src/pyloadapi/__init__.py | 2 +- src/pyloadapi/api.py | 192 +++++++++++++++++++++++++++++++++++++- src/pyloadapi/cli.py | 99 +++++++++++++++++++- src/pyloadapi/types.py | 11 ++- tests/conftest.py | 2 + tests/test_api.py | 153 ++++++++++++++++++++++++++++++ tests/test_cli.py | 127 ++++++++++++++++++++++++- 7 files changed, 575 insertions(+), 11 deletions(-) diff --git a/src/pyloadapi/__init__.py b/src/pyloadapi/__init__.py index c84a60f..00d98d4 100644 --- a/src/pyloadapi/__init__.py +++ b/src/pyloadapi/__init__.py @@ -1,6 +1,6 @@ """PyLoadAPI package.""" -__version__ = "1.3.2" +__version__ = "1.4.0" from .api import PyLoadAPI from .exceptions import CannotConnect, InvalidAuth, ParserError diff --git a/src/pyloadapi/api.py b/src/pyloadapi/api.py index b1b3a5a..d919950 100644 --- a/src/pyloadapi/api.py +++ b/src/pyloadapi/api.py @@ -7,6 +7,7 @@ """ from http import HTTPStatus +import json from json import JSONDecodeError import logging import traceback @@ -15,7 +16,7 @@ import aiohttp from .exceptions import CannotConnect, InvalidAuth, ParserError -from .types import LoginResponse, PyLoadCommand, StatusServerResponse +from .types import Destination, LoginResponse, PyLoadCommand, StatusServerResponse _LOGGER = logging.getLogger(__name__) @@ -167,6 +168,91 @@ async def get( "Executing command {command} failed due to request exception" ) from e + async def post( + self, + command: PyLoadCommand, + data: dict[str, Any], + ) -> Any: + """Execute a pyLoad API command using a POST request. + + Parameters + ---------- + command : PyLoadCommand + The pyLoad command to execute. + data : dict[str, Any] + Data to include in the request body. The values in the dictionary + will be JSON encoded. + + Returns + ------- + Any + The response data from the API. + + Raises + ------ + CannotConnect + If the request to the API fails due to a connection issue. + InvalidAuth + If the request fails due to invalid or expired authentication. + ParserError + If there's an error parsing the API response. + + Notes + ----- + This method sends an asynchronous POST request to the pyLoad API endpoint + specified by `command`, with the provided `data` dictionary. It handles + authentication errors, HTTP status checks, and parses the JSON response. + + + Example + ------- + To add a new package to pyLoad, use: + ```python + status = await pyload_api.post(PyLoadCommand.ADD_PACKAGE, data={...} + ``` + + """ + url = f"{self.api_url}api/{command}" + data = { + k: str(v) if isinstance(v, bytes) else json.dumps(v) + for k, v in data.items() + } + + try: + async with self._session.post(url, data=data) as r: + _LOGGER.debug( + "Response from %s [%s]: %s", r.url, r.status, await r.text() + ) + + if r.status == HTTPStatus.UNAUTHORIZED: + raise InvalidAuth( + "Request failed due invalid or expired authentication cookie." + ) + r.raise_for_status() + try: + data = await r.json() + except JSONDecodeError as e: + _LOGGER.debug( + "Exception: Cannot parse response for %s:\n %s", + command, + traceback.format_exc(), + ) + raise ParserError( + f"Get {command} failed during parsing of request response." + ) from e + + return data + + except (TimeoutError, aiohttp.ClientError) as e: + _LOGGER.debug( + "Exception: Cannot execute command %s:\n %s", + command, + traceback.format_exc(), + ) + raise CannotConnect( + f"Executing command {command} failed due to request exception" + ) from e + async def get_status(self) -> StatusServerResponse: """Get general status information of pyLoad. @@ -529,3 +615,107 @@ async def free_space(self) -> int: return int(r) except CannotConnect as e: raise CannotConnect("Get free space failed due to request exception") from e + + async def add_package( + self, + name: str, + links: list[str], + destination: Destination = Destination.COLLECTOR, + ) -> int: + """Add a new package to pyLoad from a list of links. + + Parameters + ---------- + name : str + The name of the package to be added. + links : list[str] + A list of download links to be included in the package. + destination : Destination, optional + The destination where the package should be stored, by default Destination.COLLECTOR. + + Returns + ------- + int + The ID of the newly created package. + + Raises + ------ + CannotConnect + If the request to add the package fails due to a connection issue. + InvalidAuth + If the request fails due to invalid or expired authentication. + ParserError + If there's an issue parsing the response from the server. + + Example + ------- + To add a new package with a couple of links to the pyLoad collector: + ```python + package_id = await pyload_api.add_package( + "test_package", + [ + "https://example.com/file1.zip", + "https://example.com/file2.iso", + ] + ) + ``` + """ + + try: + r = await self.post( + PyLoadCommand.ADD_PACKAGE, + data={ + "name": name, + "links": links, + "dest": destination, + }, + ) + return int(r) + except CannotConnect as e: + raise CannotConnect("Adding package failed due to request exception") from e + + async def upload_container( + self, + filename: str, + binary_data: bytes, + ) -> None: + """Upload a container file to pyLoad. + + Parameters + ---------- + filename : str + The name of the file to be uploaded. + binary_data : bytes + The binary content of the file. + + Returns + ------- + None + + Raises + ------ + CannotConnect + If the request to upload the container fails due to a connection issue. + InvalidAuth + If the request fails due to invalid or expired authentication. + + Example + ------- + To upload a container file to pyLoad: + ```python + await pyload_api.upload_container( + "example_container.dlc", + b"binary data of the file" + ) + ``` + """ + try: + await self.post( + PyLoadCommand.UPLOAD_CONTAINER, + data={"filename": filename, "data": binary_data}, + ) + + except CannotConnect as e: + raise CannotConnect( + "Uploading container to pyLoad failed due to request exception" + ) from e diff --git a/src/pyloadapi/cli.py b/src/pyloadapi/cli.py index 99ae594..8a84624 100644 --- a/src/pyloadapi/cli.py +++ b/src/pyloadapi/cli.py @@ -10,8 +10,8 @@ import click from pyloadapi import CannotConnect, InvalidAuth, ParserError, PyLoadAPI +from pyloadapi.types import Destination -logging.basicConfig(level=logging.INFO) _LOGGER = logging.getLogger(__name__) CONFIG_FILE_PATH = Path("~/.config/pyloadapi/.pyload_config.json").expanduser() @@ -46,7 +46,7 @@ async def init_api( return api -@click.group() +@click.group(invoke_without_command=True) @click.option("--api-url", help="Base URL of pyLoad") @click.option("--username", help="Username for pyLoad") @click.option("--password", help="Password for pyLoad") @@ -59,6 +59,9 @@ def cli( ) -> None: """CLI for interacting with pyLoad.""" + if not any([api_url, username, password, ctx.invoked_subcommand]): + click.echo(ctx.get_help()) + config = load_config() if api_url: @@ -307,3 +310,95 @@ async def _toggle_reconnect() -> None: raise click.ClickException("Unable to parse response from pyLoad") from e asyncio.run(_toggle_reconnect()) + + +@cli.command() +@click.pass_context +@click.argument( + "container", + type=click.Path( + exists=True, + readable=True, + path_type=Path, + ), +) +def upload_container(ctx: click.Context, container: Path) -> None: + """Upload a container file to pyLoad.""" + + with open(container, "rb") as f: + binary_data = f.read() + + async def _upload_container() -> None: + try: + async with aiohttp.ClientSession() as session: + api = await init_api( + session, + ctx.obj["api_url"], + ctx.obj["username"], + ctx.obj["password"], + ) + + await api.upload_container( + filename=click.format_filename(container, shorten=True), + binary_data=binary_data, + ) + + except CannotConnect as e: + raise click.ClickException("Unable to connect to pyLoad") from e + except InvalidAuth as e: + raise click.ClickException( + "Authentication failed, verify username and password" + ) from e + except ParserError as e: + raise click.ClickException("Unable to parse response from pyLoad") from e + + asyncio.run(_upload_container()) + + +@cli.command() +@click.pass_context +@click.argument("package_name") +@click.option( + "--queue/--collector", + default=True, + help="Add package to queue or collector. Defaults to queue", +) +def add_package(ctx: click.Context, package_name: str, queue: bool) -> None: + """Add a package to pyLoad.""" + + links = [] + + while value := click.prompt( + "Please enter a link", type=str, default="", show_default=False + ): + links.append(value) + + if not links: + raise click.ClickException("No links entered") + + async def _add_package() -> None: + try: + async with aiohttp.ClientSession() as session: + api = await init_api( + session, + ctx.obj["api_url"], + ctx.obj["username"], + ctx.obj["password"], + ) + + await api.add_package( + name=package_name, + links=links, + destination=Destination.QUEUE if queue else Destination.COLLECTOR, + ) + + except CannotConnect as e: + raise click.ClickException("Unable to connect to pyLoad") from e + except InvalidAuth as e: + raise click.ClickException( + "Authentication failed, verify username and password" + ) from e + except ParserError as e: + raise click.ClickException("Unable to parse response from pyLoad") from e + + asyncio.run(_add_package()) diff --git a/src/pyloadapi/types.py b/src/pyloadapi/types.py index a95d680..7c31ff3 100644 --- a/src/pyloadapi/types.py +++ b/src/pyloadapi/types.py @@ -23,7 +23,7 @@ """ -from enum import StrEnum +from enum import IntEnum, StrEnum from typing import Any, NotRequired, TypedDict, TypeVar T = TypeVar("T") @@ -139,11 +139,12 @@ class PyLoadCommand(StrEnum): RESTART = "restart" VERSION = "getServerVersion" FREESPACE = "freeSpace" - ADDPACKAGE = "addPackage" + ADD_PACKAGE = "addPackage" + UPLOAD_CONTAINER = "uploadContainer" -class Destination(StrEnum): +class Destination(IntEnum): """Destination for new Packages.""" - QUEUE = "queue" - COLLECTOR = "collector" + COLLECTOR = 0 + QUEUE = 1 diff --git a/tests/conftest.py b/tests/conftest.py index 5bcee54..ddb393a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -48,6 +48,8 @@ "reconnect": False, } +BYTE_DATA = b"BYTE_DATA" + @pytest.fixture(name="session") async def aiohttp_client_session() -> AsyncGenerator[aiohttp.ClientSession, Any]: diff --git a/tests/test_api.py b/tests/test_api.py index 635ef94..616f92c 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -12,8 +12,10 @@ import pytest from pyloadapi import CannotConnect, InvalidAuth, ParserError, PyLoadAPI +from pyloadapi.types import Destination from .conftest import ( + BYTE_DATA, TEST_API_URL, TEST_LOGIN_RESPONSE, TEST_STATUS_RESPONSE, @@ -175,3 +177,154 @@ async def json(*args: Any) -> None: with pytest.raises(expected_exception=ParserError): await getattr(pyload, method)() + + +async def test_upload_container( + pyload: PyLoadAPI, mocked_aiohttp: aioresponses +) -> None: + """Test upload_container method.""" + + mocked_aiohttp.post( + f"{TEST_API_URL}api/uploadContainer", + ) + + await pyload.upload_container("filename.dlc", BYTE_DATA) + mocked_aiohttp.assert_called_once_with( + f"{TEST_API_URL}api/uploadContainer", + method="POST", + data={"filename": '"filename.dlc"', "data": "b'BYTE_DATA'"}, + ) + + +async def test_upload_container_exception( + pyload: PyLoadAPI, mocked_aiohttp: aioresponses +) -> None: + """Test upload_container exception.""" + + mocked_aiohttp.post( + f"{TEST_API_URL}api/uploadContainer", exception=aiohttp.ClientError + ) + + with pytest.raises(expected_exception=CannotConnect): + await pyload.upload_container("filename.dlc", BYTE_DATA) + + +async def test_upload_container_unauthorized( + pyload: PyLoadAPI, mocked_aiohttp: aioresponses +) -> None: + """Test upload_container authentication error.""" + + mocked_aiohttp.post( + f"{TEST_API_URL}api/uploadContainer", status=HTTPStatus.UNAUTHORIZED + ) + + with pytest.raises(expected_exception=InvalidAuth): + await pyload.upload_container("filename.dlc", BYTE_DATA) + + +async def test_upload_container_parse_exception( + pyload: PyLoadAPI, + mocked_aiohttp: aioresponses, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test upload_container parser error.""" + + mocked_aiohttp.post(re.compile(r".*"), status=HTTPStatus.OK) + + async def json(*args: Any) -> None: + raise JSONDecodeError("", "", 0) + + monkeypatch.setattr(aiohttp.ClientResponse, "json", json) + + with pytest.raises(expected_exception=ParserError): + await pyload.upload_container("filename.dlc", BYTE_DATA) + + +async def test_add_package(pyload: PyLoadAPI, mocked_aiohttp: aioresponses) -> None: + """Test add_package method.""" + + mocked_aiohttp.post( + f"{TEST_API_URL}api/addPackage", + payload=1, + ) + + await pyload.add_package( + name="Package Name", + links=[ + "https://example.com/file1.zip", + "https://example.com/file2.iso", + ], + destination=Destination.COLLECTOR, + ) + mocked_aiohttp.assert_called_once_with( + f"{TEST_API_URL}api/addPackage", + method="POST", + data={ + "name": '"Package Name"', + "links": '["https://example.com/file1.zip", "https://example.com/file2.iso"]', + "dest": "0", + }, + ) + + +async def test_add_package_exception( + pyload: PyLoadAPI, mocked_aiohttp: aioresponses +) -> None: + """Test add_package with exception.""" + + mocked_aiohttp.post( + f"{TEST_API_URL}api/addPackage", payload=1, exception=aiohttp.ClientError + ) + with pytest.raises(expected_exception=CannotConnect): + await pyload.add_package( + name="Package Name", + links=[ + "https://example.com/file1.zip", + "https://example.com/file2.iso", + ], + destination=Destination.COLLECTOR, + ) + + +async def test_add_package_unauthorized( + pyload: PyLoadAPI, mocked_aiohttp: aioresponses +) -> None: + """Test add_package authentication error.""" + + mocked_aiohttp.post( + f"{TEST_API_URL}api/addPackage", payload=1, status=HTTPStatus.UNAUTHORIZED + ) + with pytest.raises(expected_exception=InvalidAuth): + await pyload.add_package( + name="Package Name", + links=[ + "https://example.com/file1.zip", + "https://example.com/file2.iso", + ], + destination=Destination.COLLECTOR, + ) + + +async def test_add_package_parse_exception( + pyload: PyLoadAPI, + mocked_aiohttp: aioresponses, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test add_package parser error.""" + + mocked_aiohttp.post(re.compile(r".*"), status=HTTPStatus.OK) + + async def json(*args: Any) -> None: + raise JSONDecodeError("", "", 0) + + monkeypatch.setattr(aiohttp.ClientResponse, "json", json) + + with pytest.raises(expected_exception=ParserError): + await pyload.add_package( + name="Package Name", + links=[ + "https://example.com/file1.zip", + "https://example.com/file2.iso", + ], + destination=Destination.COLLECTOR, + ) diff --git a/tests/test_cli.py b/tests/test_cli.py index e1dd549..946bcdf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -9,7 +9,7 @@ from pyloadapi import CannotConnect, InvalidAuth, ParserError from pyloadapi.cli import cli -from .conftest import TEST_API_URL, TEST_PASSWORD, TEST_USERNAME +from .conftest import BYTE_DATA, TEST_API_URL, TEST_PASSWORD, TEST_USERNAME @pytest.mark.usefixtures("mock_pyloadapi") @@ -131,11 +131,12 @@ def test_exceptions( exception: Exception, msg: str, command: str, + tmp_path: Path, ) -> None: """Test status.""" runner = CliRunner() - with runner.isolated_filesystem(): + with runner.isolated_filesystem(temp_dir=tmp_path): mock_pyloadapi.get_status.side_effect = exception mock_pyloadapi.unpause.side_effect = exception mock_pyloadapi.pause.side_effect = exception @@ -153,3 +154,125 @@ def test_exceptions( ) assert result.exit_code == 1 assert result.output == msg + + +@pytest.mark.usefixtures("mock_pyloadapi") +def test_upload_container( + tmp_path: Path, +) -> None: + """Test status.""" + + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path) as tmp: + dlc = Path(tmp, "container.dlc") + with open(dlc, "wb") as f: + f.write(BYTE_DATA) + + result = runner.invoke( + cli, + args=f"--username {TEST_USERNAME} --password {TEST_PASSWORD} --api-url {TEST_API_URL} upload-container {dlc.as_posix()}", + ) + assert result.exit_code == 0 + + +@pytest.mark.parametrize( + ("exception", "msg"), + [ + (CannotConnect, "Error: Unable to connect to pyLoad\n"), + (InvalidAuth, "Error: Authentication failed, verify username and password\n"), + (ParserError, "Error: Unable to parse response from pyLoad\n"), + ], +) +def test_upload_container_exceptions( + mock_pyloadapi: MagicMock, + exception: Exception, + msg: str, + tmp_path: Path, +) -> None: + """Test status.""" + + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path) as tmp: + dlc = Path(tmp, "container.dlc") + with open(dlc, "wb") as f: + f.write(BYTE_DATA) + + mock_pyloadapi.upload_container.side_effect = exception + + result = runner.invoke( + cli, + args=f"--username {TEST_USERNAME} --password {TEST_PASSWORD} --api-url {TEST_API_URL} upload-container {dlc.as_posix()}", + ) + assert result.exit_code == 1 + assert result.output == msg + + +@pytest.mark.usefixtures("mock_pyloadapi") +def test_add_package(tmp_path: Path) -> None: + """Test add-package.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path): + result = runner.invoke( + cli, + args=f"--username {TEST_USERNAME} --password {TEST_PASSWORD} --api-url {TEST_API_URL} add-package Test-Package", + input=("http://example.com/file1.zip\n" "http://example.com/file2.iso\n\n"), + ) + assert result.exit_code == 0 + assert result.output == ( + "Please enter a link: http://example.com/file1.zip\n" + "Please enter a link: http://example.com/file2.iso\n" + "Please enter a link: \n" + ) + + +@pytest.mark.usefixtures("mock_pyloadapi") +def test_add_package_no_links(tmp_path: Path) -> None: + """Test add-package aborts if no links provided.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path): + result = runner.invoke( + cli, + args=f"--username {TEST_USERNAME} --password {TEST_PASSWORD} --api-url {TEST_API_URL} add-package Test-Package", + input="\n", + ) + assert result.exit_code == 1 + assert result.output == "Please enter a link: \nError: No links entered\n" + + +@pytest.mark.parametrize( + ("exception", "msg"), + [ + (CannotConnect, "Error: Unable to connect to pyLoad\n"), + (InvalidAuth, "Error: Authentication failed, verify username and password\n"), + (ParserError, "Error: Unable to parse response from pyLoad\n"), + ], +) +def test_add_package_exceptions( + mock_pyloadapi: MagicMock, + tmp_path: Path, + exception: Exception, + msg: str, +) -> None: + """Test add-package aborts if no links provided.""" + runner = CliRunner() + with runner.isolated_filesystem(temp_dir=tmp_path): + mock_pyloadapi.add_package.side_effect = exception + result = runner.invoke( + cli, + args=f"--username {TEST_USERNAME} --password {TEST_PASSWORD} --api-url {TEST_API_URL} add-package Test-Package", + input=("http://example.com/file1.zip\n" "http://example.com/file2.iso\n\n"), + ) + assert result.exit_code == 1 + assert msg in result.output + + +@pytest.mark.usefixtures("mock_pyloadapi") +def test_get_help() -> None: + """Test that invoking cli without commands and params returns help.""" + runner = CliRunner() + with runner.isolated_filesystem(): + result = runner.invoke( + cli, + ) + assert result.exit_code == 0 + assert "Usage: cli [OPTIONS] COMMAND [ARGS]..." in result.output