Merge branch '2.0' of github.com:nornir-automation/nornir into pydantic

nornir/core/__init__.py

@@ -1,12 +1,12 @@
 import logging
 import logging.config
 from multiprocessing.dummy import Pool
-from typing import Type
 
 from nornir.core.configuration import Config
-from nornir.core.connections import ConnectionPlugin
 from nornir.core.task import AggregatedResult, Task
-from nornir.plugins import connections
+from nornir.plugins.connections import register_default_connection_plugins
+
+register_default_connection_plugins()
 
 
 class Data(object):
@@ -16,12 +16,10 @@ class Data(object):
 
     Attributes:
         failed_hosts (list): Hosts that have failed to run a task properly
-        available_connections (dict): Dictionary holding available connection plugins
     """
 
     def __init__(self):
         self.failed_hosts = set()
-        self.available_connections = connections.available_connections
 
     def recover_host(self, host):
         """Remove ``host`` from list of failed hosts."""
@@ -47,8 +45,6 @@ class Nornir(object):
         dry_run(``bool``): Whether if we are testing the changes or not
         config (:obj:`nornir.core.configuration.Config`): Configuration object
         config_file (``str``): Path to Yaml configuration file
-        available_connections (``dict``): dict of connection types that will be made available.
-            Defaults to :obj:`nornir.plugins.tasks.connections.available_connections`
 
     Attributes:
         inventory (:obj:`nornir.core.inventory.Inventory`): Inventory to work with
@@ -58,14 +54,7 @@ class Nornir(object):
     """
 
     def __init__(
-        self,
-        inventory,
-        dry_run,
-        config=None,
-        config_file=None,
-        available_connections=None,
-        logger=None,
-        data=None,
+        self, inventory, dry_run, config=None, config_file=None, logger=None, data=None
     ):
         self.logger = logger or logging.getLogger("nornir")
 
