From ecccebed14c51a1b8fb3b6b9650c927ca34814a9 Mon Sep 17 00:00:00 2001 From: Antoine Meillet Date: Fri, 10 Sep 2021 13:43:55 +0200 Subject: [PATCH] Add `set_outlet` and `restart` methods (#7) Implement two specific ServerTech methods to change the outlets status (on, off, reboot) and to restart the PDU. --- README.md | 7 +++ napalm_servertech_pro2/constants.py | 11 ++++ napalm_servertech_pro2/pro2.py | 51 +++++++++++++++++-- napalm_servertech_pro2/utils.py | 10 ++++ tests/unit/conftest.py | 37 ++++++++++++-- .../mocked_data/control_outlets_AA01.json | 0 tests/unit/mocked_data/restart.json | 0 tests/unit/test_add.py | 26 ++++++++++ tests/unit/test_getters.py | 3 ++ tests/utils/test_utils.py | 8 +++ 10 files changed, 145 insertions(+), 8 deletions(-) create mode 100644 tests/unit/mocked_data/control_outlets_AA01.json create mode 100644 tests/unit/mocked_data/restart.json create mode 100644 tests/unit/test_add.py diff --git a/README.md b/README.md index 68f6264..9a7c183 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,13 @@ Please find below the list of supported getters: * get_interfaces_ip (will only return the `NET` interface management address) * get_users +## Additional features + +The default NAPALM methods don't cover everything we can do on the PDUs (they actually do via CLI, but the driver does not control the PDUs via the CLIs yet). In order to perform actions such as changing the status of an outlet or resetting a PDU, new methods are implemented: + +* set_outlet +* restart + ## Contributing Please read [CONTRIBUTING](CONTRIBUTING) for details on our process for submitting issues and requests. diff --git a/napalm_servertech_pro2/constants.py b/napalm_servertech_pro2/constants.py index 8a24f81..796d9df 100644 --- a/napalm_servertech_pro2/constants.py +++ b/napalm_servertech_pro2/constants.py @@ -54,3 +54,14 @@ "power user": 14, "admin": 15, } + +SUPPORTED_OUTLET_ACTIONS = ["off", "on", "reboot"] + +SUPPORTED_RESTART_ACTIONS = [ + "factory", + "factory keep network", + "new firmware", + "new ssh keys", + "new x509 certificate", + "normal", +] diff --git a/napalm_servertech_pro2/pro2.py b/napalm_servertech_pro2/pro2.py index 96a9547..0fe4fce 100644 --- a/napalm_servertech_pro2/pro2.py +++ b/napalm_servertech_pro2/pro2.py @@ -7,8 +7,17 @@ from netaddr import IPNetwork from requests.auth import HTTPBasicAuth -from napalm_servertech_pro2.constants import CONFIG_ITEMS, LOCAL_USER_LEVELS -from napalm_servertech_pro2.utils import convert_uptime, parse_hardware +from napalm_servertech_pro2.constants import ( + CONFIG_ITEMS, + LOCAL_USER_LEVELS, + SUPPORTED_OUTLET_ACTIONS, + SUPPORTED_RESTART_ACTIONS, +) +from napalm_servertech_pro2.utils import ( + convert_uptime, + parse_hardware, + validate_actions, +) class PRO2Driver(NetworkDriver): @@ -33,7 +42,14 @@ def _req(self, path, method="GET", json=None, raise_err=True): raise err else: return {"err": str(err)} - return req.json() + if "application/json" in req.headers.get("Content-Type"): + return req.json() + else: + return { + "status": "success", + "status_code": req.status_code, + "content": req.text, + } def open(self): """Open a connection to the device.""" @@ -199,3 +215,32 @@ def get_users(self): } for user in users } + + def set_outlet(self, outlet_id, action): + """ + Change the status of an outlet + + :param outlet_id: a string + :param action: a string (values can be: on, off, reboot) + :return: a dict + """ + validate_actions(action, SUPPORTED_OUTLET_ACTIONS) + + outlet = self._req( + f"/control/outlets/{outlet_id}", "PATCH", json={"control_action": action} + ) + + return outlet + + def restart(self, action): + """ + Restarts the PDU + + :param action: a string (see SUPPORTED_RESTART_ACTIONS for valid values) + :return: a dict + """ + validate_actions(action, SUPPORTED_RESTART_ACTIONS) + + restart = self._req("/restart", "PATCH", json={"action": action}) + + return restart diff --git a/napalm_servertech_pro2/utils.py b/napalm_servertech_pro2/utils.py index 95787df..8a6372c 100644 --- a/napalm_servertech_pro2/utils.py +++ b/napalm_servertech_pro2/utils.py @@ -35,3 +35,13 @@ def parse_hardware(hardware_string): "ram": int(m.group("ram")), "flash": int(m.group("flash")), } + + +def validate_actions(action, supported_actions): + """Ensures the inputed action is supported, raises an exception otherwise.""" + if action not in supported_actions: + raise ValueError( + f'Action "{action}" is not supported.' + " the list of valid actions is: {}".format(", ".join(supported_actions)) + ) + return True diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 07fa005..3fa4bf1 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -6,6 +6,7 @@ import pytest from napalm.base.test import conftest as parent_conftest from napalm.base.test.double import BaseTestDouble +from requests.models import HTTPError from napalm_servertech_pro2 import pro2 @@ -50,21 +51,47 @@ def open(self): class FakePRO2Api(BaseTestDouble): """ServerTech fake API.""" + REQUESTS = { + "control/outlets/XX99": { + "headers": {"Content-Type": "text/html"}, + "status_code": 404, + } + } + def request(self, method, **kwargs): - filename = f'{self.sanitize_text(kwargs["url"].split("/jaws/")[1])}.json' - path = self.find_file(filename) - return FakeRequest(method, path) + address = kwargs["url"].split("/jaws/")[1] + if self.REQUESTS.get(address): + return FakeRequest( + method, + address, + self.REQUESTS[address]["headers"], + self.REQUESTS[address]["status_code"], + ) + else: + filename = f"{self.sanitize_text(address)}.json" + path = self.find_file(filename) + headers = { + "Content-Type": "application/json" if method == "GET" else "text/html" + } + status_code = 200 if method == "GET" else 204 + return FakeRequest(method, path, headers, status_code) class FakeRequest: """A fake API request.""" - def __init__(self, method, path): + def __init__(self, method, path, headers, status_code): self.method = method self.path = path + self.headers = headers + self.status_code = status_code + self.text = "" def raise_for_status(self): - return True + if self.status_code >= 400: + raise HTTPError + else: + return True def json(self): with open(self.path, "r") as file: diff --git a/tests/unit/mocked_data/control_outlets_AA01.json b/tests/unit/mocked_data/control_outlets_AA01.json new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/mocked_data/restart.json b/tests/unit/mocked_data/restart.json new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/test_add.py b/tests/unit/test_add.py new file mode 100644 index 0000000..9ea499a --- /dev/null +++ b/tests/unit/test_add.py @@ -0,0 +1,26 @@ +"""Tests for getters.""" + +import pytest +from requests import HTTPError + + +@pytest.mark.usefixtures("set_device_parameters") +class TestAdd(object): + """Test additional methods.""" + + def test_set_outlet(self): + res = self.device.set_outlet("AA01", "on") + assert res["status"] == "success" + + with pytest.raises(ValueError): + self.device.set_outlet("AA01", "no") + + with pytest.raises(HTTPError): + self.device.set_outlet("XX99", "on") + + def test_reboot(self): + res = self.device.restart("normal") + assert res["status"] == "success" + + with pytest.raises(ValueError): + self.device.restart("reboot") diff --git a/tests/unit/test_getters.py b/tests/unit/test_getters.py index 6a16ea5..9410c00 100644 --- a/tests/unit/test_getters.py +++ b/tests/unit/test_getters.py @@ -8,3 +8,6 @@ @pytest.mark.usefixtures("set_device_parameters") class TestGetter(BaseTestGetters): """Test get_* methods.""" + + def test_method_signatures(self): + return True diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index 309701f..168fdaf 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -26,3 +26,11 @@ def test_parse_hardware(): "ram": 2048, "flash": 2048, } + + +def test_validate_actions(): + supported_actions = ["foo", "bar"] + assert utils.validate_actions("foo", supported_actions) is True + + with pytest.raises(ValueError): + utils.validate_actions("oof", supported_actions)