From 268c4495ddd728f09f5c958024076211ad9f0233 Mon Sep 17 00:00:00 2001
From: dexteradeus <dexteradeus@users.noreply.github.com>
Date: Wed, 20 Jun 2018 11:18:08 -0400
Subject: [PATCH] Add support for custom device handlers

Allow a custom device handler to be supplied to the manager via the
device_params dictionary.
---
 ncclient/manager.py       | 7 +++++++
 test/unit/test_manager.py | 8 ++++++++
 2 files changed, 15 insertions(+)

diff --git a/ncclient/manager.py b/ncclient/manager.py
index 302edcda..dd307f50 100644
--- a/ncclient/manager.py
+++ b/ncclient/manager.py
@@ -71,6 +71,10 @@ def make_device_handler(device_params):
     if device_params is None:
         device_params = {}
 
+    handler = device_params.get('handler', None)
+    if handler:
+        return handler(device_params)
+
     device_name = device_params.get("name", "default")
     # Attempt to import device handler class. All device handlers are
     # in a module called "ncclient.devices.<devicename>" and in a class named
@@ -98,6 +102,9 @@ def connect_ssh(*args, **kwds):
     To invoke advanced vendor related operation add device_params =
         {'name':'<vendor_alias>'} in connection paramerers. For the time,
         'junos' and 'nexus' are supported for Juniper and Cisco Nexus respectively.
+
+    A custom device handler can be provided with device_params =
+        {'handler':<handler class>} in connection paramerers.
     """
     # Extract device parameter dict, if it was passed into this function. Need to
     # remove it from kwds, since the session.connect() doesn't like extra stuff in
diff --git a/test/unit/test_manager.py b/test/unit/test_manager.py
index 8ba395b1..f69988fb 100644
--- a/test/unit/test_manager.py
+++ b/test/unit/test_manager.py
@@ -1,6 +1,7 @@
 import unittest
 from mock import patch, MagicMock
 from ncclient import manager
+from ncclient.devices.junos import JunosDeviceHandler
 
 
 class TestManager(unittest.TestCase):
@@ -66,6 +67,13 @@ def test_make_device_handler(self):
             device_handler.__class__.__name__,
             "JunosDeviceHandler")
 
+    def test_make_device_handler_provided_handler(self):
+        device_handler = manager.make_device_handler(
+            {'handler': JunosDeviceHandler})
+        self.assertEqual(
+            device_handler.__class__.__name__,
+            "JunosDeviceHandler")
+
     @patch('ncclient.operations.LockContext')
     def test_manager_locked(self, mock_lock):
         conn = manager.Manager(None, None, timeout=20)