Repository: etcd-io/etcd
File: tests/antithesis/test-template/robustness/main.go (Go)

Code before:
	if *local {
		hosts = []string{localetcd0, localetcd1, localetcd2}
	}

	ctx := context.Background()
	baseTime := time.Now()
	duration := time.Duration(robustnessrand.RandRange(5, 60) * int64(time.Second))
	testRobustness(ctx, hosts, baseTime, duration)
}

func testRobustness(ctx context.Context, hosts []string, baseTime time.Time, duration time.Duration) {
	lg, err := zap.NewProduction()
	if err != nil {
		panic(err)
	}
	lg.Info("Start traffic generation")
	reports := runTraffic(ctx, lg, hosts, baseTime, duration)
	lg.Info("Completed traffic generation")

	validateConfig := validate.Config{ExpectRevisionUnique: traffic.EtcdAntithesis.ExpectUniqueRevision()}
	result := validate.ValidateAndReturnVisualize(lg, validateConfig, reports, nil, 5*time.Minute)
	err = result.Linearization.Visualize(lg, "history.html")
	if err != nil {
		lg.Error("Failed to save visualization", zap.Error(result.Error))
	}
	assert.Always(result.Error == nil, "Robustness validation failed", map[string]any{"error": result.Error})
	lg.Info("Completed robustness validation")
}

func runTraffic(ctx context.Context, lg *zap.Logger, hosts []string, baseTime time.Time, duration time.Duration) []report.ClientReport {
	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
	finish := closeAfter(ctx, duration)
	ids := identity.NewIDProvider()
	storage := identity.NewLeaseIDStorage()
	concurrencyLimiter := traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)

	r, err := traffic.CheckEmptyDatabaseAtStart(ctx, lg, hosts, ids, baseTime)
	if err != nil {
		lg.Fatal("Failed empty database at start check", zap.Error(err))
	}

	reports := []report.ClientReport{r}
	var mux sync.Mutex
	var wg sync.WaitGroup
	for i := 0; i < profile.ClientCount; i++ {
		c := connect([]string{hosts[i%len(hosts)]}, ids, baseTime)
		defer c.Close()
		wg.Add(1)
		go func(c *client.RecordingClient) {
			defer wg.Done()
			defer c.Close()

Review comment (suggestion):
```suggestion
	assert.Always(result.Error == nil, "Robustness validation passes", map[string]any{"error": result.Error})
```

Code after:
	if *local {
		hosts = []string{localetcd0, localetcd1, localetcd2}
	}

	ctx := context.Background()
	baseTime := time.Now()
	duration := time.Duration(robustnessrand.RandRange(5, 60) * int64(time.Second))
	testRobustness(ctx, hosts, baseTime, duration)
}

func testRobustness(ctx context.Context, hosts []string, baseTime time.Time, duration time.Duration) {
	lg, err := zap.NewProduction()
	if err != nil {
		panic(err)
	}
	lg.Info("Start traffic generation")
	reports := runTraffic(ctx, lg, hosts, baseTime, duration)
	lg.Info("Completed traffic generation")

	validateConfig := validate.Config{ExpectRevisionUnique: traffic.EtcdAntithesis.ExpectUniqueRevision()}
	result := validate.ValidateAndReturnVisualize(lg, validateConfig, reports, nil, 5*time.Minute)
	err = result.Linearization.Visualize(lg, "history.html")
	if err != nil {
		lg.Error("Failed to save visualization", zap.Error(err))
	}
	assert.Always(result.Error == nil, "Robustness validation passes", map[string]any{"error": result.Error})
	lg.Info("Completed robustness validation", zap.Error(result.Error))
}

func runTraffic(ctx context.Context, lg *zap.Logger, hosts []string, baseTime time.Time, duration time.Duration) []report.ClientReport {
	limiter := rate.NewLimiter(rate.Limit(profile.MaximalQPS), profile.BurstableQPS)
	finish := closeAfter(ctx, duration)
	ids := identity.NewIDProvider()
	storage := identity.NewLeaseIDStorage()
	concurrencyLimiter := traffic.NewConcurrencyLimiter(profile.MaxNonUniqueRequestConcurrency)

	r, err := traffic.CheckEmptyDatabaseAtStart(ctx, lg, hosts, ids, baseTime)
	if err != nil {
		lg.Fatal("Failed empty database at start check", zap.Error(err))
	}

	reports := []report.ClientReport{r}
	var mux sync.Mutex
	var wg sync.WaitGroup
	for i := 0; i < profile.ClientCount; i++ {
		c := connect([]string{hosts[i%len(hosts)]}, ids, baseTime)
		defer c.Close()
		wg.Add(1)
		go func(c *client.RecordingClient) {
			defer wg.Done()
			defer c.Close()