highperfocused/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2025-04-11 05:19:09 +02: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:
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}")

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 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