Add support for custom PDU outlet count

The number of outlet may now be specified at the constructor call of the
PDU class.  If defined, it will override the PDU class specific outlet
count.

Also, the SNMPPDUHarness.stop method was not protected with a lock.
Without this, the stop method simply fails to stop the thread if called
just after the start method.

Reviewed-By: Maxime Dupuis <mdupuis@internap.com>

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]))