AI/open-webui
1
0
Fork 1
mirror of https://git.mirrors.martin98.com/https://github.com/open-webui/open-webui synced 2025-08-18 07:25:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

gcs tests pass, updated get_file and delete_file to retrieve blob instead of instantiating blob

Browse Source
This commit is contained in:
kahghi 2025-01-21 12:53:49 +08:00
parent 4b56c15a3f
commit 55cc127b03
2 changed files with 47 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
backend/open_webui/storage/provider.py
View File

@ -168,7 +168,7 @@ class GCSStorageProvider(StorageProvider):
try: try:
filename = file_path.removeprefix("gs://").split("/")[1] filename = file_path.removeprefix("gs://").split("/")[1]
local_file_path = f"{UPLOAD_DIR}/{filename}" 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) blob.download_to_filename(local_file_path)
return local_file_path return local_file_path
@ -179,7 +179,7 @@ class GCSStorageProvider(StorageProvider):
"""Handles deletion of the file from GCS storage.""" """Handles deletion of the file from GCS storage."""
try: try:
filename = file_path.removeprefix("gs://").split("/")[1] filename = file_path.removeprefix("gs://").split("/")[1]
blob = self.bucket.blob(filename) blob = self.bucket.get_blob(filename)
blob.delete() blob.delete()
except NotFound as e: except NotFound as e:
raise RuntimeError(f"Error deleting file from GCS: {e}") raise RuntimeError(f"Error deleting file from GCS: {e}")

100
backend/open_webui/test/apps/webui/storage/test_provider.py
View File

@ -7,6 +7,7 @@ from moto import mock_aws
from open_webui.storage import provider from open_webui.storage import provider
from gcp_storage_emulator.server import create_server from gcp_storage_emulator.server import create_server
from google.cloud import storage from google.cloud import storage
from google.cloud.exceptions import NotFound
def mock_upload_dir(monkeypatch, tmp_path): def mock_upload_dir(monkeypatch, tmp_path):
@ -211,7 +212,7 @@ class TestGCSStorageProvider:
contents, gcs_file_path = self.Storage.upload_file( contents, gcs_file_path = self.Storage.upload_file(
io.BytesIO(self.file_content), self.filename 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() assert self.file_content == object.download_as_bytes()
# local checks # local checks
assert (upload_dir / self.filename).exists() assert (upload_dir / self.filename).exists()
@ -222,60 +223,49 @@ class TestGCSStorageProvider:
with pytest.raises(ValueError): with pytest.raises(ValueError):
self.Storage.upload_file(self.file_bytesio_empty, self.filename) self.Storage.upload_file(self.file_bytesio_empty, self.filename)
# def test_get_file(self, monkeypatch, tmp_path): def test_get_file(self, monkeypatch, tmp_path, setup):
# upload_dir = mock_upload_dir(monkeypatch, tmp_path) upload_dir = mock_upload_dir(monkeypatch, tmp_path)
# self.s3_client.create_bucket(Bucket=self.Storage.bucket_name) self.Storage.gcs_client, self.Storage.bucket = setup
# contents, s3_file_path = self.Storage.upload_file( contents, gcs_file_path = self.Storage.upload_file(
# io.BytesIO(self.file_content), self.filename io.BytesIO(self.file_content), self.filename
# ) )
# file_path = self.Storage.get_file(s3_file_path) file_path = self.Storage.get_file(gcs_file_path)
# assert file_path == str(upload_dir / self.filename) assert file_path == str(upload_dir / self.filename)
# assert (upload_dir / self.filename).exists() assert (upload_dir / self.filename).exists()
# def test_delete_file(self, monkeypatch, tmp_path): def test_delete_file(self, monkeypatch, tmp_path, setup):
# upload_dir = mock_upload_dir(monkeypatch, tmp_path) upload_dir = mock_upload_dir(monkeypatch, tmp_path)
# self.s3_client.create_bucket(Bucket=self.Storage.bucket_name) self.Storage.gcs_client, self.Storage.bucket = setup
# contents, s3_file_path = self.Storage.upload_file( contents, gcs_file_path = self.Storage.upload_file(
# io.BytesIO(self.file_content), self.filename io.BytesIO(self.file_content), self.filename
# ) )
# assert (upload_dir / self.filename).exists() # ensure that local directory has the uploaded file as well
# self.Storage.delete_file(s3_file_path) assert (upload_dir / self.filename).exists()
# assert not (upload_dir / self.filename).exists() assert self.Storage.bucket.get_blob(self.filename).name == self.filename
# with pytest.raises(ClientError) as exc: self.Storage.delete_file(gcs_file_path)
# self.s3_client.Object(self.Storage.bucket_name, self.filename).load() # check that deleting file from gcs will delete the local file as well
# error = exc.value.response["Error"] assert not (upload_dir / self.filename).exists()
# assert error["Code"] == "404" assert self.Storage.bucket.get_blob(self.filename) == None
# assert error["Message"] == "Not Found"
# def test_delete_all_files(self, monkeypatch, tmp_path): def test_delete_all_files(self, monkeypatch, tmp_path, setup):
# upload_dir = mock_upload_dir(monkeypatch, tmp_path) upload_dir = mock_upload_dir(monkeypatch, tmp_path)
# # create 2 files self.Storage.gcs_client, self.Storage.bucket = setup
# self.s3_client.create_bucket(Bucket=self.Storage.bucket_name) # create 2 files
# self.Storage.upload_file(io.BytesIO(self.file_content), self.filename) self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
# object = self.s3_client.Object(self.Storage.bucket_name, self.filename) object = self.Storage.bucket.get_blob(self.filename)
# assert self.file_content == object.get()["Body"].read() assert (upload_dir / self.filename).exists()
# assert (upload_dir / self.filename).exists() assert (upload_dir / self.filename).read_bytes() == self.file_content
# assert (upload_dir / self.filename).read_bytes() == self.file_content assert self.Storage.bucket.get_blob(self.filename).name == self.filename
# self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra) assert self.file_content == object.download_as_bytes()
# object = self.s3_client.Object(self.Storage.bucket_name, self.filename_extra) self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra)
# assert self.file_content == object.get()["Body"].read() object = self.Storage.bucket.get_blob(self.filename_extra)
# assert (upload_dir / self.filename).exists() assert (upload_dir / self.filename_extra).exists()
# assert (upload_dir / self.filename).read_bytes() == self.file_content 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() self.Storage.delete_all_files()
# assert not (upload_dir / self.filename).exists() assert not (upload_dir / self.filename).exists()
# with pytest.raises(ClientError) as exc: assert not (upload_dir / self.filename_extra).exists()
# self.s3_client.Object(self.Storage.bucket_name, self.filename).load() assert self.Storage.bucket.get_blob(self.filename) == None
# error = exc.value.response["Error"] assert self.Storage.bucket.get_blob(self.filename_extra) == None
# 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()