UTs Fixes

PlaytikaOSS · Jul 28, 2024 · 3911584 · 3911584
1 parent 1a64899
commit 3911584
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/tests/test_cmab.py b/tests/test_cmab.py
@@ -328,7 +328,14 @@ def test_cmab_get_state(mu, sigma, n_features):
     cmab = CmabBernoulli(actions=actions)
     expected_state = json.loads(
         json.dumps(
-            {"actions": actions, "strategy": {}, "predict_with_proba": False, "predict_actions_randomly": False},
+            {
+                "actions": actions,
+                "strategy": {},
+                "predict_with_proba": False,
+                "predict_actions_randomly": False,
+                "epsilon": None,
+                "default_action": None,
+            },
             default=dict,
         )
     )
@@ -547,6 +554,8 @@ def test_cmab_bai_get_state(mu, sigma, n_features, exploit_p: Float01):
                 "strategy": {"exploit_p": exploit_p},
                 "predict_with_proba": False,
                 "predict_actions_randomly": False,
+                "epsilon": None,
+                "default_action": None,
             },
             default=dict,
         )
@@ -790,6 +799,8 @@ def test_cmab_cc_get_state(
                 "strategy": {"subsidy_factor": subsidy_factor},
                 "predict_with_proba": True,
                 "predict_actions_randomly": False,
+                "epsilon": None,
+                "default_action": None,
             },
             default=dict,
         )

diff --git a/tests/test_smab.py b/tests/test_smab.py
@@ -206,7 +206,12 @@ def test_smab_get_state(a, b, c, d):
     actions = {"action1": Beta(n_successes=a, n_failures=b), "action2": Beta(n_successes=c, n_failures=d)}
     smab = SmabBernoulli(actions=actions)
 
-    expected_state = {"actions": actions, "strategy": {}}
+    expected_state = {
+        "actions": actions,
+        "strategy": {},
+        "epsilon": None,
+        "default_action": None,
+    }
 
     class_name, smab_state = smab.get_state()
     assert class_name == "SmabBernoulli"