fix: address review comments on test reliability and correctness

- Replace time.Sleep with deadline-based polling in TestInputModule_PublishMessages and TestBentoTrigger_StartStop
- Remove unused wantType field from TestBentoPlugin_CreateModule test table
- Add context timeout to TestProcessorStep_ExecuteWithInvalidBloblang
- Fix TestBentoTrigger_CallbackError to actually return a non-nil error and verify clean stop
- Fix count() Bloblang calls to include required name argument
- Add testmain_test.go to register Bento pure components for tests

Co-authored-by: intel352 <77607+intel352@users.noreply.github.com>

Author: Copilot

go.mod

@@ -30,6 +30,7 @@ require (
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0 // indirect
 	github.com/IBM/sarama v1.46.3 // indirect
 	github.com/Jeffail/gabs/v2 v2.7.0 // indirect
+	github.com/Jeffail/grok v1.1.0 // indirect
 	github.com/Jeffail/shutdown v1.0.0 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/OneOfOne/xxhash v1.2.8 // indirect
@@ -104,9 +105,12 @@ require (
 	github.com/hashicorp/go-secure-stdlib/strutil v0.1.2 // indirect
 	github.com/hashicorp/go-sockaddr v1.0.7 // indirect
 	github.com/hashicorp/go-uuid v1.0.3 // indirect
+	github.com/hashicorp/golang-lru/arc/v2 v2.0.7 // indirect
+	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
 	github.com/hashicorp/hcl v1.0.1-vault-7 // indirect
 	github.com/hashicorp/vault/api v1.22.0 // indirect
 	github.com/hashicorp/yamux v0.1.2 // indirect
+	github.com/influxdata/go-syslog/v3 v3.0.0 // indirect
 	github.com/itchyny/gojq v0.12.18 // indirect
 	github.com/itchyny/timefmt-go v0.1.7 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
@@ -118,8 +122,10 @@ require (
 	github.com/jcmturner/gofork v1.7.6 // indirect
 	github.com/jcmturner/gokrb5/v8 v8.4.4 // indirect
 	github.com/jcmturner/rpc/v2 v2.0.3 // indirect
+	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/klauspost/compress v1.18.1 // indirect
+	github.com/klauspost/pgzip v1.2.6 // indirect
 	github.com/matoous/go-nanoid/v2 v2.0.0 // indirect
 	github.com/mattn/go-colorable v0.1.14 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
@@ -147,9 +153,13 @@ require (
 	github.com/prometheus/client_model v0.6.2 // indirect
 	github.com/prometheus/common v0.48.0 // indirect
 	github.com/prometheus/procfs v0.12.0 // indirect
+	github.com/quipo/dependencysolver v0.0.0-20170801134659-2b009cb4ddcc // indirect
 	github.com/rcrowley/go-metrics v0.0.0-20250401214520-65e299d6c5c9 // indirect
 	github.com/redis/go-redis/v9 v9.18.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
+	github.com/rickb777/date v1.20.5 // indirect
+	github.com/rickb777/plural v1.4.1 // indirect
+	github.com/robfig/cron/v3 v3.0.1 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/ryanuber/go-glob v1.0.0 // indirect
 	github.com/segmentio/ksuid v1.0.4 // indirect

go.sum

@@ -207,6 +207,7 @@ github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
 github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/golobby/cast v1.3.3 h1:s2Lawb9RMz7YyYf8IrfMQY4IFmA1R/lgfmj97Vc6fig=
 github.com/golobby/cast v1.3.3/go.mod h1:0oDO5IT84HTXcbLDf1YXuk0xtg/cRDrxhbpWKxwtJCY=
+github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
@@ -306,6 +307,8 @@ github.com/jhump/protoreflect v1.16.0 h1:54fZg+49widqXYQ0b+usAFHbMkBGR4PpXrsHc8+
 github.com/jhump/protoreflect v1.16.0/go.mod h1:oYPd7nPvcBw/5wlDfm/AVmU9zH9BgqGCI469pGxfj/8=
 github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
 github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
+github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=
+github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
 github.com/klauspost/compress v1.18.1 h1:bcSGx7UbpBqMChDtsF28Lw6v/G94LPrrbMbdC3JH2co=
@@ -321,6 +324,7 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/leodido/ragel-machinery v0.0.0-20181214104525-299bdde78165/go.mod h1:WZxr2/6a/Ar9bMDc2rN/LJrE/hF6bXE4LPyDSIxwAfg=
 github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
 github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/linkedin/goavro/v2 v2.12.0 h1:rIQQSj8jdAUlKQh6DttK8wCRv4t4QO09g1C4aBWXslg=
@@ -377,6 +381,8 @@ github.com/oklog/run v1.2.0 h1:O8x3yXwah4A73hJdlrwo/2X6J62gE5qTMusH0dvz60E=
 github.com/oklog/run v1.2.0/go.mod h1:mgDbKRSwPhJfesJ4PntqFUbKQRZ50NgmZTSPlFA0YFk=
 github.com/oklog/ulid v1.3.1 h1:EGfNDEx6MqHz8B3uNV6QAib1UR2Lm97sHi3ocA6ESJ4=
 github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn4U=
+github.com/onsi/gomega v1.27.10 h1:naR28SdDFlqrG6kScpT8VWpu1xWY5nJRCF3XaYyBjhI=
+github.com/onsi/gomega v1.27.10/go.mod h1:RsS8tutOdbdgzbPtzzATp12yT7kM5I5aElG3evPbQ0M=
 github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040=
@@ -449,6 +455,10 @@ github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/tilinna/z85 v1.0.0 h1:uqFnJBlD01dosSeo5sK1G1YGbPuwqVHqR+12OJDRjUw=
 github.com/tilinna/z85 v1.0.0/go.mod h1:EfpFU/DUY4ddEy6CRvk2l+UQNEzHbh+bqBQS+04Nkxs=
+github.com/trivago/grok v1.0.0 h1:oV2ljyZT63tgXkmgEHg2U0jMqiKKuL0hkn49s6aRavQ=
+github.com/trivago/grok v1.0.0/go.mod h1:9t59xLInhrncYq9a3J7488NgiBZi5y5yC7bss+w4NHM=
+github.com/trivago/tgo v1.0.7 h1:uaWH/XIy9aWYWpjm2CU3RpcqZXmX2ysQ9/Go+d9gyrM=
+github.com/trivago/tgo v1.0.7/go.mod h1:w4dpD+3tzNIIiIfkWWa85w5/B77tlvdZckQ+6PkFnhc=
 github.com/urfave/cli/v2 v2.27.1 h1:8xSQ6szndafKVRmfyeUMxkNUJQMjL1F2zmsZ+qHpfho=
 github.com/urfave/cli/v2 v2.27.1/go.mod h1:8qnjx1vcq5s2/wpsqoZFndg2CE5tNFyrTvS6SinrnYQ=
 github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
@@ -596,6 +606,9 @@ gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EV
 gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc=
 gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/input_module_test.go

@@ -189,7 +189,7 @@ func TestInputModule_PublishMessages(t *testing.T) {
 		"target_topic": "test-topic",
 		"input": map[string]any{
 			"generate": map[string]any{
-				"mapping":  `root = {"id": count(), "msg": "hello"}`,
+				"mapping":  `root = {"id": count("input_id"), "msg": "hello"}`,
 				"count":    3,
 				"interval": "10ms",
 			},
@@ -207,8 +207,14 @@ func TestInputModule_PublishMessages(t *testing.T) {
 		t.Fatalf("Start() error = %v", err)
 	}
 
-	// Wait for messages to be published
-	time.Sleep(200 * time.Millisecond)
+	// Poll until 3 messages are published or deadline is reached.
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		if len(pub.GetMessages()) >= 3 {
+			break
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
 
 	stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 	defer cancel()

internal/plugin_test.go

@@ -75,35 +75,30 @@ func TestBentoPlugin_CreateModule(t *testing.T) {
 		typeName string
 		name     string
 		config   map[string]any
-		wantType string
 		wantErr  bool
 	}{
 		{
 			typeName: "bento.stream",
 			name:     "test-stream",
 			config:   map[string]any{},
-			wantType: "*internal.streamModule",
 			wantErr:  false,
 		},
 		{
 			typeName: "bento.input",
 			name:     "test-input",
 			config:   map[string]any{},
-			wantType: "*internal.inputModule",
 			wantErr:  false,
 		},
 		{
 			typeName: "bento.output",
 			name:     "test-output",
 			config:   map[string]any{},
-			wantType: "*internal.outputModule",
 			wantErr:  false,
 		},
 		{
 			typeName: "bento.broker",
 			name:     "test-broker",
 			config:   map[string]any{},
-			wantType: "*internal.brokerModule",
 			wantErr:  false,
 		},
 		{

internal/processor_step_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"testing"
+	"time"
 )
 
 func TestNewProcessorStep(t *testing.T) {
@@ -177,7 +178,8 @@ func TestProcessorStep_ExecuteWithInvalidBloblang(t *testing.T) {
 		t.Fatalf("newProcessorStep() error = %v", err)
 	}
 
-	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
 	result, err := s.Execute(ctx, map[string]any{"data": "test"}, nil, map[string]any{}, nil)
 
 	if err == nil {

internal/testmain_test.go

@@ -0,0 +1,7 @@
+package internal
+
+import (
+	// Register pure Bento components (e.g. generate input, bloblang processor)
+	// so that tests can use them without importing the full component suite.
+	_ "github.com/warpstreamlabs/bento/v4/public/components/pure"
+)

internal/trigger_test.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"errors"
 	"sync"
 	"testing"
 	"time"
@@ -169,7 +170,7 @@ func TestBentoTrigger_StartStop(t *testing.T) {
 			map[string]any{
 				"input": map[string]any{
 					"generate": map[string]any{
-						"mapping":  `root = {"id": count()}`,
+						"mapping":  `root = {"id": count("trigger_id")}`,
 						"count":    2,
 						"interval": "50ms",
 					},
@@ -188,8 +189,14 @@ func TestBentoTrigger_StartStop(t *testing.T) {
 		t.Fatalf("Start() error = %v", err)
 	}
 
-	// Wait for messages to be processed
-	time.Sleep(300 * time.Millisecond)
+	// Poll until 2 callbacks are observed or deadline is reached.
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		if len(cb.GetCalls()) >= 2 {
+			break
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
 
 	stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 	defer cancel()
@@ -314,19 +321,32 @@ func TestBentoTrigger_NoSubscriptions(t *testing.T) {
 }
 
 func TestBentoTrigger_CallbackError(t *testing.T) {
-	// Callback that returns error
-	var callbackErr error
+	// Callback that errors on the first invocation and succeeds on subsequent
+	// ones. This exercises the error path while still allowing the stream to
+	// terminate cleanly (Bento retries the NACKed message once, then continues).
+	var (
+		cbMu    sync.Mutex
+		cbCount int
+	)
 	errorCb := func(action string, data map[string]any) error {
-		return callbackErr
+		cbMu.Lock()
+		cbCount++
+		count := cbCount
+		cbMu.Unlock()
+		if count == 1 {
+			return errors.New("first callback error")
+		}
+		return nil
 	}
 
 	trigger, err := newBentoTrigger(map[string]any{
 		"subscriptions": []any{
 			map[string]any{
 				"input": map[string]any{
 					"generate": map[string]any{
-						"mapping": `root = {"test": "data"}`,
-						"count":   1,
+						"mapping":  `root = {"test": "data"}`,
+						"count":    3,
+						"interval": "10ms",
 					},
 				},
 				"workflow": "test",
@@ -337,22 +357,37 @@ func TestBentoTrigger_CallbackError(t *testing.T) {
 		t.Fatalf("newBentoTrigger() error = %v", err)
 	}
 
-	callbackErr = nil // First call succeeds
-
 	ctx := context.Background()
 	if err := trigger.Start(ctx); err != nil {
 		t.Fatalf("Start() error = %v", err)
 	}
 
-	time.Sleep(100 * time.Millisecond)
+	// Poll until at least 2 callbacks are observed (1 error + 1 success).
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		cbMu.Lock()
+		count := cbCount
+		cbMu.Unlock()
+		if count >= 2 {
+			break
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
 
-	stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
-	// Should stop cleanly even if callback had errors
+	// Should stop cleanly even when a callback returned an error.
 	if err := trigger.Stop(stopCtx); err != nil {
 		t.Errorf("Stop() error = %v", err)
 	}
+
+	cbMu.Lock()
+	count := cbCount
+	cbMu.Unlock()
+	if count == 0 {
+		t.Error("expected callback to be invoked at least once")
+	}
 }
 
 func TestBentoTrigger_StopWithoutStart(t *testing.T) {