vetrag/controller_test.go

package main

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

func TestChatService_HandleChat(t *testing.T) {
	// Setup mock dependencies
	mockLLM := &MockLLMClient{
		ExtractKeywordsFunc: func(ctx context.Context, message string) (map[string]interface{}, error) {
			return map[string]interface{}{
				"translate": "test translation",
				"keyword":   []string{"test", "keyword"},
				"animal":    "dog",
			}, nil
		},
		DisambiguateBestMatchFunc: func(ctx context.Context, message string, candidates []Visit) (string, error) {
			return "visit1", nil
		},
	}

	mockDB := &MockVisitDB{
		FindCandidatesFunc: func(keywords []string) ([]Visit, error) {
			return []Visit{
				{ID: "visit1", Procedures: []Procedure{{Name: "Test", Price: 100, DurationMin: 30}}},
				{ID: "visit2", Procedures: []Procedure{{Name: "Test2", Price: 200, DurationMin: 60}}},
			}, nil
		},
		FindByIdFunc: func(id string) (Visit, error) {
			return Visit{
				ID: id,
				Procedures: []Procedure{
					{Name: "Test", Price: 100, DurationMin: 30},
				},
				Notes: "Test notes",
			}, nil
		},
	}

	mockRepo := &MockChatRepository{}

	// Create service with mocks
	svc := NewChatService(mockLLM, mockDB, mockRepo)

	// Create test context
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)

	// Mock request body
	reqBody := `{"message": "I need a test visit"}`
	c.Request = httptest.NewRequest(http.MethodPost, "/chat", strings.NewReader(reqBody))
	c.Request.Header.Set("Content-Type", "application/json")

	// Call the handler
	svc.HandleChat(c)

	// Validate response
	assert.Equal(t, http.StatusOK, w.Code)

	var resp ChatResponse
	err := json.Unmarshal(w.Body.Bytes(), &resp)
	assert.NoError(t, err)

	assert.NotNil(t, resp.Match)
	assert.Equal(t, "visit1", *resp.Match)
	assert.Equal(t, 100, resp.TotalPrice)
	assert.Equal(t, 30, resp.TotalDuration)
	assert.Equal(t, "Test notes", resp.Notes)
}

type MockVisitDB struct {
	FindCandidatesFunc func(keywords []string) ([]Visit, error)
	FindByIdFunc       func(id string) (Visit, error)
}

func (m *MockVisitDB) FindCandidates(keywords []string) ([]Visit, error) {
	if m.FindCandidatesFunc != nil {
		return m.FindCandidatesFunc(keywords)
	}
	return nil, nil
}

func (m *MockVisitDB) FindById(id string) (Visit, error) {
	if m.FindByIdFunc != nil {
		return m.FindByIdFunc(id)
	}
	return Visit{}, nil
}

type MockChatRepository struct {
	SaveChatInteractionFunc   func(ctx context.Context, interaction ChatInteraction) error
	ListChatInteractionsFunc  func(ctx context.Context, limit, offset int) ([]ChatInteraction, error)
	SaveLLMRawEventFunc       func(ctx context.Context, correlationID, phase, raw string) error
	ListLLMRawEventsFunc      func(ctx context.Context, correlationID string, limit, offset int) ([]RawLLMEvent, error)
	SaveKnowledgeModelFunc    func(ctx context.Context, text string) error
	ListKnowledgeModelsFunc   func(ctx context.Context, limit, offset int) ([]knowledgeModelMeta, error)
	GetKnowledgeModelTextFunc func(ctx context.Context, id int64) (string, error)
	GetUserByUsernameFunc     func(ctx context.Context, username string) (*User, error)
	CountUsersFunc            func(ctx context.Context) (int, error)
	CreateUserFunc            func(ctx context.Context, username, passwordHash string) error
}

func (m *MockChatRepository) SaveChatInteraction(ctx context.Context, interaction ChatInteraction) error {
	if m.SaveChatInteractionFunc != nil {
		return m.SaveChatInteractionFunc(ctx, interaction)
	}
	return nil
}

func (m *MockChatRepository) ListChatInteractions(ctx context.Context, limit, offset int) ([]ChatInteraction, error) {
	if m.ListChatInteractionsFunc != nil {
		return m.ListChatInteractionsFunc(ctx, limit, offset)
	}
	return []ChatInteraction{}, nil
}

func (m *MockChatRepository) SaveLLMRawEvent(ctx context.Context, correlationID, phase, raw string) error {
	if m.SaveLLMRawEventFunc != nil {
		return m.SaveLLMRawEventFunc(ctx, correlationID, phase, raw)
	}
	return nil
}

func (m *MockChatRepository) ListLLMRawEvents(ctx context.Context, correlationID string, limit, offset int) ([]RawLLMEvent, error) {
	if m.ListLLMRawEventsFunc != nil {
		return m.ListLLMRawEventsFunc(ctx, correlationID, limit, offset)
	}
	return []RawLLMEvent{}, nil
}

func (m *MockChatRepository) SaveKnowledgeModel(ctx context.Context, text string) error {
	if m.SaveKnowledgeModelFunc != nil {
		return m.SaveKnowledgeModelFunc(ctx, text)
	}
	return nil
}

func (m *MockChatRepository) ListKnowledgeModels(ctx context.Context, limit, offset int) ([]knowledgeModelMeta, error) {
	if m.ListKnowledgeModelsFunc != nil {
		return m.ListKnowledgeModelsFunc(ctx, limit, offset)
	}
	return []knowledgeModelMeta{}, nil
}

func (m *MockChatRepository) GetKnowledgeModelText(ctx context.Context, id int64) (string, error) {
	if m.GetKnowledgeModelTextFunc != nil {
		return m.GetKnowledgeModelTextFunc(ctx, id)
	}
	return "", nil
}

func (m *MockChatRepository) GetUserByUsername(ctx context.Context, username string) (*User, error) {
	if m.GetUserByUsernameFunc != nil {
		return m.GetUserByUsernameFunc(ctx, username)
	}
	return nil, nil
}

func (m *MockChatRepository) CountUsers(ctx context.Context) (int, error) {
	if m.CountUsersFunc != nil {
		return m.CountUsersFunc(ctx)
	}
	return 0, nil
}

func (m *MockChatRepository) CreateUser(ctx context.Context, username, passwordHash string) error {
	if m.CreateUserFunc != nil {
		return m.CreateUserFunc(ctx, username, passwordHash)
	}
	return nil
}