Merge pull request #41 from JoProvost/pdu_size_support

Add support for custom PDU outlet count

virtualpdu/pdu/__init__.py

@@ -130,7 +130,9 @@ class PDU(object):
     outlet_features = [PDUOutletControl]
     general_features = []
 
-    def __init__(self, name, core):
+    def __init__(self, name, core, outlet_count=None):
+        if outlet_count is not None:
+            self.outlet_count = outlet_count
         self.name = name
 
         mapping = {}

virtualpdu/pdu/apc_rackpdu.py

@@ -98,7 +98,7 @@ class APCRackPDUOutletConfigIndex(PDUOutletFeature):
 
 
 class APCRackPDU(PDU):
-    outlet_count = 8
+    outlet_count = 24
     outlet_index_start = 1
     outlet_features = [APCRackPDUOutletControl, APCRackPDUOutletName,
                        APCRackPDUOutletConfigIndex, APCRackPDUOutletState]

virtualpdu/pdu/pysnmp_handler.py

@@ -127,21 +127,28 @@ class SNMPPDUHarness(threading.Thread):
         self.listen_port = listen_port
         self.transportDispatcher = AsyncoreDispatcher()
 
+        self._lock = threading.Lock()
+        self._stop_requested = False
+
     def run(self):
-        self.logger.info("Starting PDU '{}' on {}:{}".format(
-            self.pdu.name, self.listen_address, self.listen_port)
-        )
-        self.transportDispatcher.registerRecvCbFun(
-            self.snmp_handler.message_handler)
+        with self._lock:
+            if self._stop_requested:
+                return
 
-        # UDP/IPv4
-        self.transportDispatcher.registerTransport(
-            udp.domainName,
-            udp.UdpSocketTransport().openServerMode(
-                (self.listen_address, self.listen_port))
-        )
+            self.logger.info("Starting PDU '{}' on {}:{}".format(
+                self.pdu.name, self.listen_address, self.listen_port)
+            )
+            self.transportDispatcher.registerRecvCbFun(
+                self.snmp_handler.message_handler)
 
-        self.transportDispatcher.jobStarted(1)
+            # UDP/IPv4
+            self.transportDispatcher.registerTransport(
+                udp.domainName,
+                udp.UdpSocketTransport().openServerMode(
+                    (self.listen_address, self.listen_port))
+            )
+
+            self.transportDispatcher.jobStarted(1)
 
         try:
             # Dispatcher will never finish as job#1 never reaches zero
@@ -150,4 +157,9 @@ class SNMPPDUHarness(threading.Thread):
             self.transportDispatcher.closeDispatcher()
 
     def stop(self):
-        self.transportDispatcher.jobFinished(1)
+        with self._lock:
+            self._stop_requested = True
+            try:
+                self.transportDispatcher.jobFinished(1)
+            except KeyError:
+                pass  # The job is not started yet and will not start

virtualpdu/tests/integration/pdu/test_pysnmp_handler.py

@@ -105,3 +105,17 @@ class TestSNMPPDUHarness(base.TestCase):
         self.assertEqual(42, val)
 
         harness.stop()
+
+    def test_start_stop_threadsafety(self):
+        mock_pdu = mock.Mock()
+        port = randint(20000, 30000)
+        harness = pysnmp_handler.SNMPPDUHarness(pdu=mock_pdu,
+                                                listen_address='127.0.0.1',
+                                                listen_port=port,
+                                                community='bleh')
+
+        harness.start()
+        harness.stop()
+        harness.join(timeout=5)
+
+        self.assertFalse(harness.isAlive())

virtualpdu/tests/unit/pdu/test_apc_rackpdu.py

@@ -22,6 +22,7 @@ from virtualpdu.tests.unit.pdu.base_pdu_test_cases import BasePDUTests
 
 class TestAPCRackPDU(base.TestCase, BasePDUTests):
     pdu_class = apc_rackpdu.APCRackPDU
+    outlet_control_class = apc_rackpdu.APCRackPDUOutletControl
     outlet_control_oid = \
         apc_rackpdu.rPDU_outlet_control_outlet_command \
         + (apc_rackpdu.APCRackPDU.outlet_index_start,)
@@ -100,3 +101,17 @@ class TestAPCRackPDU(base.TestCase, BasePDUTests):
         self.core_mock.get_pdu_outlet_state.assert_called_with(
             pdu='my_pdu',
             outlet=1)
+
+    def test_as_many_outlets_as_specified_by_constructor(self):
+        pdu = self.pdu_class(name='test_pdu',
+                             core=self.core_mock,
+                             outlet_count=10)
+        self.assertEqual(10, pdu.outlet_count)
+        self.assertEqual(10, len([oid for oid in pdu.oid_mapping.values()
+                                  if type(oid) is self.outlet_control_class]))
+
+    def test_as_many_outlets_as_specified_by_type(self):
+        self.assertEqual(self.pdu_class.outlet_count, self.pdu.outlet_count)
+        self.assertEqual(self.pdu_class.outlet_count,
+                         len([oid for oid in self.pdu.oid_mapping.values()
+                              if type(oid) is self.outlet_control_class]))

virtualpdu/tests/unit/pdu/test_baytech_mrp27.py

@@ -21,6 +21,7 @@ from virtualpdu.tests.unit.pdu.base_pdu_test_cases import BasePDUTests
 
 class TestBaytechMRP27PDU(base.TestCase, BasePDUTests):
     pdu_class = baytech_mrp27.BaytechMRP27PDU
+    outlet_control_class = baytech_mrp27.BaytechMRP27PDUOutletControl
     outlet_control_oid = \
         baytech_mrp27.sBTA_modules_RPC_outlet_state \
         + (1, baytech_mrp27.BaytechMRP27PDU.outlet_index_start,)
@@ -36,3 +37,17 @@ class TestBaytechMRP27PDU(base.TestCase, BasePDUTests):
             univ.ObjectIdentifier(baytech_mrp27.sBTA),
             self.pdu.oid_mapping[sysObjectID].value
         )
+
+    def test_as_many_outlets_as_specified_by_constructor(self):
+        pdu = self.pdu_class(name='test_pdu',
+                             core=self.core_mock,
+                             outlet_count=10)
+        self.assertEqual(10, pdu.outlet_count)
+        self.assertEqual(10, len([oid for oid in pdu.oid_mapping.values()
+                                  if type(oid) is self.outlet_control_class]))
+
+    def test_as_many_outlets_as_specified_by_type(self):
+        self.assertEqual(self.pdu_class.outlet_count, self.pdu.outlet_count)
+        self.assertEqual(self.pdu_class.outlet_count,
+                         len([oid for oid in self.pdu.oid_mapping.values()
+                              if type(oid) is self.outlet_control_class]))