优化结构

2026-08-30 01:08:10 +08:00
parent 509df03d61
commit e0c35e4e9b
23 changed files with 1513 additions and 38 deletions
@@ -0,0 +1,116 @@
+package mq
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestKafkaRoundTripIntegration(t *testing.T) {
+	rawBrokers := strings.TrimSpace(os.Getenv("KRA_KAFKA_TEST_BROKERS"))
+	if rawBrokers == "" {
+		t.Skip("KRA_KAFKA_TEST_BROKERS is not configured")
+	}
+	brokers := []string{"127.0.0.1:1"}
+	for _, broker := range strings.Split(rawBrokers, ",") {
+		if broker = strings.TrimSpace(broker); broker != "" {
+			brokers = append(brokers, broker)
+		}
+	}
+	if len(brokers) == 1 {
+		t.Fatal("KRA_KAFKA_TEST_BROKERS contains no broker addresses")
+	}
+
+	suffix := fmt.Sprintf("%d", time.Now().UnixNano())
+	topics := []string{"kra-review-qos1-" + suffix, "kra-review-qos0-" + suffix}
+	groupID := "kra-review-" + suffix
+	newClient := func(clientID string) *Kafka {
+		client, err := NewKafka(KafkaConfig{
+			Enabled:                true,
+			Brokers:                brokers,
+			ClientID:               clientID,
+			GroupID:                groupID,
+			StartOffset:            "earliest",
+			MinBytes:               1,
+			MaxBytes:               1 << 20,
+			MaxWait:                100 * time.Millisecond,
+			ConnectTimeout:         5 * time.Second,
+			ReconnectInterval:      500 * time.Millisecond,
+			AllowAutoTopicCreation: true,
+		})
+		if err != nil {
+			t.Fatalf("NewKafka(%s): %v", clientID, err)
+		}
+		return client
+	}
+
+	client := newClient("kra-review-producer-" + suffix)
+	t.Cleanup(func() { _ = client.Close() })
+	firstPayloads := [][]byte{[]byte("qos1-first"), []byte("qos0-first")}
+	qosValues := []byte{AtLeastOnce, AtMostOnce}
+	for index, topic := range topics {
+		if err := client.Publish(context.Background(), topic, firstPayloads[index], qosValues[index], false); err != nil {
+			t.Fatalf("publish first message to %s: %v", topic, err)
+		}
+	}
+
+	firstReceived := []chan Message{make(chan Message, 1), make(chan Message, 1)}
+	for index, topic := range topics {
+		index := index
+		if err := client.Subscribe(context.Background(), topic, qosValues[index], func(_ context.Context, message Message) {
+			firstReceived[index] <- message
+		}); err != nil {
+			t.Fatalf("subscribe to %s: %v", topic, err)
+		}
+	}
+	for index, received := range firstReceived {
+		message := waitKafkaMessage(t, received, 30*time.Second)
+		if string(message.Payload) != string(firstPayloads[index]) || message.QoS != qosValues[index] {
+			t.Fatalf("first message[%d] = payload %q qos %d", index, message.Payload, message.QoS)
+		}
+	}
+	// QoS 1 commits after the handler returns. Give the synchronous commit time
+	// to finish before closing the readers and rejoining the same group.
+	time.Sleep(500 * time.Millisecond)
+	if err := client.Unsubscribe(context.Background(), topics...); err != nil {
+		t.Fatal(err)
+	}
+
+	replayClient := newClient("kra-review-replay-" + suffix)
+	t.Cleanup(func() { _ = replayClient.Close() })
+	replayed := []chan Message{make(chan Message, 2), make(chan Message, 2)}
+	for index, topic := range topics {
+		index := index
+		if err := replayClient.Subscribe(context.Background(), topic, qosValues[index], func(_ context.Context, message Message) {
+			replayed[index] <- message
+		}); err != nil {
+			t.Fatalf("replay subscribe to %s: %v", topic, err)
+		}
+	}
+	secondPayloads := [][]byte{[]byte("qos1-second"), []byte("qos0-second")}
+	for index, topic := range topics {
+		if err := replayClient.Publish(context.Background(), topic, secondPayloads[index], qosValues[index], false); err != nil {
+			t.Fatalf("publish second message to %s: %v", topic, err)
+		}
+	}
+	for index, received := range replayed {
+		message := waitKafkaMessage(t, received, 30*time.Second)
+		if string(message.Payload) != string(secondPayloads[index]) {
+			t.Fatalf("message[%d] replayed uncommitted payload %q, want %q", index, message.Payload, secondPayloads[index])
+		}
+	}
+}
+
+func waitKafkaMessage(t *testing.T, messages <-chan Message, timeout time.Duration) Message {
+	t.Helper()
+	select {
+	case message := <-messages:
+		return message
+	case <-time.After(timeout):
+		t.Fatal("timed out waiting for kafka message")
+		return Message{}
+	}
+}