query-orchestration/internal/document/batch/service_test.go

package batch_test

import (
	"testing"

	"queryorchestration/internal/database/repository"
	documentbatch "queryorchestration/internal/document/batch"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/test"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type TestConfig struct {
	serviceconfig.BaseConfig
}

func TestUpdateProgress(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_progress"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Progress",
	})
	require.NoError(t, err)

	// Create test batch with total documents
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test updating progress
	err = service.UpdateProgress(t.Context(), clientID, batchID, 5, 2, 1)
	require.NoError(t, err)

	// Verify progress was updated
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, int32(5), batch.ProcessedDocuments)
	assert.Equal(t, int32(2), batch.FailedDocuments)
	assert.Equal(t, int32(1), batch.InvalidTypeDocuments)
	assert.Equal(t, int32(80), batch.ProgressPercent) // (5+2+1)/10 * 100 = 80%
}

func TestUpdateProgressZeroTotalDocuments(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_zero"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Zero",
	})
	require.NoError(t, err)

	// Create test batch with zero total documents
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   0,
	})
	require.NoError(t, err)

	// Test updating progress should fail with zero total documents
	err = service.UpdateProgress(t.Context(), clientID, batchID, 5, 2, 1)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "batch has zero total documents")
}

func TestMarkCompleted(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_complete"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Complete",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test marking as completed
	err = service.MarkCompleted(t.Context(), batchID)
	require.NoError(t, err)

	// Verify status was updated
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, repository.BatchStatusCompleted, batch.Status)
}

func TestMarkFailed(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_failed"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Failed",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test marking as failed
	err = service.MarkFailed(t.Context(), batchID)
	require.NoError(t, err)

	// Verify status was updated
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, repository.BatchStatusFailed, batch.Status)
}

func TestAddFailedFilename(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_filename"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Filename",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test adding failed filename
	err = service.AddFailedFilename(t.Context(), batchID, "failed_file.pdf")
	require.NoError(t, err)

	// Get the full details to check failed filenames
	details, err := service.Get(t.Context(), clientID, batchID)
	require.NoError(t, err)
	assert.Contains(t, details.FailedFilenames, "failed_file.pdf")
}