实现 durable assignment 消费循环

internal/assignmentqueue/jetstream_test.go

@@ -53,11 +53,16 @@ func TestPublisherUsesBackendSubjectAndAssignmentDeduplication(t *testing.T) {
 }
 
 type fakeHandler struct {
-	done bool
-	err  error
+	done      bool
+	err       error
+	remaining int
 }
 
-func (h fakeHandler) Handle(context.Context, taskassignment.Assignment) (bool, error) {
+func (h *fakeHandler) Handle(context.Context, taskassignment.Assignment) (bool, error) {
+	if h.remaining > 0 {
+		h.remaining--
+		return false, nil
+	}
 	return h.done, h.err
 }
 
@@ -86,7 +91,7 @@ func encodedAssignment(t *testing.T) []byte {
 
 func TestProcessorAcknowledgesOnlyCompletedAssignment(t *testing.T) {
 	message := &fakeMessage{data: encodedAssignment(t)}
-	processor := Processor{TrustDomain: "ddupan.top", Handler: fakeHandler{done: true}}
+	processor := Processor{TrustDomain: "ddupan.top", Handler: &fakeHandler{done: true}}
 	if err := processor.Process(context.Background(), message); err != nil {
 		t.Fatal(err)
 	}
@@ -97,7 +102,7 @@ func TestProcessorAcknowledgesOnlyCompletedAssignment(t *testing.T) {
 
 func TestProcessorKeepsRunningAssignmentPending(t *testing.T) {
 	message := &fakeMessage{data: encodedAssignment(t)}
-	processor := Processor{TrustDomain: "ddupan.top", Handler: fakeHandler{}}
+	processor := Processor{TrustDomain: "ddupan.top", Handler: &fakeHandler{}}
 	if err := processor.Process(context.Background(), message); err != nil {
 		t.Fatal(err)
 	}
@@ -110,7 +115,7 @@ func TestProcessorRetriesBackendFailureAndTerminatesPoisonMessage(t *testing.T)
 	retry := &fakeMessage{data: encodedAssignment(t)}
 	processor := Processor{
 		TrustDomain: "ddupan.top",
-		Handler:     fakeHandler{err: errors.New("backend unavailable")},
+		Handler:     &fakeHandler{err: errors.New("backend unavailable")},
 		RetryDelay:  time.Minute,
 	}
 	if err := processor.Process(context.Background(), retry); err == nil {
@@ -128,3 +133,52 @@ func TestProcessorRetriesBackendFailureAndTerminatesPoisonMessage(t *testing.T)
 		t.Fatalf("poison message = %#v", poison)
 	}
 }
+
+func TestProcessUntilDoneReconcilesWithoutRedelivery(t *testing.T) {
+	message := &fakeMessage{data: encodedAssignment(t)}
+	processor := Processor{
+		TrustDomain:  "ddupan.top",
+		Handler:      &fakeHandler{done: true, remaining: 2},
+		PollInterval: time.Millisecond,
+	}
+	if err := processor.ProcessUntilDone(context.Background(), message); err != nil {
+		t.Fatal(err)
+	}
+	if message.inProgress != 2 || message.acked != 1 || message.nacked != 0 {
+		t.Fatalf("message = %#v", message)
+	}
+}
+
+func TestProcessUntilDoneLeavesMessagePendingOnShutdown(t *testing.T) {
+	message := &fakeMessage{data: encodedAssignment(t)}
+	processor := Processor{
+		TrustDomain:  "ddupan.top",
+		Handler:      &fakeHandler{},
+		PollInterval: time.Hour,
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if err := processor.ProcessUntilDone(ctx, message); !errors.Is(err, context.Canceled) {
+		t.Fatalf("error = %v", err)
+	}
+	if message.acked != 0 || message.nacked != 0 || message.terminated != 0 {
+		t.Fatalf("message = %#v", message)
+	}
+}
+
+type fakeConsumerManager struct{ config jetstream.ConsumerConfig }
+
+func (m *fakeConsumerManager) CreateOrUpdateConsumer(_ context.Context, _ string, config jetstream.ConsumerConfig) (jetstream.Consumer, error) {
+	m.config = config
+	return nil, nil
+}
+
+func TestOpenConsumerUsesIndependentDurablePerBackend(t *testing.T) {
+	manager := &fakeConsumerManager{}
+	if _, err := OpenConsumer(context.Background(), manager, "CI_RUNNER", "ci.assignment", taskassignment.BackendPod, 4); err != nil {
+		t.Fatal(err)
+	}
+	if manager.config.Durable != "pod" || manager.config.FilterSubject != "ci.assignment.pod" || manager.config.AckPolicy != jetstream.AckExplicitPolicy || manager.config.MaxAckPending != 4 {
+		t.Fatalf("config = %#v", manager.config)
+	}
+}