From f956add20c1e58c6e8186d62b47aff091abbbdfb Mon Sep 17 00:00:00 2001
From: Chell <im.chell@outlook.com>
Date: Thu, 2 May 2024 17:51:17 +0200
Subject: [PATCH] Refactor Orchestrator class in core.py to include before and
 after experiment plugins

---
 midori/core.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/midori/core.py b/midori/core.py
index d72fad2..7433041 100644
--- a/midori/core.py
+++ b/midori/core.py
@@ -16,20 +16,26 @@ def __init__(
         hostname: str,
         username: str,
         password: str,
+        before_experiment_plugins: List[Type[PluginHelper]],
         before_trial_cooling_time: int,
         setup_plugins: List[Type[PluginHelper]],
         trial_timespan: int,
         end_trial_plugins: List[Type[PluginHelper]],
         after_trial_cooling_time: int,
+        end_experiment_plugins: List[Type[PluginHelper]],
         variables: Dict[str, List[str]],
         subject_path: str,
     ) -> None:
+        self.__before_experiment_plugins: List[Type[PluginHelper]] = (
+            before_experiment_plugins
+        )
         self.__subject_path: str = subject_path
         self.__before_trial_cooling_time: int = before_trial_cooling_time
         self.__trial_timespan: int = trial_timespan
         self.__setup_plugins: List[Type[PluginHelper]] = setup_plugins
         self.__after_trial_cooling_time: int = after_trial_cooling_time
         self.__end_trial_plugins: List[Type[PluginHelper]] = end_trial_plugins
+        self.__end_experiment_plugins: List[Type[PluginHelper]] = end_experiment_plugins
         self.__ssh: paramiko.SSHClient = paramiko.SSHClient()
         self.__ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
@@ -43,6 +49,18 @@ def __init__(
 
     def run(self) -> None:
         print("Starting the experiment...")
+
+        # Before Experiment plugins
+        output = None
+        for Plugin in self.__before_experiment_plugins:
+            plugin = Plugin(
+                ssh=self.__ssh,
+                subject_path=self.__subject_path,
+                treatment="",
+                previous_output=output,
+            )
+            output = plugin.execute()
+
         print(f"Treatments: {self.treatments}")
         for treatment in self.treatments:
             if not is_a_branch_exist(
@@ -85,4 +103,15 @@ def run(self) -> None:
             # After Trial Cooling time
             pause(interval=self.__after_trial_cooling_time)
 
+        # End Experiment plugins
+        output = None
+        for Plugin in self.__end_experiment_plugins:
+            plugin = Plugin(
+                ssh=self.__ssh,
+                subject_path=self.__subject_path,
+                treatment=treatment,
+                previous_output=output,
+            )
+            output = plugin.execute()
+
         close(ssh=self.__ssh)