Merged in feature/batch-status (pull request #215)

Implement and test the batch status feature * working
2026-03-12 18:53:42 +00:00
parent 8a4029058b
commit aafe7d5b5f
39 changed files with 4252 additions and 516 deletions
@@ -8,6 +8,7 @@ import (

 	docinitrunner "queryorchestration/api/docInitRunner"
 	"queryorchestration/internal/database/repository"
+	"queryorchestration/internal/document/batch/outcome"
 	documentinit "queryorchestration/internal/document/init"
 	"queryorchestration/internal/server/runner"
 	"queryorchestration/internal/serviceconfig/objectstore"
@@ -51,6 +52,7 @@ func TestDocInitRunner(t *testing.T) {
 	runnerInstance := docinitrunner.New(&docinitrunner.Services{
 		Document: documentinit.New(cfg),
 		Queries:  cfg.GetDBQueries(),
+		Outcome:  outcome.New(cfg.GetDBQueries()),
 	})

 	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
@@ -101,3 +103,242 @@ func TestDocInitRunner(t *testing.T) {

 	test.AssertMessageBody(t, cfg, cfg.GetDocumentSyncURL(), regexp.MustCompile(`{"id":".+"}`))
 }
+
+// setupDocInitTest creates shared infrastructure (DB, S3, SQS) and returns
+// the config, runner, and a cleanup-style helper.
+func setupDocInitTest(t *testing.T) *DocInitConfig {
+	t.Helper()
+	cfg := &DocInitConfig{}
+	test.CreateDB(t, cfg)
+	acfg := test.CreateAWSContainer(t, cfg)
+
+	test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
+	cfg.DocumentSyncURL = test.CreateQueue(t, cfg, test.DocSyncRunnerName)
+
+	test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
+	test.CreateBucket(t, cfg)
+
+	return cfg
+}
+
+// uploadTestFileToS3 uploads a test file and returns the BucketKey.
+// If batchID is non-nil, it is embedded in the key so ParseBucketKey can extract it.
+func uploadTestFileToS3(t *testing.T, cfg *DocInitConfig, clientID string, batchID *uuid.UUID) objectstore.BucketKey {
+	t.Helper()
+	location := objectstore.BucketKey{
+		Location:  objectstore.Import,
+		ClientID:  clientID,
+		CreatedAt: time.Now().UTC(),
+		EntityID:  uuid.New(),
+		BatchID:   batchID,
+	}
+
+	testContent := []byte("test doc content " + uuid.New().String())
+	_, err := cfg.GetStoreClient().PutObject(t.Context(), &s3.PutObjectInput{
+		Bucket: aws.String(cfg.GetBucket()),
+		Key:    aws.String(location.String()),
+		Body:   bytes.NewReader(testContent),
+	})
+	require.NoError(t, err)
+
+	return location
+}
+
+// TestDocInitRunner_ReportsInitComplete processes a batch document and verifies
+// the outcome row is updated to init_complete with the assigned document_id.
+func TestDocInitRunner_ReportsInitComplete(t *testing.T) {
+	cfg := setupDocInitTest(t)
+	ctx := t.Context()
+
+	clientID := "batch_init_complete_" + uuid.New().String()[:8]
+	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Batch Init Complete",
+	})
+	require.NoError(t, err)
+
+	// Create batch
+	batchID, err := cfg.GetDBQueries().CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
+		ClientID:         clientID,
+		OriginalFilename: "test.zip",
+		TotalDocuments:   1,
+	})
+	require.NoError(t, err)
+
+	filename := "invoice.pdf"
+	// Insert submitted outcome
+	err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
+		BatchID:  batchID,
+		Filename: filename,
+		Column3:  repository.BatchOutcomeStatusSubmitted,
+	})
+	require.NoError(t, err)
+
+	// Upload test file and create upload record with batch ID embedded in key
+	location := uploadTestFileToS3(t, cfg, clientID, &batchID)
+
+	uploadID := uuid.New()
+	err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{
+		ID:       uploadID,
+		Clientid: clientID,
+		Bucket:   cfg.GetBucket(),
+		Key:      location.String(),
+		Part:     1,
+		Createdat: pgtype.Timestamp{
+			Valid: true,
+			Time:  time.Now().UTC(),
+		},
+		Filename: &filename,
+		BatchID:  &batchID,
+	})
+	require.NoError(t, err)
+
+	runnerInstance := docinitrunner.New(&docinitrunner.Services{
+		Document: documentinit.New(cfg),
+		Queries:  cfg.GetDBQueries(),
+		Outcome:  outcome.New(cfg.GetDBQueries()),
+	})
+
+	doc := docinitrunner.Body{
+		Bucket: cfg.GetBucket(),
+		Key:    location.String(),
+		Hash:   "unique_hash_init_complete_" + uuid.New().String()[:8],
+	}
+
+	assert.True(t, runnerInstance.Process(ctx, doc))
+
+	// Verify outcome was updated to init_complete with document_id set
+	rows, err := cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID)
+	require.NoError(t, err)
+	require.Len(t, rows, 1)
+	assert.Equal(t, repository.BatchOutcomeStatusInitComplete, rows[0].Outcome)
+	assert.NotNil(t, rows[0].DocumentID, "document_id should be set after init")
+}
+
+// TestDocInitRunner_ReportsInitDuplicate processes a batch document whose hash
+// matches an existing document and verifies the outcome is init_duplicate.
+func TestDocInitRunner_ReportsInitDuplicate(t *testing.T) {
+	cfg := setupDocInitTest(t)
+	ctx := t.Context()
+
+	clientID := "batch_init_dup_" + uuid.New().String()[:8]
+	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "Batch Init Duplicate",
+	})
+	require.NoError(t, err)
+
+	// Create an existing document with a known hash
+	existingHash := "existing_hash_" + uuid.New().String()[:8]
+	_, err = cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
+		Clientid: clientID,
+		Hash:     existingHash,
+	})
+	require.NoError(t, err)
+
+	// Create batch
+	batchID, err := cfg.GetDBQueries().CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
+		ClientID:         clientID,
+		OriginalFilename: "test.zip",
+		TotalDocuments:   1,
+	})
+	require.NoError(t, err)
+
+	filename := "duplicate.pdf"
+	// Insert submitted outcome
+	err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
+		BatchID:  batchID,
+		Filename: filename,
+		Column3:  repository.BatchOutcomeStatusSubmitted,
+	})
+	require.NoError(t, err)
+
+	// Upload test file and create upload record with batch ID in key
+	location := uploadTestFileToS3(t, cfg, clientID, &batchID)
+
+	uploadID := uuid.New()
+	err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{
+		ID:       uploadID,
+		Clientid: clientID,
+		Bucket:   cfg.GetBucket(),
+		Key:      location.String(),
+		Part:     1,
+		Createdat: pgtype.Timestamp{
+			Valid: true,
+			Time:  time.Now().UTC(),
+		},
+		Filename: &filename,
+		BatchID:  &batchID,
+	})
+	require.NoError(t, err)
+
+	runnerInstance := docinitrunner.New(&docinitrunner.Services{
+		Document: documentinit.New(cfg),
+		Queries:  cfg.GetDBQueries(),
+		Outcome:  outcome.New(cfg.GetDBQueries()),
+	})
+
+	// Use the same hash as the existing document to trigger duplicate
+	doc := docinitrunner.Body{
+		Bucket: cfg.GetBucket(),
+		Key:    location.String(),
+		Hash:   existingHash,
+	}
+
+	assert.True(t, runnerInstance.Process(ctx, doc))
+
+	// Verify outcome was updated to init_duplicate
+	rows, err := cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID)
+	require.NoError(t, err)
+	require.Len(t, rows, 1)
+	assert.Equal(t, repository.BatchOutcomeStatusInitDuplicate, rows[0].Outcome)
+	assert.NotNil(t, rows[0].DocumentID, "document_id should be set for duplicate")
+}
+
+// TestDocInitRunner_NoOpForNonBatchDocument verifies that processing a non-batch
+// document does not cause outcome errors (the Resolve call is a no-op for nil batchID).
+func TestDocInitRunner_NoOpForNonBatchDocument(t *testing.T) {
+	cfg := setupDocInitTest(t)
+	ctx := t.Context()
+
+	clientID := "nonbatch_init_" + uuid.New().String()[:8]
+	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
+		Clientid: clientID,
+		Name:     "NonBatch Init",
+	})
+	require.NoError(t, err)
+
+	location := uploadTestFileToS3(t, cfg, clientID, nil)
+
+	uploadID := uuid.New()
+	fname := "regular.pdf"
+	err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{
+		ID:       uploadID,
+		Clientid: clientID,
+		Bucket:   cfg.GetBucket(),
+		Key:      location.String(),
+		Part:     1,
+		Createdat: pgtype.Timestamp{
+			Valid: true,
+			Time:  time.Now().UTC(),
+		},
+		Filename: &fname,
+		BatchID:  nil, // Non-batch
+	})
+	require.NoError(t, err)
+
+	runnerInstance := docinitrunner.New(&docinitrunner.Services{
+		Document: documentinit.New(cfg),
+		Queries:  cfg.GetDBQueries(),
+		Outcome:  outcome.New(cfg.GetDBQueries()),
+	})
+
+	doc := docinitrunner.Body{
+		Bucket: cfg.GetBucket(),
+		Key:    location.String(),
+		Hash:   "nonbatch_hash_" + uuid.New().String()[:8],
+	}
+
+	// Should process successfully without any outcome errors
+	assert.True(t, runnerInstance.Process(ctx, doc))
+}