diff --git a/dsd/machine.py b/dsd/machine.py index 0cc2bc9..3f081eb 100644 --- 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 index b92d276..7fcbdeb 100644 --- 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 index ce49467..96fee53 100644 --- 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): """