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

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}")

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_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_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
-
-    #     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()
+    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, 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, 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()
+        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