fix tests

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>
mlflow · Jun 27, 2024 · 32f1594 · 32f1594
1 parent e332c67
commit 32f1594
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/tests/store/_unity_catalog/model_registry/test_unity_catalog_rest_store.py b/tests/store/_unity_catalog/model_registry/test_unity_catalog_rest_store.py
@@ -314,6 +314,7 @@ def test_create_model_version_with_langchain_dependencies(store, langchain_local
             secret_access_key=secret_access_key,
             session_token=session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_artifact_repo.log_artifacts.assert_called_once_with(local_dir=ANY, artifact_path="")
         _assert_create_model_version_endpoints_called(
@@ -377,6 +378,7 @@ def test_create_model_version_with_resources(store, langchain_local_model_dir_wi
             secret_access_key=secret_access_key,
             session_token=session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_artifact_repo.log_artifacts.assert_called_once_with(local_dir=ANY, artifact_path="")
         _assert_create_model_version_endpoints_called(
@@ -434,6 +436,7 @@ def test_create_model_version_with_langchain_no_dependencies(
             secret_access_key=secret_access_key,
             session_token=session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_artifact_repo.log_artifacts.assert_called_once_with(local_dir=ANY, artifact_path="")
         _assert_create_model_version_endpoints_called(
@@ -1019,6 +1022,7 @@ def test_create_model_version_aws(store, local_model_dir):
             secret_access_key=secret_access_key,
             session_token=session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_artifact_repo.log_artifacts.assert_called_once_with(local_dir=ANY, artifact_path="")
         _assert_create_model_version_endpoints_called(

diff --git a/tests/store/artifact/test_unity_catalog_models_artifact_repo.py b/tests/store/artifact/test_unity_catalog_models_artifact_repo.py
@@ -142,6 +142,7 @@ def test_uc_models_artifact_repo_download_artifacts_uses_temporary_creds_aws(mon
             secret_access_key=fake_secret_access_key,
             session_token=fake_session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_s3_repo.download_artifacts.assert_called_once_with("artifact_path", "dst_path")
         request_mock.assert_called_with(