mirror of
https://github.com/open-webui/open-webui
synced 2024-11-07 09:09:53 +00:00
239 lines
8.7 KiB
Python
239 lines
8.7 KiB
Python
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 apps.webui.models.chats import ChatForm
|
|
from apps.webui.models.chats import 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 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
|