Make sure that we can make a copy of dataset (#20)

Author: t-persson

src/environment_provider/iut/prepare.py

@@ -39,6 +39,10 @@ def __init__(self, jsontas, prepare_ruleset):
         self.prepare_ruleset = prepare_ruleset
         self.jsontas = jsontas
         self.dataset = self.jsontas.dataset
+        # Pop the config as it contains values that are not pickleable.
+        # Pickle is used by deepcopy inside of 'dataset.copy' which is required
+        # during the preparation step.
+        self.config = self.dataset._Dataset__dataset.pop("config")
 
     def execute_preparation_steps(self, iut, preparation_steps):
         """Execute the preparation steps for the environment provider on an IUT.
@@ -48,21 +52,27 @@ def execute_preparation_steps(self, iut, preparation_steps):
         :param preparation_steps: Steps to execute to prepare an IUT.
         :type preparation_steps: dict
         """
-        with self.lock:
-            dataset = self.dataset.copy()
-        jsontas = JsonTas(dataset=dataset)
-        steps = {}
-        dataset.add("iut", iut)
-        dataset.add("steps", steps)
-        for step, definition in preparation_steps.items():
-            definition = OrderedDict(**definition)
-            self.logger.info("Executing step %r", step)
-            step_result = jsontas.run(json_data=definition)
-            self.logger.info("%r", step_result)
-            if not step_result:
-                self.logger.error("Failed to execute step %r", step)
-                return False, iut
-            steps[step] = step_result
+        try:
+            with self.lock:
+                dataset = self.dataset.copy()
+                dataset.add("config", self.config)
+            jsontas = JsonTas(dataset=dataset)
+            steps = {}
+            dataset.add("iut", iut)
+            dataset.add("steps", steps)
+            for step, definition in preparation_steps.items():
+                definition = OrderedDict(**definition)
+                self.logger.info("Executing step %r", step)
+                step_result = jsontas.run(json_data=definition)
+                self.logger.info("%r", step_result)
+                if not step_result:
+                    self.logger.error("Failed to execute step %r", step)
+                    return False, iut
+                steps[step] = step_result
+        except Exception as exception:  # pylint:disable=broad-except
+            self.logger.error("Failure when preparing IUT %r", iut)
+            self.logger.error("%r", exception)
+            return False, iut
         return True, iut
 
     def prepare(self, iuts):
@@ -73,27 +83,31 @@ def prepare(self, iuts):
         :return: Prepared IUTs.
         :rtype: list
         """
-        if not self.prepare_ruleset:
-            self.logger.info("No defined preparation rule.")
-            return iuts
-        thread_pool = ThreadPool()
+        try:
+            if not self.prepare_ruleset:
+                self.logger.info("No defined preparation rule.")
+                return iuts
+            thread_pool = ThreadPool()
 
-        stages = self.prepare_ruleset.get("stages", {})
-        steps = stages.get("environment_provider", {}).get("steps", {})
-        results = []
-        for iut in reversed(iuts):
-            self.logger.info("Preparing IUT %r", iut)
-            results.append(
-                thread_pool.apply_async(
-                    self.execute_preparation_steps, args=(iut, deepcopy(steps))
+            stages = self.prepare_ruleset.get("stages", {})
+            steps = stages.get("environment_provider", {}).get("steps", {})
+            results = []
+            for iut in reversed(iuts):
+                self.logger.info("Preparing IUT %r", iut)
+                results.append(
+                    thread_pool.apply_async(
+                        self.execute_preparation_steps, args=(iut, deepcopy(steps))
+                    )
                 )
-            )
-        for result in results:
-            success, iut = result.get()
-            if not success:
-                self.logger.error("Unable to prepare %r.", iut)
-                iuts.remove(iut)
-            else:
-                iut.update(**deepcopy(stages))
-        self.dataset.add("iuts", deepcopy(iuts))
-        return iuts
+            for result in results:
+                success, iut = result.get()
+                if not success:
+                    self.logger.error("Unable to prepare %r.", iut)
+                    iuts.remove(iut)
+                else:
+                    iut.update(**deepcopy(stages))
+            self.dataset.add("iuts", deepcopy(iuts))
+            return iuts
+        finally:
+            # Re-add the config that was popped in __init__.
+            self.dataset.add("config", self.config)