瀏覽代碼

add function to retrieve the storage provider so it can be easily tested

Rodrigo Agundez 3 月之前
父節點
當前提交
dd16c3d5c1
共有 1 個文件被更改,包括 11 次插入6 次删除
  1. 11 6
      backend/open_webui/storage/provider.py

+ 11 - 6
backend/open_webui/storage/provider.py

@@ -138,9 +138,14 @@ class S3StorageProvider(StorageProvider):
         LocalStorageProvider.delete_all_files()
         LocalStorageProvider.delete_all_files()
 
 
 
 
-if STORAGE_PROVIDER == "local":
-    Storage = LocalStorageProvider()
-elif STORAGE_PROVIDER == "s3":
-    Storage = S3StorageProvider()
-else:
-    raise RuntimeError(f"Unsupported storage provider: {STORAGE_PROVIDER}")
+def get_storage_provider(storage_provider: str):
+    if storage_provider == "local":
+        Storage = LocalStorageProvider()
+    elif storage_provider == "s3":
+        Storage = S3StorageProvider()
+    else:
+        raise RuntimeError(f"Unsupported storage provider: {storage_provider}")
+    return Storage
+
+
+Storage = get_storage_provider(STORAGE_PROVIDER)