chore: update literals

capitalone · Sep 30, 2024 · 416d09f · 416d09f
1 parent 9232332
commit 416d09f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/rubicon_ml/client/artifact.py b/rubicon_ml/client/artifact.py
@@ -71,7 +71,7 @@ def _get_data(self):
     @failsafe
     def get_data(
         self,
-        deserialize: Optional[Literal["h2o_binary", "h2o_mojo", "pickle", "xgboost"]] = None,
+        deserialize: Optional[Literal["h2o", "h2o_binary", "h2o_mojo", "pickle", "xgboost"]] = None,
         unpickle: bool = False,  # TODO: deprecate & move to `deserialize`
     ):
         """Loads the data associated with this artifact and

diff --git a/tests/unit/client/test_artifact_client.py b/tests/unit/client/test_artifact_client.py
@@ -163,6 +163,7 @@ def test_download_location(mock_open, project_client):
 @pytest.mark.parametrize(
     ["use_mojo", "deserialization_method"],
     [
+        (False, "h2o"),
         (False, "h2o_binary"),
         (True, "h2o_mojo"),
     ],