add tests

aimclub · maypink · Sep 12, 2023 · Mar 17, 2023 · Mar 17, 2023 · Mar 21, 2023
commit 46241dedc8ae5bcf5cde9a397a22f8c6d4e2fa5c
diff --git a/golem/core/optimisers/adaptive/mab_agents/mab_agent.py b/golem/core/optimisers/adaptive/mab_agents/mab_agent.py
@@ -1,6 +1,7 @@
 import os.path
 import pickle
 import random
+from pathlib import Path
 from typing import Union, Sequence, Optional
 
 from mabwiser.mab import MAB, LearningPolicy
@@ -58,18 +59,21 @@ def partial_fit(self, experience: ExperienceBuffer):
 
     def save(self):
         """ Saves bandit to specified file. """
-        path_to_save = os.path.join(self._path_to_save, 'MAB')
-        os.makedirs(path_to_save, exist_ok=True)
         # to get file name
-        mabs_num = [int(e.split('_')[0]) for e in os.listdir(path_to_save) if e.split('_')[0].isdigit()]
-        if not mabs_num:
-            max_saved_mab = 0
+        if not self._path_to_save.endswith('.pkl'):
+            path_to_save = os.path.join(self._path_to_save, 'MAB')
+            os.makedirs(path_to_save, exist_ok=True)
+            mabs_num = [int(e.split('_')[0]) for e in os.listdir(path_to_save) if e.split('_')[0].isdigit()]
+            if not mabs_num:
+                max_saved_mab = 0
+            else:
+                max_saved_mab = max(mabs_num) + 1
+            path_to_file = os.path.join(path_to_save, f'{max_saved_mab}_mab.pkl')
         else:
-            max_saved_mab = max(mabs_num) + 1
-        path_to_file = os.path.join(path_to_save, f'{max_saved_mab}_mab.pkl')
+            path_to_file = self._path_to_save
         with open(path_to_file, 'wb') as f:
             pickle.dump(self, f)
-        self._log.info(f"MAB was save to {self._path_to_save}")
+        self._log.info(f"MAB was saved to {self._path_to_save}")
 
     @staticmethod
     def load(path: str):

diff --git a/test/unit/adaptive/test_mab_agent.py b/test/unit/adaptive/test_mab_agent.py
@@ -0,0 +1,24 @@
+import os.path
+from pathlib import Path
+
+from golem.core.optimisers.adaptive.mab_agents.mab_agent import MultiArmedBanditAgent
+
+
+def test_save_mab():
+    """ Tests if MAB is saved. """
+    file_name = 'test_mab.pkl'
+    path_to_save = os.path.join(Path(__file__).parent, file_name)
+    mab = MultiArmedBanditAgent(actions=[0, 1, 2],
+                                n_jobs=1,
+                                path_to_save=path_to_save)
+    mab.save()
+    assert file_name in os.listdir(Path(__file__).parent)
+
+
+def test_load_mab():
+    """ Tests if MAB is loaded. """
+    file_name = 'test_mab.pkl'
+    path_to_load = os.path.join(Path(__file__).parent, file_name)
+    mab = MultiArmedBanditAgent.load(path=path_to_load)
+    assert isinstance(mab, MultiArmedBanditAgent)
+    os.remove(path_to_load)