fix: Preflights are loaded with their runtime name (log2timeline#361)

moettle · Nov 11, 2020 · 1ab0fa0 · 1ab0fa0
1 parent 26f321b
commit 1ab0fa0
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 4 deletions.
diff --git a/dftimewolf/lib/state.py b/dftimewolf/lib/state.py
@@ -253,11 +253,13 @@ def RunPreflights(self):
  """Runs preflight modules."""
  for preflight_definition in self.recipe.get('preflights', []):
  preflight_name = preflight_definition['name']
+ runtime_name = preflight_definition.get('runtime_name', preflight_name)
+
  args = preflight_definition.get('args', {})
 
  new_args = utils.ImportArgsFromDict(
  args, self.command_line_options, self.config)
- preflight = self._module_pool[preflight_name]
+ preflight = self._module_pool[runtime_name]
  try:
  preflight.SetUp(**new_args)
  preflight.Process()
@@ -268,7 +270,8 @@ def CleanUpPreflights(self):
  """Executes any cleanup actions defined in preflight modules."""
  for preflight_definition in self.recipe.get('preflights', []):
  preflight_name = preflight_definition['name']
- preflight = self._module_pool[preflight_name]
+ runtime_name = preflight_definition.get('runtime_name', preflight_name)
+ preflight = self._module_pool[runtime_name]
  try:
  preflight.CleanUp()
  finally:

diff --git a/tests/lib/state.py b/tests/lib/state.py
@@ -59,7 +59,7 @@ def testLoadRecipeWithRuntimeNames(self):
  self.assertIn('DummyModule2', test_state._module_pool)
  self.assertIn('DummyModule1-2', test_state._module_pool)
  self.assertIn('DummyModule2-2', test_state._module_pool)
- self.assertIn('DummyPreflightModule', test_state._module_pool)
+ self.assertIn('DummyPreflightModule-runtime', test_state._module_pool)
  self.assertEqual(len(test_state._module_pool), 5)
 
  def testStoreContainer(self):
@@ -92,6 +92,17 @@ def testProcessPreflightModules(self, mock_setup, mock_process):
  mock_setup.assert_called_with()
  mock_process.assert_called_with()
 
+ @mock.patch('tests.test_modules.modules.DummyPreflightModule.Process')
+ @mock.patch('tests.test_modules.modules.DummyPreflightModule.SetUp')
+ def testProcessNamedPreflightModules(self, mock_setup, mock_process):
+ """Tests that preflight's process function is called correctly."""
+ test_state = state.DFTimewolfState(config.Config)
+ test_state.command_line_options = {}
+ test_state.LoadRecipe(test_recipe.named_modules_contents)
+ test_state.RunPreflights()
+ mock_setup.assert_called_with()
+ mock_process.assert_called_with()
+
  @mock.patch('tests.test_modules.modules.DummyPreflightModule.CleanUp')
  def testCleanupPreflightModules(self, mock_cleanup):
  """Tests that preflight's process function is called correctly."""

diff --git a/tests/test_modules/test_recipe.py b/tests/test_modules/test_recipe.py
@@ -24,7 +24,8 @@
  'dummy_recipe',
  'short_description': 'Nothing to see here.',
  'preflights': [{
- 'name': 'DummyPreflightModule'
+ 'name': 'DummyPreflightModule',
+ 'runtime_name': 'DummyPreflightModule-runtime'
  }],
  'modules': [{
  'wants': [],