✨ Introduce seed-telegrams tool for generating synthetic telegrams for testing and development. Enhance README with detailed usage instructions, command-line parameters, and examples. Add unit tests for telegram generation and status selection logic to ensure robustness. Update documentation to reflect new features and usage scenarios.

2025-11-16 17:24:16 +08:00
parent 06b2495d77
commit 12fda00df9
9 changed files with 730 additions and 80 deletions
@@ -23,6 +23,26 @@ var (
 	bodyCategories     = []string{"ARR", "DEP", "CNL", "DLA", "FPL"}
 )

+// SeedConfig controls how telegrams are generated and dispatched.
+type SeedConfig struct {
+	Count        int
+	Mode         string
+	IntervalMin  time.Duration
+	IntervalMax  time.Duration
+	Duration     time.Duration
+	CategoryFlag string
+	StatusFlag   string
+	ErrorReason  string
+	HeaderFormat string
+	DryRun       bool
+}
+
+// publishFunc abstracts the publish side-effect so that it can be swapped in tests.
+type publishFunc func(*telegram) error
+
+// sleepFunc is used instead of time.Sleep so tests can stub out real sleeping.
+var sleepFunc = time.Sleep
+
 func main() {
 	natsURL := flag.String("nats-url", "nats://127.0.0.1:4222", "NATS server URL (empty skips publish)")
 	subject := flag.String("subject", "telegram.raw", "Subject to publish telegrams to")
@@ -36,6 +56,10 @@ func main() {
 	jsStream := flag.String("stream", "", "JetStream stream (optional when --jetstream)")
 	jsSubject := flag.String("js-subject", "", "Override subject for JetStream publish (defaults to --subject)")
 	headerFormat := flag.String("header-format", "json", "Metadata header encoding: json|none")
+	mode := flag.String("mode", "burst", "Seed mode: burst|interval|mixed")
+	intervalMin := flag.Duration("interval-min", time.Second, "Minimum interval between messages in interval/mixed modes")
+	intervalMax := flag.Duration("interval-max", 2*time.Second, "Maximum interval between messages in interval/mixed modes")
+	duration := flag.Duration("duration", 0, "Total duration for interval/mixed modes (0 = rely on --count only)")
 	flag.Parse()

 	rand.Seed(time.Now().UnixNano())
@@ -64,62 +88,34 @@ func main() {
 		}
 	}

-	categories := bodyCategories
-	if strings.ToLower(*category) != "mixed" {
-		categories = []string{strings.ToUpper(*category)}
+	// Normalize and prepare configuration.
+	if *intervalMax < *intervalMin {
+		*intervalMax = *intervalMin
 	}

-	statuses := statusValues
-	if strings.ToLower(*status) != "random" {
-		statuses = []string{strings.ToLower(*status)}
+	cfg := SeedConfig{
+		Count:        *count,
+		Mode:         strings.ToLower(*mode),
+		IntervalMin:  *intervalMin,
+		IntervalMax:  *intervalMax,
+		Duration:     *duration,
+		CategoryFlag: *category,
+		StatusFlag:   *status,
+		ErrorReason:  *errorReason,
+		HeaderFormat: *headerFormat,
+		DryRun:       *dryRun,
 	}

-	for i := 0; i < *count; i++ {
-		cat := categories[rand.Intn(len(categories))]
-		payload, intentionallyInvalid := buildTelegram(cat)
+	categories := buildCategories(cfg.CategoryFlag, bodyCategories)
+	statuses := buildStatuses(cfg.StatusFlag, statusValues)

-		// 当状态为 random 时，根据报文是否合法来倾向选择 parsed 或 body_error
-		if strings.ToLower(*status) == "random" {
-			if intentionallyInvalid {
-				// 故意非法的报文：大概率标记为 body_error
-				if rand.Intn(100) < 80 {
-					payload.Status = "body_error"
-				} else {
-					payload.Status = statusValues[rand.Intn(len(statusValues))]
-				}
-			} else {
-				// 合法报文：大概率标记为 parsed
-				if rand.Intn(100) < 70 {
-					payload.Status = "parsed"
-				} else {
-					payload.Status = statusValues[rand.Intn(len(statusValues))]
-				}
-			}
-		} else {
-			// 非 random 模式下沿用原有逻辑
-			payload.Status = statuses[rand.Intn(len(statuses))]
-		}
-
-		payload.ErrorReason = *errorReason
-		payload.Metadata = map[string]string{
-			"message_id": payload.MessageID,
-			"category":   payload.Category,
-			"comments":   fmt.Sprintf("seeded iteration=%d", i),
-			"status":     payload.Status,
-		}
-
-		if *dryRun {
-			blob, _ := json.MarshalIndent(payload, "", "  ")
-			fmt.Println(string(blob))
-			fmt.Println("---")
-			continue
-		}
-
-		if nc != nil && !*noNATS {
+	var publisher publishFunc
+	if nc != nil && !*noNATS {
+		publisher = func(payload *telegram) error {
 			data := []byte(payload.Content)
 			msg := &nats.Msg{Subject: *subject, Data: data, Header: nats.Header{}}
 			msg.Header.Set("Nats-Msg-Id", payload.UUID)
-			if strings.ToLower(*headerFormat) == "json" {
+			if strings.ToLower(cfg.HeaderFormat) == "json" {
 				headerJSON, _ := json.Marshal(payload.Metadata)
 				msg.Header.Set("x-telegram-meta", string(headerJSON))
 			}
@@ -134,16 +130,22 @@ func main() {
 				}
 				msg.Subject = pubSubject
 				if _, err := js.PublishMsg(msg); err != nil {
-					log.Fatalf("jetstream publish: %v", err)
-				}
-			} else {
-				if err := nc.PublishMsg(msg); err != nil {
-					log.Fatalf("nats publish: %v", err)
+					return fmt.Errorf("jetstream publish: %w", err)
 				}
+				return nil
 			}
+
+			if err := nc.PublishMsg(msg); err != nil {
+				return fmt.Errorf("nats publish: %w", err)
+			}
+			return nil
 		}
 	}

+	if err := RunSeed(cfg, categories, statuses, publisher); err != nil {
+		log.Fatalf("run seed: %v", err)
+	}
+
 	if !*dryRun {
 		if nc != nil && !*noNATS {
 			log.Printf("Published %d telegram(s) to %s", *count, *subject)
@@ -151,6 +153,160 @@ func main() {
 	}
 }

+// RunSeed drives telegram generation according to the configured mode.
+// It delegates message construction to buildTelegram and side-effects to the provided publisher.
+func RunSeed(cfg SeedConfig, categories []string, statuses []string, publisher publishFunc) error {
+	if len(categories) == 0 {
+		return fmt.Errorf("no categories available")
+	}
+	if len(statuses) == 0 {
+		return fmt.Errorf("no statuses available")
+	}
+
+	switch cfg.Mode {
+	case "interval":
+		return runIntervalMode(cfg, categories, statuses, publisher)
+	case "mixed":
+		return runMixedMode(cfg, categories, statuses, publisher)
+	default:
+		// Default to burst semantics.
+		return runBurstMode(cfg, categories, statuses, publisher)
+	}
+}
+
+func runBurstMode(cfg SeedConfig, categories []string, statuses []string, publisher publishFunc) error {
+	count := cfg.Count
+	if count <= 0 {
+		return nil
+	}
+	for i := 0; i < count; i++ {
+		if err := sendTelegram(i, cfg, categories, statuses, publisher); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func runIntervalMode(cfg SeedConfig, categories []string, statuses []string, publisher publishFunc) error {
+	start := time.Now()
+	sent := 0
+
+	for {
+		if cfg.Count > 0 && sent >= cfg.Count {
+			break
+		}
+		if cfg.Duration > 0 && time.Since(start) >= cfg.Duration {
+			break
+		}
+
+		if err := sendTelegram(sent, cfg, categories, statuses, publisher); err != nil {
+			return err
+		}
+		sent++
+
+		// Compute next sleep duration within [IntervalMin, IntervalMax].
+		sleepDur := cfg.IntervalMin
+		if cfg.IntervalMax > cfg.IntervalMin {
+			delta := cfg.IntervalMax - cfg.IntervalMin
+			sleepDur = cfg.IntervalMin + time.Duration(rand.Int63n(int64(delta)+1))
+		}
+		if sleepDur > 0 && !cfg.DryRun {
+			sleepFunc(sleepDur)
+		}
+	}
+
+	return nil
+}
+
+// mixed 模式：前半段使用 interval 模式，后半段使用 burst。
+func runMixedMode(cfg SeedConfig, categories []string, statuses []string, publisher publishFunc) error {
+	// 简单策略：如果 Count>0，前半部分 interval，后半部分 burst；否则退化为 interval。
+	if cfg.Count <= 0 {
+		return runIntervalMode(cfg, categories, statuses, publisher)
+	}
+
+	half := cfg.Count / 2
+	if half == 0 {
+		// Count==1 时直接按 burst 处理。
+		return runBurstMode(cfg, categories, statuses, publisher)
+	}
+
+	intervalCfg := cfg
+	intervalCfg.Count = half
+	if err := runIntervalMode(intervalCfg, categories, statuses, publisher); err != nil {
+		return err
+	}
+
+	burstCfg := cfg
+	burstCfg.Count = cfg.Count - half
+	return runBurstMode(burstCfg, categories, statuses, publisher)
+}
+
+// sendTelegram builds a single telegram, assigns status/metadata, and either prints or publishes it.
+func sendTelegram(iteration int, cfg SeedConfig, categories []string, statuses []string, publisher publishFunc) error {
+	cat := categories[rand.Intn(len(categories))]
+	payload, intentionallyInvalid := buildTelegram(cat)
+
+	payload.Status = chooseStatus(intentionallyInvalid, cfg.StatusFlag, statuses)
+	payload.ErrorReason = cfg.ErrorReason
+	payload.Metadata = map[string]string{
+		"message_id": payload.MessageID,
+		"category":   payload.Category,
+		"comments":   fmt.Sprintf("seeded iteration=%d", iteration),
+		"status":     payload.Status,
+	}
+
+	if cfg.DryRun {
+		blob, _ := json.MarshalIndent(payload, "", "  ")
+		fmt.Println(string(blob))
+		fmt.Println("---")
+		return nil
+	}
+
+	if publisher != nil {
+		return publisher(payload)
+	}
+	return nil
+}
+
+// buildCategories returns the effective categories list based on the CLI flag.
+func buildCategories(flagValue string, all []string) []string {
+	if strings.ToLower(flagValue) == "mixed" || flagValue == "" {
+		return all
+	}
+	return []string{strings.ToUpper(flagValue)}
+}
+
+// buildStatuses returns the effective statuses list based on the CLI flag.
+func buildStatuses(flagValue string, all []string) []string {
+	if strings.ToLower(flagValue) == "random" || flagValue == "" {
+		return all
+	}
+	return []string{strings.ToLower(flagValue)}
+}
+
+// chooseStatus encapsulates the status selection logic, including the special \"random\" behaviour.
+func chooseStatus(intentionallyInvalid bool, statusFlag string, statuses []string) string {
+	if strings.ToLower(statusFlag) != "random" {
+		return statuses[rand.Intn(len(statuses))]
+	}
+
+	// random 模式：根据报文是否合法，对 parsed/body_error 做倾向性选择。
+	if intentionallyInvalid {
+		// 非法报文：大概率 body_error。
+		if rand.Intn(100) < 80 {
+			return "body_error"
+		}
+	} else {
+		// 合法报文：大概率 parsed。
+		if rand.Intn(100) < 70 {
+			return "parsed"
+		}
+	}
+
+	return statuses[rand.Intn(len(statuses))]
+}
+
 type telegram struct {
 	UUID        string            `json:"uuid"`
 	MessageID   string            `json:"message_id"`
@@ -0,0 +1,211 @@
+package main
+
+import (
+	"testing"
+	"time"
+)
+
+// Test buildBody generates syntactically plausible bodies for each category.
+func TestBuildBodyFormats(t *testing.T) {
+	tests := []struct {
+		category string
+		prefix   string
+		suffix   string
+	}{
+		{"ARR", "(ARR-", ")"},
+		{"DEP", "(DEP-", ")"},
+		{"CNL", "(CNL-", ")"},
+		{"DLA", "(DLA-", ")"},
+		{"FPL", "(FPL-", ")"},
+	}
+
+	for _, tt := range tests {
+		body, invalid := buildBody(tt.category)
+		if !invalid && body == "" {
+			t.Fatalf("category %s: expected non-empty body", tt.category)
+		}
+		if body[0:len(tt.prefix)] != tt.prefix {
+			t.Fatalf("category %s: expected prefix %q, got %q", tt.category, tt.prefix, body[0:len(tt.prefix)])
+		}
+		if body[len(body)-len(tt.suffix):] != tt.suffix {
+			t.Fatalf("category %s: expected suffix %q, got %q", tt.category, tt.suffix, body[len(body)-len(tt.suffix):])
+		}
+	}
+}
+
+// Test buildTelegram wires header and body into a full telegram.
+func TestBuildTelegramBasicFields(t *testing.T) {
+	tg, _ := buildTelegram("ARR")
+	if tg.Category != "ARR" {
+		t.Fatalf("expected category ARR, got %s", tg.Category)
+	}
+	if tg.MessageID == "" {
+		t.Fatalf("expected non-empty MessageID")
+	}
+	if tg.UUID == "" {
+		t.Fatalf("expected non-empty UUID")
+	}
+	if tg.Content == "" {
+		t.Fatalf("expected non-empty Content")
+	}
+	if tg.ReceivedAt.IsZero() {
+		t.Fatalf("expected ReceivedAt to be set")
+	}
+	if tg.Metadata != nil {
+		t.Fatalf("expected Metadata to be nil from buildTelegram")
+	}
+}
+
+func TestBuildCategories(t *testing.T) {
+	all := []string{"ARR", "DEP", "CNL", "DLA", "FPL"}
+
+	got := buildCategories("mixed", all)
+	if len(got) != len(all) {
+		t.Fatalf("expected all categories for mixed, got %d", len(got))
+	}
+
+	got = buildCategories("", all)
+	if len(got) != len(all) {
+		t.Fatalf("expected all categories for empty flag, got %d", len(got))
+	}
+
+	got = buildCategories("arr", all)
+	if len(got) != 1 || got[0] != "ARR" {
+		t.Fatalf("expected single category ARR, got %#v", got)
+	}
+}
+
+func TestBuildStatuses(t *testing.T) {
+	all := []string{"parsed", "header_error", "body_error"}
+
+	got := buildStatuses("random", all)
+	if len(got) != len(all) {
+		t.Fatalf("expected all statuses for random, got %d", len(got))
+	}
+
+	got = buildStatuses("", all)
+	if len(got) != len(all) {
+		t.Fatalf("expected all statuses for empty flag, got %d", len(got))
+	}
+
+	got = buildStatuses("parsed", all)
+	if len(got) != 1 || got[0] != "parsed" {
+		t.Fatalf("expected single status parsed, got %#v", got)
+	}
+}
+
+func TestChooseStatusNonRandomUsesProvidedList(t *testing.T) {
+	statuses := []string{"parsed"}
+	for i := 0; i < 10; i++ {
+		got := chooseStatus(false, "parsed", statuses)
+		if got != "parsed" {
+			t.Fatalf("expected parsed, got %s", got)
+		}
+	}
+}
+
+func TestChooseStatusRandomOnlyReturnsKnownStatuses(t *testing.T) {
+	statuses := []string{"parsed", "header_error", "body_error"}
+
+	for i := 0; i < 100; i++ {
+		got := chooseStatus(false, "random", statuses)
+		if !contains(statuses, got) && got != "body_error" && got != "parsed" {
+			t.Fatalf("unexpected status from chooseStatus: %s", got)
+		}
+	}
+}
+
+func contains(list []string, v string) bool {
+	for _, s := range list {
+		if s == v {
+			return true
+		}
+	}
+	return false
+}
+
+// Test RunSeed burst mode publishes Count messages.
+func TestRunSeedBurstPublishesCount(t *testing.T) {
+	cfg := SeedConfig{
+		Count:  5,
+		Mode:   "burst",
+		DryRun: false,
+	}
+	categories := []string{"ARR"}
+	statuses := []string{"parsed"}
+
+	var published int
+	pub := func(*telegram) error {
+		published++
+		return nil
+	}
+
+	if err := RunSeed(cfg, categories, statuses, pub); err != nil {
+		t.Fatalf("RunSeed burst returned error: %v", err)
+	}
+	if published != cfg.Count {
+		t.Fatalf("expected %d published messages, got %d", cfg.Count, published)
+	}
+}
+
+// Test RunSeed interval mode respects Count and does not sleep in DryRun.
+func TestRunSeedIntervalRespectsCount(t *testing.T) {
+	origSleep := sleepFunc
+	defer func() { sleepFunc = origSleep }()
+	sleepFunc = func(time.Duration) {}
+
+	cfg := SeedConfig{
+		Count:       3,
+		Mode:        "interval",
+		IntervalMin: 10 * time.Millisecond,
+		IntervalMax: 20 * time.Millisecond,
+		DryRun:      true,
+	}
+	categories := []string{"ARR"}
+	statuses := []string{"parsed"}
+
+	var published int
+	pub := func(*telegram) error {
+		published++
+		return nil
+	}
+
+	if err := RunSeed(cfg, categories, statuses, pub); err != nil {
+		t.Fatalf("RunSeed interval returned error: %v", err)
+	}
+	if published != cfg.Count {
+		t.Fatalf("expected %d published messages, got %d", cfg.Count, published)
+	}
+}
+
+// Test RunSeed mixed mode still results in exactly Count messages when Count>0.
+func TestRunSeedMixedPublishesCount(t *testing.T) {
+	origSleep := sleepFunc
+	defer func() { sleepFunc = origSleep }()
+	sleepFunc = func(time.Duration) {}
+
+	cfg := SeedConfig{
+		Count:       6,
+		Mode:        "mixed",
+		IntervalMin: 1 * time.Millisecond,
+		IntervalMax: 2 * time.Millisecond,
+		DryRun:      true,
+	}
+	categories := []string{"ARR"}
+	statuses := []string{"parsed"}
+
+	var published int
+	pub := func(*telegram) error {
+		published++
+		return nil
+	}
+
+	if err := RunSeed(cfg, categories, statuses, pub); err != nil {
+		t.Fatalf("RunSeed mixed returned error: %v", err)
+	}
+	if published != cfg.Count {
+		t.Fatalf("expected %d published messages, got %d", cfg.Count, published)
+	}
+}
+
+