diff --git a/pkg/grillkafka/assertions.go b/pkg/grillkafka/assertions.go index ac79629..3db7c8b 100644 --- a/pkg/grillkafka/assertions.go +++ b/pkg/grillkafka/assertions.go @@ -12,7 +12,7 @@ import ( func (gk *Kafka) AssertCount(topicName string, expectedCount int) grill.Assertion { return grill.AssertionFunc(func() error { group := fmt.Sprintf("%s_%d_%s", "oh_my_test_helper", rand.Intn(1000), time.Now()) - consumer, err := gk.NewConsumer(group, topicName, time.Second*5) + consumer, err := gk.NewConsumer(group, topicName, time.Second) if err != nil { return err } @@ -33,7 +33,7 @@ func (gk *Kafka) AssertCount(topicName string, expectedCount int) grill.Assertio func (gk *Kafka) AssertMessageCount(topic string, message Message, expectedCount int) grill.Assertion { return grill.AssertionFunc(func() error { group := fmt.Sprintf("%s_%d_%s", "oh_my_test_helper", rand.Intn(1000), time.Now()) - consumer, err := gk.NewConsumer(group, topic, time.Second*5) + consumer, err := gk.NewConsumer(group, topic, time.Second) if err != nil { return err } diff --git a/pkg/grillkafka/kafka_test.go b/pkg/grillkafka/kafka_test.go index 49956f2..4fa8b45 100644 --- a/pkg/grillkafka/kafka_test.go +++ b/pkg/grillkafka/kafka_test.go @@ -52,7 +52,7 @@ func Test_GrillKafka(t *testing.T) { }, Assertions: []grill.Assertion{ grill.AssertOutput(true), - grill.Try(time.Second*30, 3, helper.AssertCount("test_topic", 1)), + grill.Try(time.Second*30, 3, helper.AssertCount("test_topic", 1), grill.WithCheckFrequency(1*time.Second)), grill.Try(time.Second*30, 3, helper.AssertMessageCount("test_topic_error", testMessage, 1)), }, Cleaners: []grill.Cleaner{ diff --git a/try.go b/try.go index 24999f9..387c592 100644 --- a/try.go +++ b/try.go @@ -5,22 +5,38 @@ import ( "time" ) -func Try(deadline time.Duration, minSuccess int, assertion Assertion) Assertion { - return &tryAssertion{ - assertion: assertion, - deadline: deadline, - minSuccess: minSuccess, +type TryOption func(*tryAssertion) + +// WithCheckFrequency sets a custom frequency for assertion checks. +// If not set, frequency is derived from deadline and minSuccess. +func WithCheckFrequency(frequency time.Duration) TryOption { + return func(t *tryAssertion) { + t.checkFrequency = frequency + } +} + +func Try(deadline time.Duration, minSuccess int, assertion Assertion, options ...TryOption) Assertion { + t := &tryAssertion{ + assertion: assertion, + deadline: deadline, + minSuccess: minSuccess, + checkFrequency: deadline / time.Duration(minSuccess*3+3), + } + for _, opt := range options { + opt(t) } + return t } type tryAssertion struct { - assertion Assertion - deadline time.Duration - minSuccess int + assertion Assertion + deadline time.Duration + minSuccess int + checkFrequency time.Duration } func (assert *tryAssertion) Assert() error { - checkC := time.Tick(assert.deadline / time.Duration(assert.minSuccess*3+3)) + checkC := time.Tick(assert.checkFrequency) quitC := time.Tick(assert.deadline) var successCount = 0 var errors []string