chore: 并入 iotlcm-worker-svr 历史到 worker-svr/（原提交 96f0f4e）

worker-svr/internal/library/rabbitmq/rabbitmq_test.go

@@ -0,0 +1,170 @@
+package rabbitmq
+
+import (
+	"context"
+	"reflect"
+	"sync"
+	"testing"
+)
+
+type fakeChannel struct {
+	mu          sync.Mutex
+	cancelCalls int
+	closeCalls  int
+	closed      bool
+}
+
+func (f *fakeChannel) Cancel(string, bool) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.cancelCalls++
+	return nil
+}
+
+func (f *fakeChannel) Close() error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.closeCalls++
+	f.closed = true
+	return nil
+}
+
+func (f *fakeChannel) IsClosed() bool {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.closed
+}
+
+func TestCloseIsIdempotent(t *testing.T) {
+	resetRabbitMQState(t)
+
+	channel := &fakeChannel{}
+	done := make(chan struct{})
+	close(done)
+	registerConsumer(&managedConsumer{
+		channel: channel,
+		tag:     "test-consumer",
+		done:    done,
+	})
+
+	if err := Close(context.Background()); err != nil {
+		t.Fatalf("first Close() error = %v", err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if err := Close(ctx); err != nil {
+		t.Fatalf("second Close() error = %v", err)
+	}
+
+	channel.mu.Lock()
+	defer channel.mu.Unlock()
+	if channel.cancelCalls != 1 {
+		t.Fatalf("Cancel() calls = %d, want 1", channel.cancelCalls)
+	}
+	if channel.closeCalls != 1 {
+		t.Fatalf("Close() calls = %d, want 1", channel.closeCalls)
+	}
+}
+
+func TestCloseSkipsAlreadyClosedChannel(t *testing.T) {
+	resetRabbitMQState(t)
+
+	channel := &fakeChannel{closed: true}
+	done := make(chan struct{})
+	close(done)
+	registerConsumer(&managedConsumer{
+		channel: channel,
+		tag:     "closed-consumer",
+		done:    done,
+	})
+
+	if err := Close(context.Background()); err != nil {
+		t.Fatalf("Close() error = %v", err)
+	}
+
+	channel.mu.Lock()
+	defer channel.mu.Unlock()
+	if channel.cancelCalls != 0 || channel.closeCalls != 0 {
+		t.Fatalf("closed channel calls: Cancel=%d Close=%d, want 0", channel.cancelCalls, channel.closeCalls)
+	}
+}
+
+func TestCloseHonorsContextWhileWaiting(t *testing.T) {
+	resetRabbitMQState(t)
+
+	channel := &fakeChannel{}
+	registerConsumer(&managedConsumer{
+		channel: channel,
+		tag:     "blocked-consumer",
+		done:    make(chan struct{}),
+	})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if err := Close(ctx); err == nil {
+		t.Fatal("Close() error = nil, want context cancellation error")
+	}
+
+	channel.mu.Lock()
+	defer channel.mu.Unlock()
+	if channel.closeCalls != 1 {
+		t.Fatalf("Close() calls = %d, want forced cleanup", channel.closeCalls)
+	}
+}
+
+func TestDeliveryMetadataContext(t *testing.T) {
+	expected := DeliveryMetadata{
+		DeliveryTag: 42,
+		Redelivered: true,
+		Consumer:    "batch-consumer",
+		Queue:       "iotlcm.batch.execute",
+	}
+	ctx := WithDeliveryMetadata(context.Background(), expected)
+	actual, ok := DeliveryMetadataFromContext(ctx)
+	if !ok {
+		t.Fatal("delivery metadata not found in context")
+	}
+	if !reflect.DeepEqual(actual, expected) {
+		t.Fatalf("delivery metadata = %+v, want %+v", actual, expected)
+	}
+	if _, ok = DeliveryMetadataFromContext(context.Background()); ok {
+		t.Fatal("unexpected delivery metadata in empty context")
+	}
+}
+
+func resetRabbitMQState(t *testing.T) {
+	t.Helper()
+	consumerMu.Lock()
+	consumers = nil
+	consumerMu.Unlock()
+	connMu.Lock()
+	conn = nil
+	connMu.Unlock()
+	t.Cleanup(func() {
+		consumerMu.Lock()
+		consumers = nil
+		consumerMu.Unlock()
+		connMu.Lock()
+		conn = nil
+		connMu.Unlock()
+	})
+}
+
+func TestUnregisterConsumerPreservesShutdownSnapshot(t *testing.T) {
+	consumerMu.Lock()
+	original := consumers
+	first, second, third := &managedConsumer{}, &managedConsumer{}, &managedConsumer{}
+	consumers = []*managedConsumer{first, second, third}
+	snapshot := consumers
+	consumerMu.Unlock()
+	defer func() { consumerMu.Lock(); consumers = original; consumerMu.Unlock() }()
+	unregisterConsumer(second)
+	consumerMu.Lock()
+	defer consumerMu.Unlock()
+	if len(consumers) != 2 || consumers[0] != first || consumers[1] != third {
+		t.Fatal("consumer not removed")
+	}
+	if snapshot[1] != second || snapshot[2] != third {
+		t.Fatal("shutdown snapshot mutated")
+	}
+}