mirror of
https://github.com/open-webui/open-webui.git
synced 2025-12-16 03:47:49 +01:00
refac: mv backend files to /open_webui dir
This commit is contained in:
0
backend/open_webui/test/__init__.py
Normal file
0
backend/open_webui/test/__init__.py
Normal file
200
backend/open_webui/test/apps/webui/routers/test_auths.py
Normal file
200
backend/open_webui/test/apps/webui/routers/test_auths.py
Normal file
@@ -0,0 +1,200 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestAuths(AbstractPostgresTest):
|
||||
BASE_PATH = "/api/v1/auths"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from open_webui.apps.webui.models.auths import Auths
|
||||
from open_webui.apps.webui.models.users import Users
|
||||
|
||||
cls.users = Users
|
||||
cls.auths = Auths
|
||||
|
||||
def test_get_session_user(self):
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"id": "1",
|
||||
"name": "John Doe",
|
||||
"email": "john.doe@openwebui.com",
|
||||
"role": "user",
|
||||
"profile_image_url": "/user.png",
|
||||
}
|
||||
|
||||
def test_update_profile(self):
|
||||
from open_webui.utils.utils import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password=get_password_hash("old_password"),
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="user",
|
||||
)
|
||||
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/update/profile"),
|
||||
json={"name": "John Doe 2", "profile_image_url": "/user2.png"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
db_user = self.users.get_user_by_id(user.id)
|
||||
assert db_user.name == "John Doe 2"
|
||||
assert db_user.profile_image_url == "/user2.png"
|
||||
|
||||
def test_update_password(self):
|
||||
from open_webui.utils.utils import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password=get_password_hash("old_password"),
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="user",
|
||||
)
|
||||
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/update/password"),
|
||||
json={"password": "old_password", "new_password": "new_password"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
old_auth = self.auths.authenticate_user(
|
||||
"john.doe@openwebui.com", "old_password"
|
||||
)
|
||||
assert old_auth is None
|
||||
new_auth = self.auths.authenticate_user(
|
||||
"john.doe@openwebui.com", "new_password"
|
||||
)
|
||||
assert new_auth is not None
|
||||
|
||||
def test_signin(self):
|
||||
from open_webui.utils.utils import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password=get_password_hash("password"),
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="user",
|
||||
)
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/signin"),
|
||||
json={"email": "john.doe@openwebui.com", "password": "password"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == user.id
|
||||
assert data["name"] == "John Doe"
|
||||
assert data["email"] == "john.doe@openwebui.com"
|
||||
assert data["role"] == "user"
|
||||
assert data["profile_image_url"] == "/user.png"
|
||||
assert data["token"] is not None and len(data["token"]) > 0
|
||||
assert data["token_type"] == "Bearer"
|
||||
|
||||
def test_signup(self):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/signup"),
|
||||
json={
|
||||
"name": "John Doe",
|
||||
"email": "john.doe@openwebui.com",
|
||||
"password": "password",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] is not None and len(data["id"]) > 0
|
||||
assert data["name"] == "John Doe"
|
||||
assert data["email"] == "john.doe@openwebui.com"
|
||||
assert data["role"] in ["admin", "user", "pending"]
|
||||
assert data["profile_image_url"] == "/user.png"
|
||||
assert data["token"] is not None and len(data["token"]) > 0
|
||||
assert data["token_type"] == "Bearer"
|
||||
|
||||
def test_add_user(self):
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/add"),
|
||||
json={
|
||||
"name": "John Doe 2",
|
||||
"email": "john.doe2@openwebui.com",
|
||||
"password": "password2",
|
||||
"role": "admin",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] is not None and len(data["id"]) > 0
|
||||
assert data["name"] == "John Doe 2"
|
||||
assert data["email"] == "john.doe2@openwebui.com"
|
||||
assert data["role"] == "admin"
|
||||
assert data["profile_image_url"] == "/user.png"
|
||||
assert data["token"] is not None and len(data["token"]) > 0
|
||||
assert data["token_type"] == "Bearer"
|
||||
|
||||
def test_get_admin_details(self):
|
||||
self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password="password",
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="admin",
|
||||
)
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.get(self.create_url("/admin/details"))
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"name": "John Doe",
|
||||
"email": "john.doe@openwebui.com",
|
||||
}
|
||||
|
||||
def test_create_api_key_(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password="password",
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="admin",
|
||||
)
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(self.create_url("/api_key"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["api_key"] is not None
|
||||
assert len(data["api_key"]) > 0
|
||||
|
||||
def test_delete_api_key(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password="password",
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="admin",
|
||||
)
|
||||
self.users.update_user_api_key_by_id(user.id, "abc")
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.delete(self.create_url("/api_key"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == True
|
||||
db_user = self.users.get_user_by_id(user.id)
|
||||
assert db_user.api_key is None
|
||||
|
||||
def test_get_api_key(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password="password",
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="admin",
|
||||
)
|
||||
self.users.update_user_api_key_by_id(user.id, "abc")
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.get(self.create_url("/api_key"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"api_key": "abc"}
|
||||
236
backend/open_webui/test/apps/webui/routers/test_chats.py
Normal file
236
backend/open_webui/test/apps/webui/routers/test_chats.py
Normal file
@@ -0,0 +1,236 @@
|
||||
import uuid
|
||||
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestChats(AbstractPostgresTest):
|
||||
BASE_PATH = "/api/v1/chats"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
|
||||
def setup_method(self):
|
||||
super().setup_method()
|
||||
from open_webui.apps.webui.models.chats import ChatForm, Chats
|
||||
|
||||
self.chats = Chats
|
||||
self.chats.insert_new_chat(
|
||||
"2",
|
||||
ChatForm(
|
||||
**{
|
||||
"chat": {
|
||||
"name": "chat1",
|
||||
"description": "chat1 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
def test_get_session_user_chat_list(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
first_chat = response.json()[0]
|
||||
assert first_chat["id"] is not None
|
||||
assert first_chat["title"] == "New Chat"
|
||||
assert first_chat["created_at"] is not None
|
||||
assert first_chat["updated_at"] is not None
|
||||
|
||||
def test_delete_all_user_chats(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(self.chats.get_chats()) == 0
|
||||
|
||||
def test_get_user_chat_list_by_user_id(self):
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url("/list/user/2"))
|
||||
assert response.status_code == 200
|
||||
first_chat = response.json()[0]
|
||||
assert first_chat["id"] is not None
|
||||
assert first_chat["title"] == "New Chat"
|
||||
assert first_chat["created_at"] is not None
|
||||
assert first_chat["updated_at"] is not None
|
||||
|
||||
def test_create_new_chat(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/new"),
|
||||
json={
|
||||
"chat": {
|
||||
"name": "chat2",
|
||||
"description": "chat2 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["archived"] is False
|
||||
assert data["chat"] == {
|
||||
"name": "chat2",
|
||||
"description": "chat2 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
}
|
||||
assert data["user_id"] == "2"
|
||||
assert data["id"] is not None
|
||||
assert data["share_id"] is None
|
||||
assert data["title"] == "New Chat"
|
||||
assert data["updated_at"] is not None
|
||||
assert data["created_at"] is not None
|
||||
assert len(self.chats.get_chats()) == 2
|
||||
|
||||
def test_get_user_chats(self):
|
||||
self.test_get_session_user_chat_list()
|
||||
|
||||
def test_get_user_archived_chats(self):
|
||||
self.chats.archive_all_chats_by_user_id("2")
|
||||
from open_webui.apps.webui.internal.db import Session
|
||||
|
||||
Session.commit()
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/all/archived"))
|
||||
assert response.status_code == 200
|
||||
first_chat = response.json()[0]
|
||||
assert first_chat["id"] is not None
|
||||
assert first_chat["title"] == "New Chat"
|
||||
assert first_chat["created_at"] is not None
|
||||
assert first_chat["updated_at"] is not None
|
||||
|
||||
def test_get_all_user_chats_in_db(self):
|
||||
with mock_webui_user(id="4"):
|
||||
response = self.fast_api_client.get(self.create_url("/all/db"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
|
||||
def test_get_archived_session_user_chat_list(self):
|
||||
self.test_get_user_archived_chats()
|
||||
|
||||
def test_archive_all_chats(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(self.create_url("/archive/all"))
|
||||
assert response.status_code == 200
|
||||
assert len(self.chats.get_archived_chats_by_user_id("2")) == 1
|
||||
|
||||
def test_get_shared_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
self.chats.update_chat_share_id_by_id(chat_id, chat_id)
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url(f"/share/{chat_id}"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == chat_id
|
||||
assert data["chat"] == {
|
||||
"name": "chat1",
|
||||
"description": "chat1 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
}
|
||||
assert data["id"] == chat_id
|
||||
assert data["share_id"] == chat_id
|
||||
assert data["title"] == "New Chat"
|
||||
|
||||
def test_get_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url(f"/{chat_id}"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == chat_id
|
||||
assert data["chat"] == {
|
||||
"name": "chat1",
|
||||
"description": "chat1 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
}
|
||||
assert data["share_id"] is None
|
||||
assert data["title"] == "New Chat"
|
||||
assert data["user_id"] == "2"
|
||||
|
||||
def test_update_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url(f"/{chat_id}"),
|
||||
json={
|
||||
"chat": {
|
||||
"name": "chat2",
|
||||
"description": "chat2 description",
|
||||
"tags": ["tag2", "tag4"],
|
||||
"title": "Just another title",
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == chat_id
|
||||
assert data["chat"] == {
|
||||
"name": "chat2",
|
||||
"title": "Just another title",
|
||||
"description": "chat2 description",
|
||||
"tags": ["tag2", "tag4"],
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
}
|
||||
assert data["share_id"] is None
|
||||
assert data["title"] == "Just another title"
|
||||
assert data["user_id"] == "2"
|
||||
|
||||
def test_delete_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(self.create_url(f"/{chat_id}"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() is True
|
||||
|
||||
def test_clone_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url(f"/{chat_id}/clone"))
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] != chat_id
|
||||
assert data["chat"] == {
|
||||
"branchPointMessageId": "1",
|
||||
"description": "chat1 description",
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
"name": "chat1",
|
||||
"originalChatId": chat_id,
|
||||
"tags": ["tag1", "tag2"],
|
||||
"title": "Clone of New Chat",
|
||||
}
|
||||
assert data["share_id"] is None
|
||||
assert data["title"] == "Clone of New Chat"
|
||||
assert data["user_id"] == "2"
|
||||
|
||||
def test_archive_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url(f"/{chat_id}/archive"))
|
||||
assert response.status_code == 200
|
||||
|
||||
chat = self.chats.get_chat_by_id(chat_id)
|
||||
assert chat.archived is True
|
||||
|
||||
def test_share_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(self.create_url(f"/{chat_id}/share"))
|
||||
assert response.status_code == 200
|
||||
|
||||
chat = self.chats.get_chat_by_id(chat_id)
|
||||
assert chat.share_id is not None
|
||||
|
||||
def test_delete_shared_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
share_id = str(uuid.uuid4())
|
||||
self.chats.update_chat_share_id_by_id(chat_id, share_id)
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(self.create_url(f"/{chat_id}/share"))
|
||||
assert response.status_code
|
||||
|
||||
chat = self.chats.get_chat_by_id(chat_id)
|
||||
assert chat.share_id is None
|
||||
105
backend/open_webui/test/apps/webui/routers/test_documents.py
Normal file
105
backend/open_webui/test/apps/webui/routers/test_documents.py
Normal file
@@ -0,0 +1,105 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestDocuments(AbstractPostgresTest):
|
||||
BASE_PATH = "/api/v1/documents"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from open_webui.apps.webui.models.documents import Documents
|
||||
|
||||
cls.documents = Documents
|
||||
|
||||
def test_documents(self):
|
||||
# Empty database
|
||||
assert len(self.documents.get_docs()) == 0
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
|
||||
# Create a new document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/create"),
|
||||
json={
|
||||
"name": "doc_name",
|
||||
"title": "doc title",
|
||||
"collection_name": "custom collection",
|
||||
"filename": "doc_name.pdf",
|
||||
"content": "",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["name"] == "doc_name"
|
||||
assert len(self.documents.get_docs()) == 1
|
||||
|
||||
# Get the document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/doc?name=doc_name"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["collection_name"] == "custom collection"
|
||||
assert data["name"] == "doc_name"
|
||||
assert data["title"] == "doc title"
|
||||
assert data["filename"] == "doc_name.pdf"
|
||||
assert data["content"] == {}
|
||||
|
||||
# Create another document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/create"),
|
||||
json={
|
||||
"name": "doc_name 2",
|
||||
"title": "doc title 2",
|
||||
"collection_name": "custom collection 2",
|
||||
"filename": "doc_name2.pdf",
|
||||
"content": "",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["name"] == "doc_name 2"
|
||||
assert len(self.documents.get_docs()) == 2
|
||||
|
||||
# Get all documents
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
|
||||
# Update the first document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/doc/update?name=doc_name"),
|
||||
json={"name": "doc_name rework", "title": "updated title"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == "doc_name rework"
|
||||
assert data["title"] == "updated title"
|
||||
|
||||
# Tag the first document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/doc/tags"),
|
||||
json={
|
||||
"name": "doc_name rework",
|
||||
"tags": [{"name": "testing-tag"}, {"name": "another-tag"}],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == "doc_name rework"
|
||||
assert data["content"] == {
|
||||
"tags": [{"name": "testing-tag"}, {"name": "another-tag"}]
|
||||
}
|
||||
assert len(self.documents.get_docs()) == 2
|
||||
|
||||
# Delete the first document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(
|
||||
self.create_url("/doc/delete?name=doc_name rework")
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert len(self.documents.get_docs()) == 1
|
||||
61
backend/open_webui/test/apps/webui/routers/test_models.py
Normal file
61
backend/open_webui/test/apps/webui/routers/test_models.py
Normal file
@@ -0,0 +1,61 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestModels(AbstractPostgresTest):
|
||||
BASE_PATH = "/api/v1/models"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from open_webui.apps.webui.models.models import Model
|
||||
|
||||
cls.models = Model
|
||||
|
||||
def test_models(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/add"),
|
||||
json={
|
||||
"id": "my-model",
|
||||
"base_model_id": "base-model-id",
|
||||
"name": "Hello World",
|
||||
"meta": {
|
||||
"profile_image_url": "/static/favicon.png",
|
||||
"description": "description",
|
||||
"capabilities": None,
|
||||
"model_config": {},
|
||||
},
|
||||
"params": {},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(
|
||||
self.create_url(query_params={"id": "my-model"})
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()[0]
|
||||
assert data["id"] == "my-model"
|
||||
assert data["name"] == "Hello World"
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(
|
||||
self.create_url("/delete?id=my-model")
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
91
backend/open_webui/test/apps/webui/routers/test_prompts.py
Normal file
91
backend/open_webui/test/apps/webui/routers/test_prompts.py
Normal file
@@ -0,0 +1,91 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestPrompts(AbstractPostgresTest):
|
||||
BASE_PATH = "/api/v1/prompts"
|
||||
|
||||
def test_prompts(self):
|
||||
# Get all prompts
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
|
||||
# Create a two new prompts
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/create"),
|
||||
json={
|
||||
"command": "/my-command",
|
||||
"title": "Hello World",
|
||||
"content": "description",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/create"),
|
||||
json={
|
||||
"command": "/my-command2",
|
||||
"title": "Hello World 2",
|
||||
"content": "description 2",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all prompts
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
|
||||
# Get prompt by command
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/command/my-command"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["command"] == "/my-command"
|
||||
assert data["title"] == "Hello World"
|
||||
assert data["content"] == "description"
|
||||
assert data["user_id"] == "2"
|
||||
|
||||
# Update prompt
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/command/my-command2/update"),
|
||||
json={
|
||||
"command": "irrelevant for request",
|
||||
"title": "Hello World Updated",
|
||||
"content": "description Updated",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["command"] == "/my-command2"
|
||||
assert data["title"] == "Hello World Updated"
|
||||
assert data["content"] == "description Updated"
|
||||
assert data["user_id"] == "3"
|
||||
|
||||
# Get prompt by command
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/command/my-command2"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["command"] == "/my-command2"
|
||||
assert data["title"] == "Hello World Updated"
|
||||
assert data["content"] == "description Updated"
|
||||
assert data["user_id"] == "3"
|
||||
|
||||
# Delete prompt
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(
|
||||
self.create_url("/command/my-command/delete")
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all prompts
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
167
backend/open_webui/test/apps/webui/routers/test_users.py
Normal file
167
backend/open_webui/test/apps/webui/routers/test_users.py
Normal file
@@ -0,0 +1,167 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
def _get_user_by_id(data, param):
|
||||
return next((item for item in data if item["id"] == param), None)
|
||||
|
||||
|
||||
def _assert_user(data, id, **kwargs):
|
||||
user = _get_user_by_id(data, id)
|
||||
assert user is not None
|
||||
comparison_data = {
|
||||
"name": f"user {id}",
|
||||
"email": f"user{id}@openwebui.com",
|
||||
"profile_image_url": f"/user{id}.png",
|
||||
"role": "user",
|
||||
**kwargs,
|
||||
}
|
||||
for key, value in comparison_data.items():
|
||||
assert user[key] == value
|
||||
|
||||
|
||||
class TestUsers(AbstractPostgresTest):
|
||||
BASE_PATH = "/api/v1/users"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from open_webui.apps.webui.models.users import Users
|
||||
|
||||
cls.users = Users
|
||||
|
||||
def setup_method(self):
|
||||
super().setup_method()
|
||||
self.users.insert_new_user(
|
||||
id="1",
|
||||
name="user 1",
|
||||
email="user1@openwebui.com",
|
||||
profile_image_url="/user1.png",
|
||||
role="user",
|
||||
)
|
||||
self.users.insert_new_user(
|
||||
id="2",
|
||||
name="user 2",
|
||||
email="user2@openwebui.com",
|
||||
profile_image_url="/user2.png",
|
||||
role="user",
|
||||
)
|
||||
|
||||
def test_users(self):
|
||||
# Get all users
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, "1")
|
||||
_assert_user(data, "2")
|
||||
|
||||
# update role
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/update/role"), json={"id": "2", "role": "admin"}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
_assert_user([response.json()], "2", role="admin")
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, "1")
|
||||
_assert_user(data, "2", role="admin")
|
||||
|
||||
# Get (empty) user settings
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/user/settings"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() is None
|
||||
|
||||
# Update user settings
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/user/settings/update"),
|
||||
json={
|
||||
"ui": {"attr1": "value1", "attr2": "value2"},
|
||||
"model_config": {"attr3": "value3", "attr4": "value4"},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get user settings
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/user/settings"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"ui": {"attr1": "value1", "attr2": "value2"},
|
||||
"model_config": {"attr3": "value3", "attr4": "value4"},
|
||||
}
|
||||
|
||||
# Get (empty) user info
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.get(self.create_url("/user/info"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() is None
|
||||
|
||||
# Update user info
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/user/info/update"),
|
||||
json={"attr1": "value1", "attr2": "value2"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get user info
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.get(self.create_url("/user/info"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"attr1": "value1", "attr2": "value2"}
|
||||
|
||||
# Get user by id
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.get(self.create_url("/2"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"name": "user 2", "profile_image_url": "/user2.png"}
|
||||
|
||||
# Update user by id
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/2/update"),
|
||||
json={
|
||||
"name": "user 2 updated",
|
||||
"email": "user2-updated@openwebui.com",
|
||||
"profile_image_url": "/user2-updated.png",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, "1")
|
||||
_assert_user(
|
||||
data,
|
||||
"2",
|
||||
role="admin",
|
||||
name="user 2 updated",
|
||||
email="user2-updated@openwebui.com",
|
||||
profile_image_url="/user2-updated.png",
|
||||
)
|
||||
|
||||
# Delete user by id
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.delete(self.create_url("/2"))
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
data = response.json()
|
||||
_assert_user(data, "1")
|
||||
161
backend/open_webui/test/util/abstract_integration_test.py
Normal file
161
backend/open_webui/test/util/abstract_integration_test.py
Normal file
@@ -0,0 +1,161 @@
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
import docker
|
||||
import pytest
|
||||
from docker import DockerClient
|
||||
from pytest_docker.plugin import get_docker_ip
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy import text, create_engine
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_fast_api_client():
|
||||
from main import app
|
||||
|
||||
with TestClient(app) as c:
|
||||
return c
|
||||
|
||||
|
||||
class AbstractIntegrationTest:
|
||||
BASE_PATH = None
|
||||
|
||||
def create_url(self, path="", query_params=None):
|
||||
if self.BASE_PATH is None:
|
||||
raise Exception("BASE_PATH is not set")
|
||||
parts = self.BASE_PATH.split("/")
|
||||
parts = [part.strip() for part in parts if part.strip() != ""]
|
||||
path_parts = path.split("/")
|
||||
path_parts = [part.strip() for part in path_parts if part.strip() != ""]
|
||||
query_parts = ""
|
||||
if query_params:
|
||||
query_parts = "&".join(
|
||||
[f"{key}={value}" for key, value in query_params.items()]
|
||||
)
|
||||
query_parts = f"?{query_parts}"
|
||||
return "/".join(parts + path_parts) + query_parts
|
||||
|
||||
@classmethod
|
||||
def setup_class(cls):
|
||||
pass
|
||||
|
||||
def setup_method(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def teardown_class(cls):
|
||||
pass
|
||||
|
||||
def teardown_method(self):
|
||||
pass
|
||||
|
||||
|
||||
class AbstractPostgresTest(AbstractIntegrationTest):
|
||||
DOCKER_CONTAINER_NAME = "postgres-test-container-will-get-deleted"
|
||||
docker_client: DockerClient
|
||||
|
||||
@classmethod
|
||||
def _create_db_url(cls, env_vars_postgres: dict) -> str:
|
||||
host = get_docker_ip()
|
||||
user = env_vars_postgres["POSTGRES_USER"]
|
||||
pw = env_vars_postgres["POSTGRES_PASSWORD"]
|
||||
port = 8081
|
||||
db = env_vars_postgres["POSTGRES_DB"]
|
||||
return f"postgresql://{user}:{pw}@{host}:{port}/{db}"
|
||||
|
||||
@classmethod
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
try:
|
||||
env_vars_postgres = {
|
||||
"POSTGRES_USER": "user",
|
||||
"POSTGRES_PASSWORD": "example",
|
||||
"POSTGRES_DB": "openwebui",
|
||||
}
|
||||
cls.docker_client = docker.from_env()
|
||||
cls.docker_client.containers.run(
|
||||
"postgres:16.2",
|
||||
detach=True,
|
||||
environment=env_vars_postgres,
|
||||
name=cls.DOCKER_CONTAINER_NAME,
|
||||
ports={5432: ("0.0.0.0", 8081)},
|
||||
command="postgres -c log_statement=all",
|
||||
)
|
||||
time.sleep(0.5)
|
||||
|
||||
database_url = cls._create_db_url(env_vars_postgres)
|
||||
os.environ["DATABASE_URL"] = database_url
|
||||
retries = 10
|
||||
db = None
|
||||
while retries > 0:
|
||||
try:
|
||||
from open_webui.config import OPEN_WEBUI_DIR
|
||||
|
||||
db = create_engine(database_url, pool_pre_ping=True)
|
||||
db = db.connect()
|
||||
log.info("postgres is ready!")
|
||||
break
|
||||
except Exception as e:
|
||||
log.warning(e)
|
||||
time.sleep(3)
|
||||
retries -= 1
|
||||
|
||||
if db:
|
||||
# import must be after setting env!
|
||||
cls.fast_api_client = get_fast_api_client()
|
||||
db.close()
|
||||
else:
|
||||
raise Exception("Could not connect to Postgres")
|
||||
except Exception as ex:
|
||||
log.error(ex)
|
||||
cls.teardown_class()
|
||||
pytest.fail(f"Could not setup test environment: {ex}")
|
||||
|
||||
def _check_db_connection(self):
|
||||
from open_webui.apps.webui.internal.db import Session
|
||||
|
||||
retries = 10
|
||||
while retries > 0:
|
||||
try:
|
||||
Session.execute(text("SELECT 1"))
|
||||
Session.commit()
|
||||
break
|
||||
except Exception as e:
|
||||
Session.rollback()
|
||||
log.warning(e)
|
||||
time.sleep(3)
|
||||
retries -= 1
|
||||
|
||||
def setup_method(self):
|
||||
super().setup_method()
|
||||
self._check_db_connection()
|
||||
|
||||
@classmethod
|
||||
def teardown_class(cls) -> None:
|
||||
super().teardown_class()
|
||||
cls.docker_client.containers.get(cls.DOCKER_CONTAINER_NAME).remove(force=True)
|
||||
|
||||
def teardown_method(self):
|
||||
from open_webui.apps.webui.internal.db import Session
|
||||
|
||||
# rollback everything not yet committed
|
||||
Session.commit()
|
||||
|
||||
# truncate all tables
|
||||
tables = [
|
||||
"auth",
|
||||
"chat",
|
||||
"chatidtag",
|
||||
"document",
|
||||
"memory",
|
||||
"model",
|
||||
"prompt",
|
||||
"tag",
|
||||
'"user"',
|
||||
]
|
||||
for table in tables:
|
||||
Session.execute(text(f"TRUNCATE TABLE {table}"))
|
||||
Session.commit()
|
||||
45
backend/open_webui/test/util/mock_user.py
Normal file
45
backend/open_webui/test/util/mock_user.py
Normal file
@@ -0,0 +1,45 @@
|
||||
from contextlib import contextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
|
||||
@contextmanager
|
||||
def mock_webui_user(**kwargs):
|
||||
from open_webui.apps.webui.main import app
|
||||
|
||||
with mock_user(app, **kwargs):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def mock_user(app: FastAPI, **kwargs):
|
||||
from open_webui.utils.utils import (
|
||||
get_current_user,
|
||||
get_verified_user,
|
||||
get_admin_user,
|
||||
get_current_user_by_api_key,
|
||||
)
|
||||
from open_webui.apps.webui.models.users import User
|
||||
|
||||
def create_user():
|
||||
user_parameters = {
|
||||
"id": "1",
|
||||
"name": "John Doe",
|
||||
"email": "john.doe@openwebui.com",
|
||||
"role": "user",
|
||||
"profile_image_url": "/user.png",
|
||||
"last_active_at": 1627351200,
|
||||
"updated_at": 1627351200,
|
||||
"created_at": 162735120,
|
||||
**kwargs,
|
||||
}
|
||||
return User(**user_parameters)
|
||||
|
||||
app.dependency_overrides = {
|
||||
get_current_user: create_user,
|
||||
get_verified_user: create_user,
|
||||
get_admin_user: create_user,
|
||||
get_current_user_by_api_key: create_user,
|
||||
}
|
||||
yield
|
||||
app.dependency_overrides = {}
|
||||
Reference in New Issue
Block a user