From dd16c3d5c1bc275f5d42c2694ea2927c792fe9a0 Mon Sep 17 00:00:00 2001
From: Rodrigo Agundez <rodrigo.agundez@dyson.com>
Date: Thu, 16 Jan 2025 08:09:56 +0800
Subject: [PATCH] add function to retrieve the storage provider so it can be
 easily tested

---
 backend/open_webui/storage/provider.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/backend/open_webui/storage/provider.py b/backend/open_webui/storage/provider.py
index 3280e6519..7bfb33ece 100644
--- a/backend/open_webui/storage/provider.py
+++ b/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)