| 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.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.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 |
|
|
