Преглед изворни кода

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