95b2547c54
- Deleted session-related functions (UserID, GuestID, UserRoles, UserOrGuestID) from the agent package to streamline the codebase. - Removed RAG-related code and references, simplifying the agent's architecture. - Updated API and load functions to reflect these changes, ensuring consistency across the agent module. - Enhanced test coverage by cleaning up deprecated test cases related to removed functionalities.
1870 lines
57 KiB
Go
1870 lines
57 KiB
Go
package store
|
|
|
|
// import (
|
|
// "fmt"
|
|
// "testing"
|
|
// "time"
|
|
|
|
// "github.com/stretchr/testify/assert"
|
|
// "github.com/yaoapp/gou/connector"
|
|
// "github.com/yaoapp/xun/capsule"
|
|
// "github.com/yaoapp/yao/config"
|
|
// "github.com/yaoapp/yao/test"
|
|
// )
|
|
|
|
// func TestNewXunDefault(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
|
|
// if err != nil {
|
|
// t.Error(err)
|
|
// return
|
|
// }
|
|
|
|
// // Check history table
|
|
// has, err := capsule.Schema().HasTable("__unit_test_conversation_history")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, true, has)
|
|
|
|
// // Check chat table
|
|
// has, err = capsule.Schema().HasTable("__unit_test_conversation_chat")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, true, has)
|
|
|
|
// // Check assistant table
|
|
// has, err = capsule.Schema().HasTable("__unit_test_conversation_assistant")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, true, has)
|
|
|
|
// // Validate table structure by attempting operations
|
|
// // Test history operations
|
|
// messages := []map[string]interface{}{
|
|
// {"role": "user", "content": "test message"},
|
|
// }
|
|
// err = store.SaveHistory("test_user", messages, "test_chat", nil)
|
|
// assert.Nil(t, err)
|
|
|
|
// history, err := store.GetHistory("test_user", "test_chat")
|
|
// assert.Nil(t, err)
|
|
// assert.NotEmpty(t, history)
|
|
|
|
// // Test chat operations
|
|
// err = store.UpdateChatTitle("test_user", "test_chat", "Test Chat")
|
|
// assert.Nil(t, err)
|
|
|
|
// chat, err := store.GetChat("test_user", "test_chat")
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, chat)
|
|
|
|
// // Test assistant operations
|
|
// assistant := map[string]interface{}{
|
|
// "name": "Test Assistant",
|
|
// "type": "assistant",
|
|
// "connector": "test",
|
|
// "description": "Test Description",
|
|
// "tags": []string{"test"},
|
|
// "mentionable": true,
|
|
// "automated": true,
|
|
// }
|
|
|
|
// id, err := store.SaveAssistant(assistant)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, id)
|
|
|
|
// // Clean up test data
|
|
// err = store.DeleteChat("test_user", "test_chat")
|
|
// assert.Nil(t, err)
|
|
|
|
// err = store.DeleteAssistant(id.(string))
|
|
// assert.Nil(t, err)
|
|
// }
|
|
|
|
// func TestNewXunConnector(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
|
|
// conn, err := connector.Select("mysql")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// sch, err := conn.Schema()
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// defer sch.DropTableIfExists("__unit_test_conversation_history")
|
|
// defer sch.DropTableIfExists("__unit_test_conversation_chat")
|
|
// defer sch.DropTableIfExists("__unit_test_conversation_assistant")
|
|
// defer sch.DropTableIfExists("__unit_test_conversation_knowledge")
|
|
// defer sch.DropTableIfExists("__unit_test_conversation_attachment")
|
|
|
|
// sch.DropTableIfExists("__unit_test_conversation_history")
|
|
// sch.DropTableIfExists("__unit_test_conversation_chat")
|
|
// sch.DropTableIfExists("__unit_test_conversation_assistant")
|
|
// sch.DropTableIfExists("__unit_test_conversation_knowledge")
|
|
// sch.DropTableIfExists("__unit_test_conversation_attachment")
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "mysql",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
|
|
// if err != nil {
|
|
// t.Error(err)
|
|
// return
|
|
// }
|
|
|
|
// // Check history table
|
|
// has, err := sch.HasTable("__unit_test_conversation_history")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, true, has)
|
|
|
|
// // Check chat table
|
|
// has, err = sch.HasTable("__unit_test_conversation_chat")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, true, has)
|
|
|
|
// // Check assistant table
|
|
// has, err = sch.HasTable("__unit_test_conversation_assistant")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, true, has)
|
|
|
|
// // Test basic operations
|
|
// messages := []map[string]interface{}{
|
|
// {"role": "user", "content": "test message"},
|
|
// }
|
|
// err = store.SaveHistory("test_user", messages, "test_chat", nil)
|
|
// assert.Nil(t, err)
|
|
|
|
// history, err := store.GetHistory("test_user", "test_chat")
|
|
// assert.Nil(t, err)
|
|
// assert.NotEmpty(t, history)
|
|
|
|
// err = store.DeleteChat("test_user", "test_chat")
|
|
// assert.Nil(t, err)
|
|
// }
|
|
|
|
// func TestXunSaveAndGetHistory(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// TTL: 3600,
|
|
// })
|
|
|
|
// // save the history
|
|
// cid := "123456"
|
|
// err = store.SaveHistory("123456", []map[string]interface{}{
|
|
// {"role": "user", "name": "user1", "content": "hello"},
|
|
// {"role": "assistant", "name": "user1", "content": "Hello there, how"},
|
|
// }, cid, nil)
|
|
// assert.Nil(t, err)
|
|
|
|
// // get the history
|
|
// data, err := store.GetHistory("123456", cid)
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, 2, len(data))
|
|
// }
|
|
|
|
// func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// TTL: 3600,
|
|
// })
|
|
|
|
// // save the history with specific cid
|
|
// sid := "123456"
|
|
// cid := "789012"
|
|
// assistantID := "test-assistant-1"
|
|
// messages := []map[string]interface{}{
|
|
// {"role": "user", "name": "user1", "content": "hello"},
|
|
// {"role": "assistant", "name": "assistant1", "content": "Hi! How can I help you?"},
|
|
// }
|
|
// context := map[string]interface{}{
|
|
// "assistant_id": assistantID,
|
|
// }
|
|
// err = store.SaveHistory(sid, messages, cid, context)
|
|
// assert.Nil(t, err)
|
|
|
|
// // get the history for specific cid
|
|
// data, err := store.GetHistory(sid, cid)
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, 2, len(data))
|
|
|
|
// // Verify assistant_id is saved in chat
|
|
// chat, err := store.GetChat(sid, cid)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, assistantID, chat.Chat["assistant_id"])
|
|
|
|
// // save another message with different cid and assistant
|
|
// anotherCID := "345678"
|
|
// anotherAssistantID := "test-assistant-2"
|
|
// moreMessages := []map[string]interface{}{
|
|
// {"role": "user", "name": "user1", "content": "another message"},
|
|
// {"role": "assistant", "name": "assistant2", "content": "Hello!"},
|
|
// }
|
|
// anotherContext := map[string]interface{}{
|
|
// "assistant_id": anotherAssistantID,
|
|
// }
|
|
// err = store.SaveHistory(sid, moreMessages, anotherCID, anotherContext)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify second chat's assistant_id
|
|
// chat2, err := store.GetChat(sid, anotherCID)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, anotherAssistantID, chat2.Chat["assistant_id"])
|
|
|
|
// // get history for the first cid - should still be 2 messages
|
|
// data, err = store.GetHistory(sid, cid)
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, 2, len(data))
|
|
|
|
// // get history for the second cid - should be 2 messages
|
|
// data, err = store.GetHistory(sid, anotherCID)
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, 2, len(data))
|
|
|
|
// // get all history for the sid without specifying cid
|
|
// allData, err := store.GetHistory(sid, cid)
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, 2, len(allData))
|
|
// }
|
|
|
|
// func TestXunGetChats(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
|
|
// // Drop tables before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Create test assistants first
|
|
// assistant1 := map[string]interface{}{
|
|
// "assistant_id": "test-assistant-1",
|
|
// "name": "Test Assistant 1",
|
|
// "avatar": "avatar1.png",
|
|
// "type": "assistant",
|
|
// "connector": "test",
|
|
// }
|
|
// assistant2 := map[string]interface{}{
|
|
// "assistant_id": "test-assistant-2",
|
|
// "name": "Test Assistant 2",
|
|
// "avatar": "avatar2.png",
|
|
// "type": "assistant",
|
|
// "connector": "test",
|
|
// }
|
|
// _, err = store.SaveAssistant(assistant1)
|
|
// assert.Nil(t, err)
|
|
// _, err = store.SaveAssistant(assistant2)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Save some test chats
|
|
// sid := "test_user"
|
|
// messages := []map[string]interface{}{
|
|
// {"role": "user", "content": "test message"},
|
|
// }
|
|
|
|
// // Create chats with different dates and assistants
|
|
// for i := 0; i < 5; i++ {
|
|
// chatID := fmt.Sprintf("chat_%d", i)
|
|
// title := fmt.Sprintf("Test Chat %d", i)
|
|
// var context map[string]interface{}
|
|
|
|
// // Alternate between having assistant and no assistant
|
|
// if i%2 == 0 {
|
|
// context = map[string]interface{}{
|
|
// "assistant_id": "test-assistant-1",
|
|
// }
|
|
// } else if i%3 == 0 {
|
|
// context = map[string]interface{}{
|
|
// "assistant_id": "test-assistant-2",
|
|
// }
|
|
// }
|
|
|
|
// // Save history first to create the chat
|
|
// err = store.SaveHistory(sid, messages, chatID, context)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Update the chat title
|
|
// err = store.UpdateChatTitle(sid, chatID, title)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify chat was created with correct assistant info
|
|
// chat, err := store.GetChat(sid, chatID)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, chat)
|
|
// assert.Equal(t, chatID, chat.Chat["chat_id"])
|
|
// assert.Equal(t, title, chat.Chat["title"])
|
|
|
|
// if i%2 == 0 {
|
|
// assert.Equal(t, "test-assistant-1", chat.Chat["assistant_id"])
|
|
// assert.Equal(t, "Test Assistant 1", chat.Chat["assistant_name"])
|
|
// assert.Equal(t, "avatar1.png", chat.Chat["assistant_avatar"])
|
|
// } else if i%3 == 0 {
|
|
// assert.Equal(t, "test-assistant-2", chat.Chat["assistant_id"])
|
|
// assert.Equal(t, "Test Assistant 2", chat.Chat["assistant_name"])
|
|
// assert.Equal(t, "avatar2.png", chat.Chat["assistant_avatar"])
|
|
// } else {
|
|
// assert.Nil(t, chat.Chat["assistant_id"])
|
|
// assert.Nil(t, chat.Chat["assistant_name"])
|
|
// assert.Nil(t, chat.Chat["assistant_avatar"])
|
|
// }
|
|
// }
|
|
|
|
// // Test GetChats
|
|
// filter := ChatFilter{
|
|
// PageSize: 10,
|
|
// Order: "desc",
|
|
// }
|
|
// groups, err := store.GetChats(sid, filter)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, groups)
|
|
// assert.Greater(t, len(groups.Groups), 0)
|
|
|
|
// // Verify assistant information in chat list
|
|
// for _, group := range groups.Groups {
|
|
// for _, chat := range group.Chats {
|
|
// if assistantID, ok := chat["assistant_id"].(string); ok && assistantID != "" {
|
|
// if assistantID == "test-assistant-1" {
|
|
// assert.Equal(t, "Test Assistant 1", chat["assistant_name"])
|
|
// assert.Equal(t, "avatar1.png", chat["assistant_avatar"])
|
|
// } else if assistantID == "test-assistant-2" {
|
|
// assert.Equal(t, "Test Assistant 2", chat["assistant_name"])
|
|
// assert.Equal(t, "avatar2.png", chat["assistant_avatar"])
|
|
// }
|
|
// } else {
|
|
// assert.Nil(t, chat["assistant_name"])
|
|
// assert.Nil(t, chat["assistant_avatar"])
|
|
// }
|
|
// }
|
|
// }
|
|
|
|
// // Test with keywords
|
|
// filter.Keywords = "test"
|
|
// groups, err = store.GetChats(sid, filter)
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(groups.Groups), 0)
|
|
// }
|
|
|
|
// func TestXunDeleteChat(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Create a test chat
|
|
// sid := "test_user"
|
|
// cid := "test_chat"
|
|
// messages := []map[string]interface{}{
|
|
// {"role": "user", "content": "test message"},
|
|
// }
|
|
|
|
// // Save the chat and history
|
|
// err = store.SaveHistory(sid, messages, cid, nil)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify chat exists
|
|
// chat, err := store.GetChat(sid, cid)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, chat)
|
|
|
|
// // Delete the chat
|
|
// err = store.DeleteChat(sid, cid)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify chat is deleted
|
|
// chat, err = store.GetChat(sid, cid)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, (*ChatInfo)(nil), chat)
|
|
// }
|
|
|
|
// func TestXunDeleteAllChats(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Create multiple test chats
|
|
// sid := "test_user"
|
|
// messages := []map[string]interface{}{
|
|
// {"role": "user", "content": "test message"},
|
|
// }
|
|
|
|
// // Save multiple chats
|
|
// for i := 0; i < 3; i++ {
|
|
// cid := fmt.Sprintf("test_chat_%d", i)
|
|
// err = store.SaveHistory(sid, messages, cid, nil)
|
|
// assert.Nil(t, err)
|
|
// }
|
|
|
|
// // Verify chats exist
|
|
// response, err := store.GetChats(sid, ChatFilter{})
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, response.Total, int64(0))
|
|
|
|
// // Delete all chats
|
|
// err = store.DeleteAllChats(sid)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify all chats are deleted
|
|
// response, err = store.GetChats(sid, ChatFilter{})
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, int64(0), response.Total)
|
|
// }
|
|
|
|
// func TestXunAssistantCRUD(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
|
|
// // Drop assistant table before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Clean up any existing data
|
|
// _, err = store.DeleteAssistants(AssistantFilter{})
|
|
// assert.Nil(t, err)
|
|
|
|
// // Test case 1: JSON fields as strings
|
|
// tagsJSON := `["tag1", "tag2", "tag3"]`
|
|
// optionsJSON := `{"model": "gpt-4"}`
|
|
// placeholderJSON := `{"title": "Test Title", "description": "Test Description", "prompts": ["prompt1", "prompt2"]}`
|
|
// assistant := map[string]interface{}{
|
|
// "name": "Test Assistant",
|
|
// "type": "assistant",
|
|
// "avatar": "https://example.com/avatar.png",
|
|
// "connector": "openai",
|
|
// "description": "Test Description",
|
|
// "path": "/assistants/test",
|
|
// "sort": 100,
|
|
// "built_in": true,
|
|
// "tags": tagsJSON,
|
|
// "options": optionsJSON,
|
|
// "placeholder": placeholderJSON,
|
|
// "mentionable": true,
|
|
// "automated": true,
|
|
// }
|
|
|
|
// // Test SaveAssistant (Create) with string JSON
|
|
// v, err := store.SaveAssistant(assistant)
|
|
// assert.Nil(t, err)
|
|
// assistantID := v.(string)
|
|
// assert.NotEmpty(t, assistantID)
|
|
|
|
// // Test GetAssistant for the first assistant
|
|
// assistantData, err := store.GetAssistant(assistantID)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, assistantData)
|
|
// assert.Equal(t, "Test Assistant", assistantData["name"])
|
|
// assert.Equal(t, "assistant", assistantData["type"])
|
|
// assert.Equal(t, "https://example.com/avatar.png", assistantData["avatar"])
|
|
// assert.Equal(t, "openai", assistantData["connector"])
|
|
// assert.Equal(t, "Test Description", assistantData["description"])
|
|
// assert.Equal(t, "/assistants/test", assistantData["path"])
|
|
// assert.Equal(t, int64(100), assistantData["sort"])
|
|
// assert.Equal(t, int64(1), assistantData["built_in"])
|
|
// assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, assistantData["tags"])
|
|
// assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, assistantData["options"])
|
|
// assert.Equal(t, map[string]interface{}{
|
|
// "title": "Test Title",
|
|
// "description": "Test Description",
|
|
// "prompts": []interface{}{"prompt1", "prompt2"},
|
|
// }, assistantData["placeholder"])
|
|
// assert.Equal(t, int64(1), assistantData["mentionable"])
|
|
// assert.Equal(t, int64(1), assistantData["automated"])
|
|
|
|
// // Test case 2: JSON fields as native types
|
|
// assistant2 := map[string]interface{}{
|
|
// "name": "Test Assistant 2",
|
|
// "type": "assistant",
|
|
// "avatar": "https://example.com/avatar2.png",
|
|
// "connector": "openai",
|
|
// "description": "Test Description 2",
|
|
// "path": "/assistants/test2",
|
|
// "sort": 200,
|
|
// "built_in": false,
|
|
// "tags": []string{"tag1", "tag2", "tag3"},
|
|
// "options": map[string]interface{}{"model": "gpt-4"},
|
|
// "prompts": []string{"prompt1", "prompt2"},
|
|
// "workflow": []string{"flow1", "flow2"},
|
|
// "knowledge": []string{"file1", "file2"},
|
|
// "tools": []map[string]interface{}{{"name": "tool1"}, {"name": "tool2"}},
|
|
// "permissions": map[string]interface{}{"read": true, "write": true},
|
|
// "placeholder": map[string]interface{}{
|
|
// "title": "Test Title 2",
|
|
// "description": "Test Description 2",
|
|
// "prompts": []string{"prompt3", "prompt4"},
|
|
// },
|
|
// "mentionable": true,
|
|
// "automated": true,
|
|
// }
|
|
|
|
// // Test SaveAssistant (Create) with native types
|
|
// v, err = store.SaveAssistant(assistant2)
|
|
// assert.Nil(t, err)
|
|
// assistant2ID := v.(string)
|
|
// assert.NotEmpty(t, assistant2ID)
|
|
|
|
// // Test case 3: Test with nil JSON fields
|
|
// assistant3 := map[string]interface{}{
|
|
// "name": "Test Assistant 3",
|
|
// "type": "assistant",
|
|
// "connector": "openai",
|
|
// "description": "Test Description 3",
|
|
// "path": nil,
|
|
// "sort": 9999,
|
|
// "built_in": false,
|
|
// "tags": nil,
|
|
// "options": nil,
|
|
// "prompts": nil,
|
|
// "workflow": nil,
|
|
// "knowledge": nil,
|
|
// "tools": nil,
|
|
// "permissions": nil,
|
|
// "placeholder": nil,
|
|
// "mentionable": true,
|
|
// "automated": true,
|
|
// }
|
|
|
|
// // Test SaveAssistant (Create) with nil fields
|
|
// v, err = store.SaveAssistant(assistant3)
|
|
// assert.Nil(t, err)
|
|
// assistant3ID := v.(string)
|
|
// assert.NotEmpty(t, assistant3ID)
|
|
|
|
// // Test GetAssistant for the third assistant
|
|
// assistant3Data, err := store.GetAssistant(assistant3ID)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, assistant3Data)
|
|
// assert.Equal(t, "Test Assistant 3", assistant3Data["name"])
|
|
// assert.Nil(t, assistant3Data["tags"])
|
|
// assert.Nil(t, assistant3Data["options"])
|
|
// assert.Nil(t, assistant3Data["prompts"])
|
|
// assert.Nil(t, assistant3Data["workflow"])
|
|
// assert.Nil(t, assistant3Data["knowledge"])
|
|
// assert.Nil(t, assistant3Data["tools"])
|
|
// assert.Nil(t, assistant3Data["permissions"])
|
|
// assert.Nil(t, assistant3Data["placeholder"])
|
|
// assert.Equal(t, int64(1), assistant3Data["mentionable"])
|
|
// assert.Equal(t, int64(1), assistant3Data["automated"])
|
|
|
|
// // Test GetAssistant with non-existent ID
|
|
// nonExistentData, err := store.GetAssistant("non-existent-id")
|
|
// assert.Error(t, err)
|
|
// assert.Nil(t, nonExistentData)
|
|
// assert.Contains(t, err.Error(), "is empty")
|
|
|
|
// // Test GetAssistants to verify JSON fields are properly stored
|
|
// resp, err := store.GetAssistants(AssistantFilter{})
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 3, len(resp.Data))
|
|
|
|
// // Clean up all test data
|
|
// _, err = store.DeleteAssistants(AssistantFilter{})
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify cleanup
|
|
// resp, err = store.GetAssistants(AssistantFilter{})
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
// }
|
|
|
|
// func TestXunAssistantPagination(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
|
|
// // Drop assistant table before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Create test data for filtering tests
|
|
// testAssistants := []map[string]interface{}{}
|
|
// for i := 0; i < 25; i++ {
|
|
// assistant := map[string]interface{}{
|
|
// "name": fmt.Sprintf("Filter Test Assistant %d", i),
|
|
// "type": "assistant",
|
|
// "connector": fmt.Sprintf("connector%d", i%3),
|
|
// "description": fmt.Sprintf("Filter Test Description %d", i),
|
|
// "tags": []string{fmt.Sprintf("tag%d", i%5)},
|
|
// "built_in": i%2 == 0,
|
|
// "mentionable": i%2 == 0,
|
|
// "automated": i%3 == 0,
|
|
// "sort": 9999 - i,
|
|
// }
|
|
// id, err := store.SaveAssistant(assistant)
|
|
// assert.Nil(t, err)
|
|
// assistant["assistant_id"] = id
|
|
// testAssistants = append(testAssistants, assistant)
|
|
// }
|
|
|
|
// // Get first assistant ID for later use
|
|
// firstAssistantID := testAssistants[0]["assistant_id"].(string)
|
|
|
|
// // Test filtering with assistantIDs
|
|
// assistantIDs := []string{firstAssistantID}
|
|
// if len(testAssistants) > 1 {
|
|
// assistantIDs = append(assistantIDs, testAssistants[1]["assistant_id"].(string))
|
|
// }
|
|
|
|
// // Test multiple assistant_ids
|
|
// resp, err := store.GetAssistants(AssistantFilter{
|
|
// AssistantIDs: assistantIDs,
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, len(assistantIDs), len(resp.Data))
|
|
// for _, assistant := range resp.Data {
|
|
// found := false
|
|
// for _, id := range assistantIDs {
|
|
// if assistant["assistant_id"] == id {
|
|
// found = true
|
|
// break
|
|
// }
|
|
// }
|
|
// assert.True(t, found, "Assistant ID should be in the requested list")
|
|
// }
|
|
|
|
// // Test assistantIDs with other filters
|
|
// resp, err = store.GetAssistants(AssistantFilter{
|
|
// AssistantIDs: assistantIDs,
|
|
// Select: []string{"name", "assistant_id", "description"},
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, len(assistantIDs), len(resp.Data))
|
|
// // Verify only selected fields are returned
|
|
// for _, item := range resp.Data {
|
|
// assert.Contains(t, item, "name")
|
|
// assert.Contains(t, item, "assistant_id")
|
|
// assert.Contains(t, item, "description")
|
|
// assert.NotContains(t, item, "tags")
|
|
// assert.NotContains(t, item, "options")
|
|
// }
|
|
|
|
// // Test filtering with select fields
|
|
// resp, err = store.GetAssistants(AssistantFilter{
|
|
// Select: []string{"name", "description", "tags"},
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 10, len(resp.Data))
|
|
|
|
// // Test filtering with select fields and other filters combined
|
|
// resp, err = store.GetAssistants(AssistantFilter{
|
|
// Tags: []string{"tag0"},
|
|
// Keywords: "Filter Test",
|
|
// Select: []string{"name", "tags"},
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test combined filters
|
|
// mentionableTrue := true
|
|
// automatedTrue := true
|
|
// resp, err = store.GetAssistants(AssistantFilter{
|
|
// Tags: []string{"tag0"},
|
|
// Keywords: "Filter Test",
|
|
// Connector: "connector0",
|
|
// Mentionable: &mentionableTrue,
|
|
// Automated: &automatedTrue,
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
|
|
// // Now test the delete operations
|
|
// // Test delete by connector
|
|
// var count int64
|
|
// count, err = store.DeleteAssistants(AssistantFilter{
|
|
// Connector: "connector0",
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, count, int64(0))
|
|
|
|
// // Verify deletion
|
|
// resp, err = store.GetAssistants(AssistantFilter{
|
|
// Connector: "connector0",
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
|
|
// // Test delete by built_in status
|
|
// builtInTrue := true
|
|
// count, err = store.DeleteAssistants(AssistantFilter{
|
|
// BuiltIn: &builtInTrue,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, count, int64(0))
|
|
|
|
// // Verify deletion
|
|
// resp, err = store.GetAssistants(AssistantFilter{
|
|
// BuiltIn: &builtInTrue,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
|
|
// // Test delete by tags
|
|
// count, err = store.DeleteAssistants(AssistantFilter{
|
|
// Tags: []string{"tag1"},
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, count, int64(0))
|
|
|
|
// // Verify deletion
|
|
// resp, err = store.GetAssistants(AssistantFilter{
|
|
// Tags: []string{"tag1"},
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
|
|
// // Test delete by keywords
|
|
// count, err = store.DeleteAssistants(AssistantFilter{
|
|
// Keywords: "Filter Test",
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, count, int64(0))
|
|
|
|
// // Verify all assistants are deleted
|
|
// resp, err = store.GetAssistants(AssistantFilter{})
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
|
|
// // Test delete by assistantIDs
|
|
// // First create some test assistants
|
|
// testIDs := []string{}
|
|
// for i := 0; i < 3; i++ {
|
|
// assistant := map[string]interface{}{
|
|
// "name": fmt.Sprintf("AssistantIDs Test Assistant %d", i),
|
|
// "type": "assistant",
|
|
// "connector": "test",
|
|
// "description": fmt.Sprintf("AssistantIDs Test Description %d", i),
|
|
// "tags": []string{"test-tag"},
|
|
// "built_in": false,
|
|
// "mentionable": true,
|
|
// "automated": true,
|
|
// }
|
|
// id, err := store.SaveAssistant(assistant)
|
|
// assert.Nil(t, err)
|
|
// testIDs = append(testIDs, id.(string))
|
|
// }
|
|
|
|
// // Delete by assistantIDs
|
|
// count, err = store.DeleteAssistants(AssistantFilter{
|
|
// AssistantIDs: testIDs,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, int64(len(testIDs)), count)
|
|
|
|
// // Verify deletion
|
|
// resp, err = store.GetAssistants(AssistantFilter{
|
|
// AssistantIDs: testIDs,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
|
|
// // Verify all assistants are deleted
|
|
// resp, err = store.GetAssistants(AssistantFilter{})
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
// }
|
|
|
|
// func TestGetAssistantTags(t *testing.T) {
|
|
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Create test assistants with tags
|
|
// assistants := []map[string]interface{}{
|
|
// {
|
|
// "assistant_id": "test-assistant-1",
|
|
// "type": "assistant",
|
|
// "connector": "test",
|
|
// "tags": []string{"tag1", "tag2"},
|
|
// "name": "Test Assistant 1",
|
|
// },
|
|
// {
|
|
// "assistant_id": "test-assistant-2",
|
|
// "type": "assistant",
|
|
// "connector": "test",
|
|
// "tags": []string{"tag2", "tag3"},
|
|
// "name": "Test Assistant 2",
|
|
// },
|
|
// {
|
|
// "assistant_id": "test-assistant-3",
|
|
// "type": "assistant",
|
|
// "connector": "test",
|
|
// "tags": []string{"tag1", "tag3", "tag4"},
|
|
// "name": "Test Assistant 3",
|
|
// },
|
|
// }
|
|
|
|
// // Save test assistants
|
|
// for _, assistant := range assistants {
|
|
// _, err := store.SaveAssistant(assistant)
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// }
|
|
|
|
// // Get tags
|
|
// tags, err := store.GetAssistantTags()
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Verify results
|
|
// expectedTags := map[string]bool{
|
|
// "tag1": true,
|
|
// "tag2": true,
|
|
// "tag3": true,
|
|
// "tag4": true,
|
|
// }
|
|
|
|
// if len(tags) != len(expectedTags) {
|
|
// t.Errorf("Expected %d tags, got %d", len(expectedTags), len(tags))
|
|
// }
|
|
|
|
// for _, tag := range tags {
|
|
// value := tag.Value
|
|
// if !expectedTags[value] {
|
|
// t.Errorf("Unexpected tag found: %s", tag)
|
|
// }
|
|
// }
|
|
|
|
// // Cleanup
|
|
// for _, assistant := range assistants {
|
|
// err := store.DeleteAssistant(assistant["assistant_id"].(string))
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// }
|
|
// }
|
|
|
|
// func TestXunSaveAndGetHistoryWithSilent(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// TTL: 3600,
|
|
// })
|
|
|
|
// // save the history with silent messages
|
|
// sid := "123456"
|
|
// cid := "silent_test"
|
|
|
|
// // First save regular messages
|
|
// messages := []map[string]interface{}{
|
|
// {"role": "user", "name": "user1", "content": "hello"},
|
|
// {"role": "assistant", "name": "assistant1", "content": "Hi! How can I help you?"},
|
|
// }
|
|
// context := map[string]interface{}{
|
|
// "assistant_id": "test-assistant-1",
|
|
// }
|
|
// err = store.SaveHistory(sid, messages, cid, context)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Then save silent messages
|
|
// silentMessages := []map[string]interface{}{
|
|
// {"role": "user", "name": "user1", "content": "silent message"},
|
|
// {"role": "assistant", "name": "assistant1", "content": "This is a silent response"},
|
|
// }
|
|
// silentContext := map[string]interface{}{
|
|
// "assistant_id": "test-assistant-1",
|
|
// "silent": true,
|
|
// }
|
|
// err = store.SaveHistory(sid, silentMessages, cid, silentContext)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Get history without filter (should only return non-silent messages)
|
|
// data, err := store.GetHistory(sid, cid)
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, 2, len(data))
|
|
// for _, msg := range data {
|
|
// // Check if silent is false, handling different types
|
|
// isSilent := false
|
|
// switch v := msg["silent"].(type) {
|
|
// case bool:
|
|
// isSilent = v
|
|
// case int:
|
|
// isSilent = v != 0
|
|
// case int64:
|
|
// isSilent = v != 0
|
|
// case float64:
|
|
// isSilent = v != 0
|
|
// }
|
|
// assert.False(t, isSilent, "message should not be silent")
|
|
// }
|
|
|
|
// // Get history with silent=true filter (should return all messages)
|
|
// silentTrue := true
|
|
// filter := ChatFilter{
|
|
// Silent: &silentTrue,
|
|
// }
|
|
// allData, err := store.GetHistoryWithFilter(sid, cid, filter)
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// assert.Equal(t, 4, len(allData))
|
|
|
|
// // Count silent messages
|
|
// silentCount := 0
|
|
// for _, msg := range allData {
|
|
// // Check if silent is true, handling different types
|
|
// isSilent := false
|
|
// switch v := msg["silent"].(type) {
|
|
// case bool:
|
|
// isSilent = v
|
|
// case int:
|
|
// isSilent = v != 0
|
|
// case int64:
|
|
// isSilent = v != 0
|
|
// case float64:
|
|
// isSilent = v != 0
|
|
// }
|
|
// if isSilent {
|
|
// silentCount++
|
|
// }
|
|
// }
|
|
// assert.Equal(t, 2, silentCount)
|
|
|
|
// // Get chat with filter (should include silent messages)
|
|
// chat, err := store.GetChatWithFilter(sid, cid, filter)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 4, len(chat.History))
|
|
|
|
// // Get chat without filter (should exclude silent messages)
|
|
// chatNoSilent, err := store.GetChat(sid, cid)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 2, len(chatNoSilent.History))
|
|
// }
|
|
|
|
// func TestXunGetChatsWithSilent(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
|
|
// // Drop tables before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Create test assistant
|
|
// assistant := map[string]interface{}{
|
|
// "assistant_id": "test-assistant-1",
|
|
// "name": "Test Assistant 1",
|
|
// "avatar": "avatar1.png",
|
|
// "type": "assistant",
|
|
// "connector": "test",
|
|
// }
|
|
// _, err = store.SaveAssistant(assistant)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Save some test chats
|
|
// sid := "test_user"
|
|
// messages := []map[string]interface{}{
|
|
// {"role": "user", "content": "test message"},
|
|
// }
|
|
|
|
// // Create regular chats
|
|
// for i := 0; i < 3; i++ {
|
|
// chatID := fmt.Sprintf("regular_chat_%d", i)
|
|
// title := fmt.Sprintf("Regular Chat %d", i)
|
|
// context := map[string]interface{}{
|
|
// "assistant_id": "test-assistant-1",
|
|
// "silent": false,
|
|
// }
|
|
|
|
// // Save history to create the chat
|
|
// err = store.SaveHistory(sid, messages, chatID, context)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Update the chat title
|
|
// err = store.UpdateChatTitle(sid, chatID, title)
|
|
// assert.Nil(t, err)
|
|
// }
|
|
|
|
// // Create silent chats
|
|
// for i := 0; i < 2; i++ {
|
|
// chatID := fmt.Sprintf("silent_chat_%d", i)
|
|
// title := fmt.Sprintf("Silent Chat %d", i)
|
|
// context := map[string]interface{}{
|
|
// "assistant_id": "test-assistant-1",
|
|
// "silent": true,
|
|
// }
|
|
|
|
// // Save history to create the chat
|
|
// err = store.SaveHistory(sid, messages, chatID, context)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Update the chat title
|
|
// err = store.UpdateChatTitle(sid, chatID, title)
|
|
// assert.Nil(t, err)
|
|
// }
|
|
|
|
// // Test GetChats with default filter (should exclude silent chats)
|
|
// defaultFilter := ChatFilter{
|
|
// PageSize: 10,
|
|
// Order: "desc",
|
|
// }
|
|
// defaultGroups, err := store.GetChats(sid, defaultFilter)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, defaultGroups)
|
|
|
|
// // Count total chats in all groups
|
|
// totalDefaultChats := 0
|
|
// for _, group := range defaultGroups.Groups {
|
|
// totalDefaultChats += len(group.Chats)
|
|
// }
|
|
// assert.Equal(t, 3, totalDefaultChats, "Default filter should only return non-silent chats")
|
|
|
|
// // Test GetChats with silent=true filter (should include all chats)
|
|
// silentTrue := true
|
|
// silentFilter := ChatFilter{
|
|
// PageSize: 10,
|
|
// Order: "desc",
|
|
// Silent: &silentTrue,
|
|
// }
|
|
// silentGroups, err := store.GetChats(sid, silentFilter)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, silentGroups)
|
|
|
|
// // Count total chats in all groups
|
|
// totalSilentChats := 0
|
|
// for _, group := range silentGroups.Groups {
|
|
// totalSilentChats += len(group.Chats)
|
|
// }
|
|
// assert.Equal(t, 5, totalSilentChats, "Silent filter should return all chats")
|
|
|
|
// // Test GetChats with silent=false filter (should only include non-silent chats)
|
|
// silentFalse := false
|
|
// nonSilentFilter := ChatFilter{
|
|
// PageSize: 10,
|
|
// Order: "desc",
|
|
// Silent: &silentFalse,
|
|
// }
|
|
// nonSilentGroups, err := store.GetChats(sid, nonSilentFilter)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, nonSilentGroups)
|
|
|
|
// // Count total chats in all groups
|
|
// totalNonSilentChats := 0
|
|
// for _, group := range nonSilentGroups.Groups {
|
|
// totalNonSilentChats += len(group.Chats)
|
|
// }
|
|
// assert.Equal(t, 3, totalNonSilentChats, "Non-silent filter should only return non-silent chats")
|
|
// }
|
|
|
|
// func TestXunAttachmentCRUD(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment")
|
|
|
|
// // Drop attachment table before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Clean up any existing data
|
|
// _, err = store.DeleteAttachments(AttachmentFilter{})
|
|
// assert.Nil(t, err)
|
|
|
|
// // Test SaveAttachment (Create)
|
|
// attachment := map[string]interface{}{
|
|
// "file_id": "test-file-123",
|
|
// "uid": "user-123",
|
|
// "manager": "local",
|
|
// "content_type": "image/jpeg",
|
|
// "name": "test-image.jpg",
|
|
// "guest": false,
|
|
// "public": true,
|
|
// "gzip": false,
|
|
// "bytes": 102400,
|
|
// "scope": []string{"user", "admin"},
|
|
// "status": "uploaded",
|
|
// "progress": "100%",
|
|
// "error": nil,
|
|
// }
|
|
|
|
// v, err := store.SaveAttachment(attachment)
|
|
// assert.Nil(t, err)
|
|
// fileID := v.(string)
|
|
// assert.Equal(t, "test-file-123", fileID)
|
|
|
|
// // Test GetAttachment
|
|
// attachmentData, err := store.GetAttachment(fileID)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, attachmentData)
|
|
// assert.Equal(t, "test-file-123", attachmentData["file_id"])
|
|
// assert.Equal(t, "user-123", attachmentData["uid"])
|
|
// assert.Equal(t, "local", attachmentData["manager"])
|
|
// assert.Equal(t, "image/jpeg", attachmentData["content_type"])
|
|
// assert.Equal(t, "test-image.jpg", attachmentData["name"])
|
|
// assert.Equal(t, int64(1), attachmentData["public"])
|
|
// assert.Equal(t, []interface{}{"user", "admin"}, attachmentData["scope"])
|
|
// assert.Equal(t, "uploaded", attachmentData["status"])
|
|
// assert.Equal(t, "100%", attachmentData["progress"])
|
|
// assert.Nil(t, attachmentData["error"])
|
|
|
|
// // Test SaveAttachment (Update)
|
|
// attachment["name"] = "updated-image.jpg"
|
|
// attachment["bytes"] = 204800
|
|
// attachment["status"] = "indexing"
|
|
// attachment["progress"] = "Processing..."
|
|
// attachment["error"] = "Connection timeout"
|
|
// v, err = store.SaveAttachment(attachment)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, "test-file-123", v.(string))
|
|
|
|
// // Verify update
|
|
// attachmentData, err = store.GetAttachment(fileID)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, "updated-image.jpg", attachmentData["name"])
|
|
// assert.Equal(t, int64(204800), attachmentData["bytes"])
|
|
// assert.Equal(t, "indexing", attachmentData["status"])
|
|
// assert.Equal(t, "Processing...", attachmentData["progress"])
|
|
// assert.Equal(t, "Connection timeout", attachmentData["error"])
|
|
|
|
// // Test GetAttachments with filters
|
|
// resp, err := store.GetAttachments(AttachmentFilter{
|
|
// UID: "user-123",
|
|
// Manager: "local",
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 1, len(resp.Data))
|
|
// assert.Equal(t, "test-file-123", resp.Data[0]["file_id"])
|
|
|
|
// // Test with non-existent file
|
|
// nonExistentData, err := store.GetAttachment("non-existent-file")
|
|
// assert.Error(t, err)
|
|
// assert.Nil(t, nonExistentData)
|
|
// assert.Contains(t, err.Error(), "is empty")
|
|
|
|
// // Test DeleteAttachment
|
|
// err = store.DeleteAttachment(fileID)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify deletion
|
|
// _, err = store.GetAttachment(fileID)
|
|
// assert.Error(t, err)
|
|
// }
|
|
|
|
// func TestXunKnowledgeCRUD(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_knowledge")
|
|
|
|
// // Drop knowledge table before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_knowledge")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Clean up any existing data
|
|
// _, err = store.DeleteKnowledges(KnowledgeFilter{})
|
|
// assert.Nil(t, err)
|
|
|
|
// // Test SaveKnowledge (Create)
|
|
// knowledge := map[string]interface{}{
|
|
// "collection_id": "test-collection-123",
|
|
// "name": "Test Knowledge Collection",
|
|
// "description": "A test knowledge collection for unit tests",
|
|
// "uid": "user-123",
|
|
// "public": true,
|
|
// "readonly": false,
|
|
// "system": false,
|
|
// "sort": 100,
|
|
// "cover": "cover-image.jpg",
|
|
// "scope": []string{"user", "admin"},
|
|
// "option": map[string]interface{}{"embedding": "openai", "chunk_size": 1000},
|
|
// }
|
|
|
|
// v, err := store.SaveKnowledge(knowledge)
|
|
// assert.Nil(t, err)
|
|
// collectionID := v.(string)
|
|
// assert.Equal(t, "test-collection-123", collectionID)
|
|
|
|
// // Test GetKnowledge
|
|
// knowledgeData, err := store.GetKnowledge(collectionID)
|
|
// assert.Nil(t, err)
|
|
// assert.NotNil(t, knowledgeData)
|
|
// assert.Equal(t, "test-collection-123", knowledgeData["collection_id"])
|
|
// assert.Equal(t, "Test Knowledge Collection", knowledgeData["name"])
|
|
// assert.Equal(t, "A test knowledge collection for unit tests", knowledgeData["description"])
|
|
// assert.Equal(t, "user-123", knowledgeData["uid"])
|
|
// assert.Equal(t, int64(1), knowledgeData["public"])
|
|
// assert.Equal(t, int64(100), knowledgeData["sort"])
|
|
// assert.Equal(t, []interface{}{"user", "admin"}, knowledgeData["scope"])
|
|
// assert.Equal(t, map[string]interface{}{"embedding": "openai", "chunk_size": float64(1000)}, knowledgeData["option"])
|
|
|
|
// // Test SaveKnowledge (Update)
|
|
// knowledge["name"] = "Updated Knowledge Collection"
|
|
// knowledge["description"] = "Updated description"
|
|
// knowledge["sort"] = 200
|
|
// v, err = store.SaveKnowledge(knowledge)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, "test-collection-123", v.(string))
|
|
|
|
// // Verify update
|
|
// knowledgeData, err = store.GetKnowledge(collectionID)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, "Updated Knowledge Collection", knowledgeData["name"])
|
|
// assert.Equal(t, "Updated description", knowledgeData["description"])
|
|
// assert.Equal(t, int64(200), knowledgeData["sort"])
|
|
|
|
// // Test knowledge without sort field (should get default value 9999)
|
|
// knowledgeWithoutSort := map[string]interface{}{
|
|
// "collection_id": "test-collection-456",
|
|
// "name": "Test Knowledge Without Sort",
|
|
// "description": "Test knowledge without explicit sort value",
|
|
// "uid": "user-123",
|
|
// }
|
|
// v2, err := store.SaveKnowledge(knowledgeWithoutSort)
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify default sort value
|
|
// knowledgeData2, err := store.GetKnowledge(v2.(string))
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, int64(9999), knowledgeData2["sort"])
|
|
|
|
// // Test GetKnowledges with filters
|
|
// resp, err := store.GetKnowledges(KnowledgeFilter{
|
|
// UID: "user-123",
|
|
// Keywords: "Updated",
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 1, len(resp.Data))
|
|
// assert.Equal(t, "test-collection-123", resp.Data[0]["collection_id"])
|
|
|
|
// // Test with non-existent collection
|
|
// nonExistentData, err := store.GetKnowledge("non-existent-collection")
|
|
// assert.Error(t, err)
|
|
// assert.Nil(t, nonExistentData)
|
|
// assert.Contains(t, err.Error(), "is empty")
|
|
|
|
// // Test DeleteKnowledge
|
|
// err = store.DeleteKnowledge(collectionID)
|
|
// assert.Nil(t, err)
|
|
// err = store.DeleteKnowledge(v2.(string))
|
|
// assert.Nil(t, err)
|
|
|
|
// // Verify deletion
|
|
// _, err = store.GetKnowledge(collectionID)
|
|
// assert.Error(t, err)
|
|
// }
|
|
|
|
// func TestXunKnowledgeFiltering(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_knowledge")
|
|
|
|
// // Drop knowledge table before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_knowledge")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Create test data for filtering tests
|
|
// testKnowledges := []map[string]interface{}{}
|
|
// for i := 0; i < 15; i++ {
|
|
// knowledge := map[string]interface{}{
|
|
// "collection_id": fmt.Sprintf("test-collection-%d", i),
|
|
// "name": fmt.Sprintf("Collection %d", i),
|
|
// "description": fmt.Sprintf("Description for collection %d", i),
|
|
// "uid": fmt.Sprintf("user-%d", i%3),
|
|
// "public": i%2 == 0,
|
|
// "readonly": i%3 == 0,
|
|
// "system": i%4 == 0,
|
|
// "sort": 100 + i*10, // Different sort values for testing ordering
|
|
// "cover": fmt.Sprintf("cover%d.jpg", i),
|
|
// }
|
|
// id, err := store.SaveKnowledge(knowledge)
|
|
// assert.Nil(t, err)
|
|
// knowledge["collection_id"] = id
|
|
// testKnowledges = append(testKnowledges, knowledge)
|
|
// }
|
|
|
|
// // Test sorting functionality - should return results ordered by sort ASC then created_at DESC
|
|
// respAll, err := store.GetKnowledges(KnowledgeFilter{
|
|
// Page: 1,
|
|
// PageSize: 15,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 15, len(respAll.Data))
|
|
|
|
// // Verify sort order - first item should have the smallest sort value
|
|
// firstSort := respAll.Data[0]["sort"].(int64)
|
|
// lastSort := respAll.Data[len(respAll.Data)-1]["sort"].(int64)
|
|
// assert.LessOrEqual(t, firstSort, lastSort, "Results should be ordered by sort ASC")
|
|
|
|
// // More specific sort order verification
|
|
// for i := 1; i < len(respAll.Data); i++ {
|
|
// prevSort := respAll.Data[i-1]["sort"].(int64)
|
|
// currSort := respAll.Data[i]["sort"].(int64)
|
|
// assert.LessOrEqual(t, prevSort, currSort, "Sort order should be ascending")
|
|
// }
|
|
|
|
// // Test filtering by UID
|
|
// resp, err := store.GetKnowledges(KnowledgeFilter{
|
|
// UID: "user-0",
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by public status
|
|
// publicTrue := true
|
|
// resp, err = store.GetKnowledges(KnowledgeFilter{
|
|
// Public: &publicTrue,
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by readonly status
|
|
// readonlyTrue := true
|
|
// resp, err = store.GetKnowledges(KnowledgeFilter{
|
|
// Readonly: &readonlyTrue,
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by system status
|
|
// systemTrue := true
|
|
// resp, err = store.GetKnowledges(KnowledgeFilter{
|
|
// System: &systemTrue,
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by keywords
|
|
// resp, err = store.GetKnowledges(KnowledgeFilter{
|
|
// Keywords: "Collection 1",
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test DeleteKnowledges with filter
|
|
// count, err := store.DeleteKnowledges(KnowledgeFilter{
|
|
// UID: "user-0",
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, count, int64(0))
|
|
|
|
// // Verify deletion
|
|
// resp, err = store.GetKnowledges(KnowledgeFilter{
|
|
// UID: "user-0",
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
|
|
// // Clean up all test data
|
|
// _, err = store.DeleteKnowledges(KnowledgeFilter{})
|
|
// assert.Nil(t, err)
|
|
// }
|
|
|
|
// func TestXunAttachmentFiltering(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment")
|
|
|
|
// // Drop attachment table before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Create test data for filtering tests
|
|
// testAttachments := []map[string]interface{}{}
|
|
// for i := 0; i < 15; i++ {
|
|
// attachment := map[string]interface{}{
|
|
// "file_id": fmt.Sprintf("test-file-%d", i),
|
|
// "uid": fmt.Sprintf("user-%d", i%3),
|
|
// "manager": fmt.Sprintf("manager%d", i%2),
|
|
// "content_type": fmt.Sprintf("type/%d", i%4),
|
|
// "name": fmt.Sprintf("file%d.txt", i),
|
|
// "guest": i%2 == 0,
|
|
// "public": i%3 == 0,
|
|
// "gzip": i%4 == 0,
|
|
// "bytes": 1024 * (i + 1),
|
|
// "collection_id": fmt.Sprintf("collection-%d", i%5),
|
|
// }
|
|
// id, err := store.SaveAttachment(attachment)
|
|
// assert.Nil(t, err)
|
|
// attachment["file_id"] = id
|
|
// testAttachments = append(testAttachments, attachment)
|
|
// }
|
|
|
|
// // Test filtering by UID
|
|
// resp, err := store.GetAttachments(AttachmentFilter{
|
|
// UID: "user-0",
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by manager
|
|
// resp, err = store.GetAttachments(AttachmentFilter{
|
|
// Manager: "manager0",
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by content_type
|
|
// resp, err = store.GetAttachments(AttachmentFilter{
|
|
// ContentType: "type/0",
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by guest status
|
|
// guestTrue := true
|
|
// resp, err = store.GetAttachments(AttachmentFilter{
|
|
// Guest: &guestTrue,
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by public status
|
|
// publicTrue := true
|
|
// resp, err = store.GetAttachments(AttachmentFilter{
|
|
// Public: &publicTrue,
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test filtering by keywords
|
|
// resp, err = store.GetAttachments(AttachmentFilter{
|
|
// Keywords: "file1",
|
|
// Page: 1,
|
|
// PageSize: 10,
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, len(resp.Data), 0)
|
|
|
|
// // Test DeleteAttachments with filter
|
|
// count, err := store.DeleteAttachments(AttachmentFilter{
|
|
// Manager: "manager0",
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Greater(t, count, int64(0))
|
|
|
|
// // Verify deletion
|
|
// resp, err = store.GetAttachments(AttachmentFilter{
|
|
// Manager: "manager0",
|
|
// })
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, 0, len(resp.Data))
|
|
|
|
// // Clean up all test data
|
|
// _, err = store.DeleteAttachments(AttachmentFilter{})
|
|
// assert.Nil(t, err)
|
|
// }
|
|
|
|
// func TestXunAttachmentStatusFields(t *testing.T) {
|
|
// test.Prepare(t, config.Conf)
|
|
// defer test.Clean()
|
|
// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment")
|
|
|
|
// // Drop attachment table before test
|
|
// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment")
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Add a small delay to ensure table is created
|
|
// time.Sleep(100 * time.Millisecond)
|
|
|
|
// store, err := NewXun(Setting{
|
|
// Connector: "default",
|
|
// Prefix: "__unit_test_conversation_",
|
|
// })
|
|
// if err != nil {
|
|
// t.Fatal(err)
|
|
// }
|
|
|
|
// // Clean up any existing data
|
|
// _, err = store.DeleteAttachments(AttachmentFilter{})
|
|
// assert.Nil(t, err)
|
|
|
|
// // Test all possible enum status values
|
|
// statusValues := []string{"uploading", "uploaded", "indexing", "indexed", "upload_failed", "index_failed"}
|
|
|
|
// for i, status := range statusValues {
|
|
// // Create attachment with specific status
|
|
// attachment := map[string]interface{}{
|
|
// "file_id": fmt.Sprintf("test-file-%s-%d", status, i),
|
|
// "uid": "user-123",
|
|
// "manager": "local",
|
|
// "content_type": "image/jpeg",
|
|
// "name": fmt.Sprintf("test-%s.jpg", status),
|
|
// "guest": false,
|
|
// "public": true,
|
|
// "gzip": false,
|
|
// "bytes": 102400,
|
|
// "status": status,
|
|
// "progress": fmt.Sprintf("%s in progress", status),
|
|
// "error": nil,
|
|
// }
|
|
|
|
// // Set error message for failed statuses
|
|
// if status == "upload_failed" || status == "index_failed" {
|
|
// attachment["error"] = fmt.Sprintf("%s error occurred", status)
|
|
// }
|
|
|
|
// v, err := store.SaveAttachment(attachment)
|
|
// assert.Nil(t, err)
|
|
// fileID := v.(string)
|
|
|
|
// // Verify the attachment was saved with correct status
|
|
// attachmentData, err := store.GetAttachment(fileID)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, status, attachmentData["status"])
|
|
// assert.Equal(t, fmt.Sprintf("%s in progress", status), attachmentData["progress"])
|
|
|
|
// if status == "upload_failed" || status == "index_failed" {
|
|
// assert.Equal(t, fmt.Sprintf("%s error occurred", status), attachmentData["error"])
|
|
// } else {
|
|
// assert.Nil(t, attachmentData["error"])
|
|
// }
|
|
// }
|
|
|
|
// // Test default status value (should be "uploading")
|
|
// attachmentWithoutStatus := map[string]interface{}{
|
|
// "file_id": "test-file-default",
|
|
// "uid": "user-123",
|
|
// "manager": "local",
|
|
// "content_type": "image/jpeg",
|
|
// "name": "test-default.jpg",
|
|
// "guest": false,
|
|
// "public": true,
|
|
// "gzip": false,
|
|
// "bytes": 102400,
|
|
// // status not specified - should use default
|
|
// }
|
|
|
|
// v, err := store.SaveAttachment(attachmentWithoutStatus)
|
|
// assert.Nil(t, err)
|
|
// fileID := v.(string)
|
|
|
|
// // Verify default status
|
|
// attachmentData, err := store.GetAttachment(fileID)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, "uploading", attachmentData["status"]) // Should be default value
|
|
// assert.Nil(t, attachmentData["progress"]) // Should be null
|
|
// assert.Nil(t, attachmentData["error"]) // Should be null
|
|
|
|
// // Test updating status workflow: uploading -> uploaded -> indexing -> indexed
|
|
// workflowAttachment := map[string]interface{}{
|
|
// "file_id": "test-file-workflow",
|
|
// "uid": "user-123",
|
|
// "manager": "local",
|
|
// "content_type": "text/plain",
|
|
// "name": "workflow-test.txt",
|
|
// "status": "uploading",
|
|
// "progress": "Starting upload...",
|
|
// }
|
|
|
|
// v, err = store.SaveAttachment(workflowAttachment)
|
|
// assert.Nil(t, err)
|
|
// workflowFileID := v.(string)
|
|
|
|
// // Update to uploaded
|
|
// workflowAttachment["status"] = "uploaded"
|
|
// workflowAttachment["progress"] = "Upload completed, starting indexing..."
|
|
// _, err = store.SaveAttachment(workflowAttachment)
|
|
// assert.Nil(t, err)
|
|
|
|
// attachmentData, err = store.GetAttachment(workflowFileID)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, "uploaded", attachmentData["status"])
|
|
// assert.Equal(t, "Upload completed, starting indexing...", attachmentData["progress"])
|
|
|
|
// // Update to indexing
|
|
// workflowAttachment["status"] = "indexing"
|
|
// workflowAttachment["progress"] = "Indexing in progress..."
|
|
// _, err = store.SaveAttachment(workflowAttachment)
|
|
// assert.Nil(t, err)
|
|
|
|
// attachmentData, err = store.GetAttachment(workflowFileID)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, "indexing", attachmentData["status"])
|
|
// assert.Equal(t, "Indexing in progress...", attachmentData["progress"])
|
|
|
|
// // Update to indexed (final state)
|
|
// workflowAttachment["status"] = "indexed"
|
|
// workflowAttachment["progress"] = "Indexing completed"
|
|
// _, err = store.SaveAttachment(workflowAttachment)
|
|
// assert.Nil(t, err)
|
|
|
|
// attachmentData, err = store.GetAttachment(workflowFileID)
|
|
// assert.Nil(t, err)
|
|
// assert.Equal(t, "indexed", attachmentData["status"])
|
|
// assert.Equal(t, "Indexing completed", attachmentData["progress"])
|
|
|
|
// // Clean up test data
|
|
// _, err = store.DeleteAttachments(AttachmentFilter{})
|
|
// assert.Nil(t, err)
|
|
// }
|