fix other test

capitalone · Feb 29, 2024 · 2dea6ac · 2dea6ac
1 parent 5cc9ebd
commit 2dea6ac
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tests/fixtures.py b/tests/fixtures.py
@@ -318,7 +318,6 @@ def artifact_schema():
     return {
         "artifacts": [
             "self",
-            {"self": "custom_logging_func"},
             {"name": "object_", "data_object_attr": "object_"},
             {"name": "object_b", "data_object_func": "artifact_function"},
         ]

diff --git a/tests/unit/schema/test_schema_logger.py b/tests/unit/schema/test_schema_logger.py
@@ -113,6 +113,8 @@ def test_log_artifacts_with_schema(objects_to_log, rubicon_project, artifact_sch
     def custom_logging_func(self, obj):
         self.custom_logging_func_called = True
 
+    artifact_schema["artifacts"].append({"self": "custom_logging_func"})
+
     with mock.patch.object(
         rubicon_ml.client.experiment.Experiment,
         "custom_logging_func",
@@ -123,8 +125,8 @@ def custom_logging_func(self, obj):
         experiment = rubicon_project.log_with_schema(object_to_log)
 
     otl_artifact = experiment.artifact(name=otl_cls.__name__)
-    ao_artifact = experiment.artifact(name=artifact_schema["artifacts"][2]["name"])
-    obj_b_artifact = experiment.artifact(name=artifact_schema["artifacts"][3]["name"])
+    ao_artifact = experiment.artifact(name=artifact_schema["artifacts"][1]["name"])
+    obj_b_artifact = experiment.artifact(name=artifact_schema["artifacts"][2]["name"])
 
     assert experiment.custom_logging_func_called
     assert isinstance(otl_artifact.get_data(unpickle=True), otl_cls)