@@ -79,9 +68,6 @@ def __init__(
         else:
             self.config = config or Config()
 
-        if available_connections is not None:
-            self.data.available_connections = available_connections
-
     def __enter__(self):
         return self
 
@@ -189,10 +175,6 @@ def to_dict(self):
         """ Return a dictionary representing the object. """
         return {"data": self.data.to_dict(), "inventory": self.inventory.to_dict()}
 
-    def get_connection_type(self, connection: str) -> Type[ConnectionPlugin]:
-        """Returns the class for the given connection type."""
-        return self.data.available_connections[connection]
-
     def close_connections(self, on_good=True, on_failed=False):
         def close_connections_task(task):
             task.host.close_connections()

nornir/core/connections.py

@@ -1,8 +1,12 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, NoReturn, Optional
+from typing import Any, Dict, NoReturn, Optional, Type
 
 
 from nornir.core.configuration import Config
+from nornir.core.exceptions import (
+    ConnectionPluginAlreadyRegistered,
+    ConnectionPluginNotRegistered,
+)
 
 
 class ConnectionPlugin(ABC):
@@ -53,4 +57,63 @@ def close(self) -> NoReturn:
 
 
 class Connections(Dict[str, ConnectionPlugin]):
-    pass
+    available: Dict[str, Type[ConnectionPlugin]] = {}
+
+    @classmethod
+    def register(cls, name: str, plugin: Type[ConnectionPlugin]) -> None:
+        """Registers a connection plugin with a specified name
+
+        Args:
+            name: name of the connection plugin to register
+            plugin: defined connection plugin class
+
+        Raises:
+            :obj:`nornir.core.exceptions.ConnectionPluginAlreadyRegistered` if
+                another plugin with the specified name was already registered
+        """
+        existing_plugin = cls.available.get(name)
+        if existing_plugin is None:
+            cls.available[name] = plugin
+        elif existing_plugin != plugin:
+            raise ConnectionPluginAlreadyRegistered(
+                f"Connection plugin {plugin.__name__} can't be registered as "
+                f"{name!r} because plugin {existing_plugin.__name__} "
+                f"was already registered under this name"
+            )
+
+    @classmethod
+    def deregister(cls, name: str) -> None:
+        """Deregisters a registered connection plugin by its name
+
+        Args:
+            name: name of the connection plugin to deregister
+
+        Raises:
+            :obj:`nornir.core.exceptions.ConnectionPluginNotRegistered`
+        """
+        if name not in cls.available:
+            raise ConnectionPluginNotRegistered(
+                f"Connection {name!r} is not registered"
+            )
+        cls.available.pop(name)
+
+    @classmethod
+    def deregister_all(cls) -> None:
+        """Deregisters all registered connection plugins"""
+        cls.available = {}
+
+    @classmethod
+    def get_plugin(cls, name: str) -> Type[ConnectionPlugin]:
+        """Fetches the connection plugin by name if already registered
+
+        Args:
+            name: name of the connection plugin
+
+        Raises:
+            :obj:`nornir.core.exceptions.ConnectionPluginNotRegistered`
+        """
+        if name not in cls.available:
+            raise ConnectionPluginNotRegistered(
+                f"Connection {name!r} is not registered"
+            )
+        return cls.available[name]

nornir/core/exceptions.py

@@ -14,6 +14,14 @@ class ConnectionNotOpen(ConnectionException):
     pass
 
 
+class ConnectionPluginAlreadyRegistered(ConnectionException):
+    pass
+
+
+class ConnectionPluginNotRegistered(ConnectionException):
+    pass
+
+
 class CommandError(Exception):
     """
     Raised when there is a command error.

nornir/core/inventory.py

@@ -329,7 +329,7 @@ def open_connection(
         if connection in self.connections:
             raise ConnectionAlreadyOpen(connection)
 
-        self.connections[connection] = self.nornir.get_connection_type(connection)()
+        self.connections[connection] = self.connections.get_plugin(connection)()
         if default_to_host_attributes:
             conn_params = self.get_connection_parameters(connection)
             self.connections[connection].open(

nornir/plugins/connections/__init__.py

@@ -1,16 +1,10 @@
-from typing import Dict, TYPE_CHECKING, Type
-
-
 from .napalm import Napalm
 from .netmiko import Netmiko
 from .paramiko import Paramiko
-
-if TYPE_CHECKING:
-    from nornir.core.connections import ConnectionPlugin  # noqa
+from nornir.core.connections import Connections
 
 
-available_connections: Dict[str, Type["ConnectionPlugin"]] = {
-    "napalm": Napalm,
-    "netmiko": Netmiko,
-    "paramiko": Paramiko,
-}
+def register_default_connection_plugins() -> None:
+    Connections.register("napalm", Napalm)
+    Connections.register("netmiko", Netmiko)
+    Connections.register("paramiko", Paramiko)

tests/core/test_connections.py

@@ -1,8 +1,16 @@
 from typing import Any, Dict, Optional
 
+import pytest
+
 from nornir.core.configuration import Config
-from nornir.core.connections import ConnectionPlugin
-from nornir.core.exceptions import ConnectionAlreadyOpen, ConnectionNotOpen
+from nornir.core.connections import ConnectionPlugin, Connections
+from nornir.core.exceptions import (
+    ConnectionAlreadyOpen,
+    ConnectionNotOpen,
+    ConnectionPluginNotRegistered,
+    ConnectionPluginAlreadyRegistered,
+)
+from nornir.plugins.connections import register_default_connection_plugins
 
 
 class DummyConnectionPlugin(ConnectionPlugin):
@@ -30,6 +38,10 @@ def close(self) -> None:
         self.connection = False
 
 
+class AnotherDummyConnectionPlugin(DummyConnectionPlugin):
+    pass
+
+
 def open_and_close_connection(task):
     task.host.open_connection("dummy")
     assert "dummy" in task.host.connections
@@ -68,37 +80,38 @@ def validate_params(task, conn, params):
 
 
 class Test(object):
+    @classmethod
+    def setup_class(cls):
+        Connections.deregister_all()
+        Connections.register("dummy", DummyConnectionPlugin)
+        Connections.register("dummy_no_overrides", DummyConnectionPlugin)
+
     def test_open_and_close_connection(self, nornir):
-        nornir.data.available_connections["dummy"] = DummyConnectionPlugin
         nr = nornir.filter(name="dev2.group_1")
         r = nr.run(task=open_and_close_connection, num_workers=1)
         assert len(r) == 1
         assert not r.failed
 
     def test_open_connection_twice(self, nornir):
-        nornir.data.available_connections["dummy"] = DummyConnectionPlugin
         nr = nornir.filter(name="dev2.group_1")
         r = nr.run(task=open_connection_twice, num_workers=1)
         assert len(r) == 1
         assert not r.failed
 
     def test_close_not_opened_connection(self, nornir):
-        nornir.data.available_connections["dummy"] = DummyConnectionPlugin
         nr = nornir.filter(name="dev2.group_1")
         r = nr.run(task=close_not_opened_connection, num_workers=1)
         assert len(r) == 1
         assert not r.failed
 
     def test_context_manager(self, nornir):
-        nornir.data.available_connections["dummy"] = DummyConnectionPlugin
         with nornir.filter(name="dev2.group_1") as nr:
             nr.run(task=a_task)
             assert "dummy" in nr.inventory.hosts["dev2.group_1"].connections
         assert "dummy" not in nr.inventory.hosts["dev2.group_1"].connections
         nornir.data.reset_failed_hosts()
 
     def test_validate_params_simple(self, nornir):
-        nornir.data.available_connections["dummy_no_overrides"] = DummyConnectionPlugin
         params = {
             "hostname": "127.0.0.1",
             "username": "root",
@@ -118,7 +131,6 @@ def test_validate_params_simple(self, nornir):
         assert not r.failed
 
     def test_validate_params_overrides(self, nornir):
-        nornir.data.available_connections["dummy"] = DummyConnectionPlugin
         params = {
             "hostname": "overriden_hostname",
             "username": "root",
@@ -131,3 +143,51 @@ def test_validate_params_overrides(self, nornir):
         r = nr.run(task=validate_params, conn="dummy", params=params, num_workers=1)
         assert len(r) == 1
         assert not r.failed
+
+
+class TestConnectionPluginsRegistration(object):
+    def setup_method(self, method):
+        Connections.deregister_all()
+        Connections.register("dummy", DummyConnectionPlugin)
+        Connections.register("another_dummy", AnotherDummyConnectionPlugin)
+
+    def teardown_method(self, method):
+        Connections.deregister_all()
+        register_default_connection_plugins()
+
+    def test_count(self):
+        assert len(Connections.available) == 2
+
+    def test_register_new(self):
+        Connections.register("new_dummy", DummyConnectionPlugin)
+        assert "new_dummy" in Connections.available
+
+    def test_register_already_registered_same(self):
+        Connections.register("dummy", DummyConnectionPlugin)
+        assert Connections.available["dummy"] == DummyConnectionPlugin
+
+    def test_register_already_registered_new(self):
+        with pytest.raises(ConnectionPluginAlreadyRegistered):
+            Connections.register("dummy", AnotherDummyConnectionPlugin)
+
+    def test_deregister_existing(self):
+        Connections.deregister("dummy")
+        assert len(Connections.available) == 1
+        assert "dummy" not in Connections.available
+
+    def test_deregister_nonexistent(self):
+        with pytest.raises(ConnectionPluginNotRegistered):
+            Connections.deregister("nonexistent_dummy")
+
+    def test_deregister_all(self):
+        Connections.deregister_all()
+        assert Connections.available == {}
+
+    def test_get_plugin(self):
+        assert Connections.get_plugin("dummy") == DummyConnectionPlugin
+        assert Connections.get_plugin("another_dummy") == AnotherDummyConnectionPlugin
+        assert len(Connections.available) == 2
+
+    def test_nonexistent_plugin(self):
+        with pytest.raises(ConnectionPluginNotRegistered):
+            Connections.get_plugin("nonexistent_dummy")