Allow regex machine models definitions

piotrkilczuk · Dec 5, 2016 · 30cbd3a · 30cbd3a
1 parent 21e9bbc
commit 30cbd3a
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 41 deletions.
diff --git a/dsd/machine.py b/dsd/machine.py
@@ -1,6 +1,8 @@
 import logging
+import re
 
 import raildriver
+import toolz
 import transitions
 
 from dsd import machine_models as models
@@ -19,42 +21,40 @@
 
 
 MODEL_MAPPING = {
-    'Default': models.GenericDSDModel,
-
-    'AP_Waggonz.Class90Pack': models.Class90DSDModel,  # TODO: make generic with patterndict
-    'AP_Waggonz.Class90Pack01': models.Class90DSDModel,
-    'AP_Waggonz.Class90Pack02': models.Class90DSDModel,
-    'AP_Waggonz.Class142Pack': models.Class142APDSDModel,
-
-    'DTG.Class378Pack01': models.Class378DSDModel,
-    'DTG.Class415Pack01': models.GenericDSDModel,
-
-    'JL.WHL': models.Class37WHLSouthDSDModel,  # TODO: make more specific with patterndict
-
-    'JustTrains.NL': models.Class43JT_47_DSDModel,  # TODO: make more specific with patterndict
-    'JustTrains.Voyager': models.Class220_221DSDModel,
-
-    'Kuju.RailSimulator': models.GenericDSDModel,  # TODO: make more specific with patterndict
-
-    'RailRight.Class40Blue': models.Class40DSDModel,  # TODO: make generic with patterndict
-    'RailRight.Class40Green': models.Class40DSDModel,  # TODO: make generic with patterndict
-
-    'RSC.BrightonMainLine': models.GenericDSDModel,  # TODO: make more specific with patterndict
-    'RSC.Class47Pack01': models.Class43JT_47_DSDModel,
-    'RSC.Class66Pack02': models.Class66APDSDModel,
-    'RSC.Class70Pack01': models.GenericDSDModel,
-    'RSC.Class325Pack01': models.Class325DSDModel,
-    'RSC.Class421Pack01': models.GenericDSDModel,
-    'RSC.Class421Pack02': models.GenericDSDModel,
-    'RSC.Class422Pack01': models.GenericDSDModel,
-    'RSC.Class423Pack01': models.GenericDSDModel,
-    'RSC.Class444Pack01': models.GenericDSDModel,
-    'RSC.Class465Pack01': models.Class465DSDModel,
-    'RSC.ECMLS': models.GenericDSDModel,
-    'RSC.GEML': models.Class360DSDModel,
-    'RSC.KentHighSpeed': models.Class395DSDModel,
-
-    'Thomson.Class455Pack01': models.GenericDSDModel,
+    'AP_Waggonz\.Class90Pack': models.Class90DSDModel,  # TODO: make generic with patterndict
+    'AP_Waggonz\.Class90Pack01': models.Class90DSDModel,
+    'AP_Waggonz\.Class90Pack02': models.Class90DSDModel,
+    'AP_Waggonz\.Class142Pack': models.Class142APDSDModel,
+
+    'DTG\.Class378Pack01': models.Class378DSDModel,
+    'DTG\.Class415Pack01': models.GenericDSDModel,
+
+    'JL\.WHL': models.Class37WHLSouthDSDModel,  # TODO: make more specific with patterndict
+
+    'JustTrains\.NL': models.Class43JT_47_DSDModel,  # TODO: make more specific with patterndict
+    'JustTrains\.Voyager': models.Class220_221DSDModel,
+
+    'Kuju\.RailSimulator': models.GenericDSDModel,  # TODO: make more specific with patterndict
+
+    'RailRight\.Class40Blue': models.Class40DSDModel,  # TODO: make generic with patterndict
+    'RailRight\.Class40Green': models.Class40DSDModel,  # TODO: make generic with patterndict
+
+    'RSC\.BrightonMainLine': models.GenericDSDModel,  # TODO: make more specific with patterndict
+    'RSC\.Class47Pack01': models.Class43JT_47_DSDModel,
+    'RSC\.Class66Pack02': models.Class66APDSDModel,
+    'RSC\.Class70Pack01': models.GenericDSDModel,
+    'RSC\.Class325Pack01': models.Class325DSDModel,
+    'RSC\.Class421Pack01': models.GenericDSDModel,
+    'RSC\.Class421Pack02': models.GenericDSDModel,
+    'RSC\.Class422Pack01': models.GenericDSDModel,
+    'RSC\.Class423Pack01': models.GenericDSDModel,
+    'RSC\.Class444Pack01': models.GenericDSDModel,
+    'RSC\.Class465Pack01': models.Class465DSDModel,
+    'RSC\.ECMLS': models.GenericDSDModel,
+    'RSC\.GEML': models.Class360DSDModel,
+    'RSC\.KentHighSpeed': models.Class395DSDModel,
+
+    'Thomson\.Class455Pack01': models.GenericDSDModel,
 }
 
 
@@ -131,7 +131,10 @@ def close(self, *args, **kwargs):
         self.usb.close()
 
     def init_model(self, loco_name):
-        model_class = MODEL_MAPPING.get('{}.{}'.format(*loco_name), MODEL_MAPPING['Default'])
+        model_matches = toolz.dicttoolz.keyfilter(
+            lambda k: re.compile('^{}'.format(k)).search('{}.{}.{}'.format(*loco_name)) is not None, MODEL_MAPPING)
+        model_class = model_matches.values()[0] if model_matches else models.GenericDSDModel
+
         model = model_class(self.beeper, self.raildriver, self.raildriver_listener, self.usb)
         logging.debug('Instantiated model {}'.format(repr(model)))
         super(DSDMachine, self).__init__(model,

diff --git a/requirements.txt b/requirements.txt
@@ -1,3 +1,4 @@
 py-raildriver==1.1.3
 pywinusb==0.4.1
 transitions==0.2.9
+toolz==0.8.0
diff --git a/tests.py b/tests.py
@@ -104,21 +104,21 @@ def test_initially_loco_explicit_model(self):
         """
         Class55DSDModel = type('Class55DSDModel', (dsd.machine.models.BaseDSDModel,), {})
         with mock.patch('dsd.machine.MODEL_MAPPING', {
-            'DTG.Class 55': Class55DSDModel,
-            'Default': dsd.machine.models.GenericDSDModel,
+            'DTG\.Class 55': Class55DSDModel,
         }):
             self.machine = dsd.DSDMachine()
             self.assertIsInstance(self.machine.model, Class55DSDModel)
             self.assertFalse(self.machine.needs_restart)
 
-    def test_initially_loco_default_model(self):
+    @mock.patch('dsd.machine.models.GenericDSDModel')
+    def test_initially_loco_default_model(self, mock_default_machine_model):
         """
         If initially there is already a loco active don't set needs_restart flag.
         Use the default model if explicit model is not available.
         """
         self.machine = dsd.DSDMachine()
         self.assertFalse(self.machine.needs_restart)
-        self.assertIsInstance(self.machine.model, dsd.MODEL_MAPPING['Default'])
+        mock_default_machine_model.assert_called_with(mock.ANY, mock.ANY, mock.ANY, mock.ANY)
 
     def test_initial_state_is_inactive(self):
         """