[Unit Tests] - ProcessChatResponse

[tomsmith8]

Unit Test Coverage for "ProcessChatResponse"

---

Stakwork Run

---

Unit Test Code

---

File: /tmp/stakwork/sphinx-tribes/handlers/chat.go

---

package chat

import (
  "bytes"
  "encoding/json"
  "errors"
  "fmt"
  "net/http"
  "net/http/httptest"
  "testing"
  "time"

  "github.com/gorilla/websocket"
  "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/mock"
)

// Mock dependencies
type MockDB struct {
  mock.Mock
}

func (m *MockDB) AddChatMessage(message *db.ChatMessage) (*db.ChatMessage, error) {
  args := m.Called(message)
  return args.Get(0).(*db.ChatMessage), args.Error(1)
}

type MockWebsocketPool struct {
  mock.Mock
}

func (m *MockWebsocketPool) SendTicketMessage(message websocket.TicketMessage) error {
  args := m.Called(message)
  return args.Error(0)
}

// Test function
func TestProcessChatResponse(t *testing.T) {
  mockDB := new(MockDB)
  mockWebsocketPool := new(MockWebsocketPool)
  handler := &ChatHandler{
  	db: mockDB,
  }

  tests := []struct {
  	name           string
  	input          string
  	mockDBResponse *db.ChatMessage
  	mockDBError    error
  	mockWSResponse error
  	expectedStatus int
  	expectedBody   ChatResponse
  }{
  	{
  		name: "Valid Input",
  		input: `{
  			"value": {
  				"chatId": "validChatId",
  				"messageId": "validMessageId",
  				"response": "This is a response",
  				"sourceWebsocketId": "validWebsocketId"
  			}
  		}`,
  		mockDBResponse: &db.ChatMessage{
  			ID:        "generatedID",
  			ChatID:    "validChatId",
  			Message:   "This is a response",
  			Role:      "assistant",
  			Timestamp: time.Now(),
  			Status:    "sent",
  			Source:    "agent",
  		},
  		mockDBError:    nil,
  		mockWSResponse: nil,
  		expectedStatus: http.StatusOK,
  		expectedBody: ChatResponse{
  			Success: true,
  			Message: "Response processed successfully",
  			Data: &db.ChatMessage{
  				ID:        "generatedID",
  				ChatID:    "validChatId",
  				Message:   "This is a response",
  				Role:      "assistant",
  				Timestamp: time.Now(),
  				Status:    "sent",
  				Source:    "agent",
  			},
  		},
  	},
  	{
  		name: "Empty ChatID",
  		input: `{
  			"value": {
  				"chatId": "",
  				"messageId": "validMessageId",
  				"response": "This is a response",
  				"sourceWebsocketId": "validWebsocketId"
  			}
  		}`,
  		expectedStatus: http.StatusBadRequest,
  		expectedBody: ChatResponse{
  			Success: false,
  			Message: "ChatID is required for message creation",
  		},
  	},
  	{
  		name: "Empty Response",
  		input: `{
  			"value": {
  				"chatId": "validChatId",
  				"messageId": "validMessageId",
  				"response": "",
  				"sourceWebsocketId": "validWebsocketId"
  			}
  		}`,
  		mockDBResponse: &db.ChatMessage{
  			ID:        "generatedID",
  			ChatID:    "validChatId",
  			Message:   "",
  			Role:      "assistant",
  			Timestamp: time.Now(),
  			Status:    "sent",
  			Source:    "agent",
  		},
  		mockDBError:    nil,
  		mockWSResponse: nil,
  		expectedStatus: http.StatusOK,
  		expectedBody: ChatResponse{
  			Success: true,
  			Message: "Response processed successfully",
  			Data: &db.ChatMessage{
  				ID:        "generatedID",
  				ChatID:    "validChatId",
  				Message:   "",
  				Role:      "assistant",
  				Timestamp: time.Now(),
  				Status:    "sent",
  				Source:    "agent",
  			},
  		},
  	},
  	{
  		name: "Invalid JSON Format",
  		input: `{
  			"value": "invalidJson"
  		}`,
  		expectedStatus: http.StatusBadRequest,
  		expectedBody: ChatResponse{
  			Success: false,
  			Message: "Invalid request body",
  		},
  	},
  	{
  		name: "Database Error",
  		input: `{
  			"value": {
  				"chatId": "validChatId",
  				"messageId": "validMessageId",
  				"response": "This is a response",
  				"sourceWebsocketId": "validWebsocketId"
  			}
  		}`,
  		mockDBResponse: nil,
  		mockDBError:    errors.New("database error"),
  		expectedStatus: http.StatusInternalServerError,
  		expectedBody: ChatResponse{
  			Success: false,
  			Message: "Failed to save response message: database error",
  		},
  	},
  	{
  		name: "WebSocket Error",
  		input: `{
  			"value": {
  				"chatId": "validChatId",
  				"messageId": "validMessageId",
  				"response": "This is a response",
  				"sourceWebsocketId": "validWebsocketId"
  			}
  		}`,
  		mockDBResponse: &db.ChatMessage{
  			ID:        "generatedID",
  			ChatID:    "validChatId",
  			Message:   "This is a response",
  			Role:      "assistant",
  			Timestamp: time.Now(),
  			Status:    "sent",
  			Source:    "agent",
  		},
  		mockDBError:    nil,
  		mockWSResponse: errors.New("websocket error"),
  		expectedStatus: http.StatusOK,
  		expectedBody: ChatResponse{
  			Success: true,
  			Message: "Response processed successfully",
  			Data: &db.ChatMessage{
  				ID:        "generatedID",
  				ChatID:    "validChatId",
  				Message:   "This is a response",
  				Role:      "assistant",
  				Timestamp: time.Now(),
  				Status:    "sent",
  				Source:    "agent",
  			},
  		},
  	},
  	{
  		name: "Large Response Message",
  		input: fmt.Sprintf(`{
  			"value": {
  				"chatId": "validChatId",
  				"messageId": "validMessageId",
  				"response": "%s",
  				"sourceWebsocketId": "validWebsocketId"
  			}
  		}`, string(make([]byte, 10000))),
  		mockDBResponse: &db.ChatMessage{
  			ID:        "generatedID",
  			ChatID:    "validChatId",
  			Message:   string(make([]byte, 10000)),
  			Role:      "assistant",
  			Timestamp: time.Now(),
  			Status:    "sent",
  			Source:    "agent",
  		},
  		mockDBError:    nil,
  		mockWSResponse: nil,
  		expectedStatus: http.StatusOK,
  		expectedBody: ChatResponse{
  			Success: true,
  			Message: "Response processed successfully",
  			Data: &db.ChatMessage{
  				ID:        "generatedID",
  				ChatID:    "validChatId",
  				Message:   string(make([]byte, 10000)),
  				Role:      "assistant",
  				Timestamp: time.Now(),
  				Status:    "sent",
  				Source:    "agent",
  			},
  		},
  	},
  	{
  		name: "Missing SourceWebsocketID",
  		input: `{
  			"value": {
  				"chatId": "validChatId",
  				"messageId": "validMessageId",
  				"response": "This is a response"
  			}
  		}`,
  		mockDBResponse: &db.ChatMessage{
  			ID:        "generatedID",
  			ChatID:    "validChatId",
  			Message:   "This is a response",
  			Role:      "assistant",
  			Timestamp: time.Now(),
  			Status:    "sent",
  			Source:    "agent",
  		},
  		mockDBError:    nil,
  		mockWSResponse: nil,
  		expectedStatus: http.StatusOK,
  		expectedBody: ChatResponse{
  			Success: true,
  			Message: "Response processed successfully",
  			Data: &db.ChatMessage{
  				ID:        "generatedID",
  				ChatID:    "validChatId",
  				Message:   "This is a response",
  				Role:      "assistant",
  				Timestamp: time.Now(),
  				Status:    "sent",
  				Source:    "agent",
  			},
  		},
  	},
  	{
  		name: "All Fields Empty",
  		input: `{
  			"value": {
  				"chatId": "",
  				"messageId": "",
  				"response": "",
  				"sourceWebsocketId": ""
  			}
  		}`,
  		expectedStatus: http.StatusBadRequest,
  		expectedBody: ChatResponse{
  			Success: false,
  			Message: "ChatID is required for message creation",
  		},
  	},
  	{
  		name: "Missing MessageID",
  		input: `{
  			"value": {
  				"chatId": "validChatId",
  				"response": "This is a response",
  				"sourceWebsocketId": "validWebsocketId"
  			}
  		}`,
  		mockDBResponse: &db.ChatMessage{
  			ID:        "generatedID",
  			ChatID:    "validChatId",
  			Message:   "This is a response",
  			Role:      "assistant",
  			Timestamp: time.Now(),
  			Status:    "sent",
  			Source:    "agent",
  		},
  		mockDBError:    nil,
  		mockWSResponse: nil,
  		expectedStatus: http.StatusOK,
  		expectedBody: ChatResponse{
  			Success: true,
  			Message: "Response processed successfully",
  			Data: &db.ChatMessage{
  				ID:        "generatedID",
  				ChatID:    "validChatId",
  				Message:   "This is a response",
  				Role:      "assistant",
  				Timestamp: time.Now(),
  				Status:    "sent",
  				Source:    "agent",
  			},
  		},
  	},
  	{
  		name: "Invalid ChatID Format",
  		input: `{
  			"value": {
  				"chatId": "invalid!ChatId",
  				"messageId": "validMessageId",
  				"response": "This is a response",
  				"sourceWebsocketId": "validWebsocketId"
  			}
  		}`,
  		expectedStatus: http.StatusBadRequest,
  		expectedBody: ChatResponse{
  			Success: false,
  			Message: "Invalid ChatID format",
  		},
  	},
  }

  for _, tt := range tests {
  	t.Run(tt.name, func(t *testing.T) {
  		req := httptest.NewRequest(http.MethodPost, "/response", bytes.NewBufferString(tt.input))
  		w := httptest.NewRecorder()

  		if tt.mockDBResponse != nil || tt.mockDBError != nil {
  			mockDB.On("AddChatMessage", mock.Anything).Return(tt.mockDBResponse, tt.mockDBError)
  		}

  		if tt.mockWSResponse != nil {
  			mockWebsocketPool.On("SendTicketMessage", mock.Anything).Return(tt.mockWSResponse)
  		}

  		handler.ProcessChatResponse(w, req)

  		resp := w.Result()
  		defer resp.Body.Close()

  		assert.Equal(t, tt.expectedStatus, resp.StatusCode)

  		var responseBody ChatResponse
  		err := json.NewDecoder(resp.Body).Decode(&responseBody)
  		assert.NoError(t, err)
  		assert.Equal(t, tt.expectedBody.Success, responseBody.Success)
  		assert.Equal(t, tt.expectedBody.Message, responseBody.Message)

  		if tt.expectedBody.Data != nil {
  			assert.NotNil(t, responseBody.Data)
  			assert.Equal(t, tt.expectedBody.Data.ChatID, responseBody.Data.ChatID)
  			assert.Equal(t, tt.expectedBody.Data.Message, responseBody.Data.Message)
  		} else {
  			assert.Nil(t, responseBody.Data)
  		}

  		mockDB.AssertExpectations(t)
  		mockWebsocketPool.AssertExpectations(t)
  	})
  }
}

[aliraza556]

@tomsmith8 assign me?

[sophieturner0]

@tomsmith8 I can help?

[MahtabBukhari]

@tomsmith8 assign me

[MirzaHanan]

@tomsmith8 assign me>

[AhsanFarooqDev]

@tomsmith8 can I help

[saithsab877]

@tomsmith8 assign please?