fix: allow destination directory to exist before we copy.

PiperOrigin-RevId: 609578936
googleapis · Feb 23, 2024 · 0b55762 · 0b55762
1 parent 2206d16
commit 0b55762
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/google/cloud/aiplatform/utils/prediction_utils.py b/google/cloud/aiplatform/utils/prediction_utils.py
@@ -150,4 +150,4 @@ def download_model_artifacts(artifact_uri: str) -> None:
                 blob.download_to_filename(name_without_prefix)
     else:
         # Copy files to the current working directory.
-        shutil.copytree(artifact_uri, ".")
+        shutil.copytree(artifact_uri, ".", dirs_exist_ok=True)
diff --git a/tests/unit/aiplatform/test_utils.py b/tests/unit/aiplatform/test_utils.py
@@ -965,7 +965,7 @@ def test_download_model_artifacts_not_gcs_uri(
         prediction_utils.download_model_artifacts(model_dir_name)
 
         assert not mock_storage_client.called
-        copy_tree_mock.assert_called_once_with(model_dir_name, ".")
+        copy_tree_mock.assert_called_once_with(model_dir_name, ".", dirs_exist_ok=True)
 
 
 @pytest.fixture(scope="function")