diff --git a/backend/open_webui/storage/provider.py b/backend/open_webui/storage/provider.py index 1eb893564..1aeec0e1e 100644 --- a/backend/open_webui/storage/provider.py +++ b/backend/open_webui/storage/provider.py @@ -168,7 +168,7 @@ class GCSStorageProvider(StorageProvider): try: filename = file_path.removeprefix("gs://").split("/")[1] local_file_path = f"{UPLOAD_DIR}/{filename}" - blob = self.bucket.blob(filename) + blob = self.bucket.get_blob(filename) blob.download_to_filename(local_file_path) return local_file_path @@ -179,7 +179,7 @@ class GCSStorageProvider(StorageProvider): """Handles deletion of the file from GCS storage.""" try: filename = file_path.removeprefix("gs://").split("/")[1] - blob = self.bucket.blob(filename) + blob = self.bucket.get_blob(filename) blob.delete() except NotFound as e: raise RuntimeError(f"Error deleting file from GCS: {e}") diff --git a/backend/open_webui/test/apps/webui/storage/test_provider.py b/backend/open_webui/test/apps/webui/storage/test_provider.py index c2f8e9e9a..37b608d37 100644 --- a/backend/open_webui/test/apps/webui/storage/test_provider.py +++ b/backend/open_webui/test/apps/webui/storage/test_provider.py @@ -7,6 +7,7 @@ from moto import mock_aws from open_webui.storage import provider from gcp_storage_emulator.server import create_server from google.cloud import storage +from google.cloud.exceptions import NotFound def mock_upload_dir(monkeypatch, tmp_path): @@ -211,7 +212,7 @@ class TestGCSStorageProvider: contents, gcs_file_path = self.Storage.upload_file( io.BytesIO(self.file_content), self.filename ) - object = self.Storage.bucket.blob(self.filename) + object = self.Storage.bucket.get_blob(self.filename) assert self.file_content == object.download_as_bytes() # local checks assert (upload_dir / self.filename).exists() @@ -222,60 +223,49 @@ class TestGCSStorageProvider: with pytest.raises(ValueError): self.Storage.upload_file(self.file_bytesio_empty, self.filename) - # def test_get_file(self, monkeypatch, tmp_path): - # upload_dir = mock_upload_dir(monkeypatch, tmp_path) - # self.s3_client.create_bucket(Bucket=self.Storage.bucket_name) - # contents, s3_file_path = self.Storage.upload_file( - # io.BytesIO(self.file_content), self.filename - # ) - # file_path = self.Storage.get_file(s3_file_path) - # assert file_path == str(upload_dir / self.filename) - # assert (upload_dir / self.filename).exists() + def test_get_file(self, monkeypatch, tmp_path, setup): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + self.Storage.gcs_client, self.Storage.bucket = setup + contents, gcs_file_path = self.Storage.upload_file( + io.BytesIO(self.file_content), self.filename + ) + file_path = self.Storage.get_file(gcs_file_path) + assert file_path == str(upload_dir / self.filename) + assert (upload_dir / self.filename).exists() - # def test_delete_file(self, monkeypatch, tmp_path): - # upload_dir = mock_upload_dir(monkeypatch, tmp_path) - # self.s3_client.create_bucket(Bucket=self.Storage.bucket_name) - # contents, s3_file_path = self.Storage.upload_file( - # io.BytesIO(self.file_content), self.filename - # ) - # assert (upload_dir / self.filename).exists() - # self.Storage.delete_file(s3_file_path) - # assert not (upload_dir / self.filename).exists() - # with pytest.raises(ClientError) as exc: - # self.s3_client.Object(self.Storage.bucket_name, self.filename).load() - # error = exc.value.response["Error"] - # assert error["Code"] == "404" - # assert error["Message"] == "Not Found" + def test_delete_file(self, monkeypatch, tmp_path, setup): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + self.Storage.gcs_client, self.Storage.bucket = setup + contents, gcs_file_path = self.Storage.upload_file( + io.BytesIO(self.file_content), self.filename + ) + # ensure that local directory has the uploaded file as well + assert (upload_dir / self.filename).exists() + assert self.Storage.bucket.get_blob(self.filename).name == self.filename + self.Storage.delete_file(gcs_file_path) + # check that deleting file from gcs will delete the local file as well + assert not (upload_dir / self.filename).exists() + assert self.Storage.bucket.get_blob(self.filename) == None - # def test_delete_all_files(self, monkeypatch, tmp_path): - # upload_dir = mock_upload_dir(monkeypatch, tmp_path) - # # create 2 files - # self.s3_client.create_bucket(Bucket=self.Storage.bucket_name) - # self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) - # object = self.s3_client.Object(self.Storage.bucket_name, self.filename) - # assert self.file_content == object.get()["Body"].read() - # assert (upload_dir / self.filename).exists() - # assert (upload_dir / self.filename).read_bytes() == self.file_content - # self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra) - # object = self.s3_client.Object(self.Storage.bucket_name, self.filename_extra) - # assert self.file_content == object.get()["Body"].read() - # assert (upload_dir / self.filename).exists() - # assert (upload_dir / self.filename).read_bytes() == self.file_content + def test_delete_all_files(self, monkeypatch, tmp_path, setup): + upload_dir = mock_upload_dir(monkeypatch, tmp_path) + self.Storage.gcs_client, self.Storage.bucket = setup + # create 2 files + self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) + object = self.Storage.bucket.get_blob(self.filename) + assert (upload_dir / self.filename).exists() + assert (upload_dir / self.filename).read_bytes() == self.file_content + assert self.Storage.bucket.get_blob(self.filename).name == self.filename + assert self.file_content == object.download_as_bytes() + self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra) + object = self.Storage.bucket.get_blob(self.filename_extra) + assert (upload_dir / self.filename_extra).exists() + assert (upload_dir / self.filename_extra).read_bytes() == self.file_content + assert self.Storage.bucket.get_blob(self.filename_extra).name == self.filename_extra + assert self.file_content == object.download_as_bytes() - # self.Storage.delete_all_files() - # assert not (upload_dir / self.filename).exists() - # with pytest.raises(ClientError) as exc: - # self.s3_client.Object(self.Storage.bucket_name, self.filename).load() - # error = exc.value.response["Error"] - # assert error["Code"] == "404" - # assert error["Message"] == "Not Found" - # assert not (upload_dir / self.filename_extra).exists() - # with pytest.raises(ClientError) as exc: - # self.s3_client.Object(self.Storage.bucket_name, self.filename_extra).load() - # error = exc.value.response["Error"] - # assert error["Code"] == "404" - # assert error["Message"] == "Not Found" - - # self.Storage.delete_all_files() - # assert not (upload_dir / self.filename).exists() - # assert not (upload_dir / self.filename_extra).exists() + self.Storage.delete_all_files() + assert not (upload_dir / self.filename).exists() + assert not (upload_dir / self.filename_extra).exists() + assert self.Storage.bucket.get_blob(self.filename) == None + assert self.Storage.bucket.get_blob(self.filename_extra) == None