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.

Author: inetAnt

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.
 

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",
+]

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

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

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:

tests/unit/mocked_data/control_outlets_AA01.json

@@ -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")

tests/unit/mocked_data/restart.json

@@ -8,3 +8,6 @@
 @pytest.mark.usefixtures("set_device_parameters")
 class TestGetter(BaseTestGetters):
     """Test get_* methods."""
+
+    def test_method_signatures(self):
+        return True

tests/unit/test_add.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)