add assertion helper for extension keys (#920)

yanmxa · web-flow · commit e15d03d99ad0 · 2023-07-17T14:17:55.000-04:00
* add assertion helper for extension keys

Signed-off-by: myan &lt;myan@redhat.com&gt;

* fixed the fmt issue

Signed-off-by: myan &lt;myan@redhat.com&gt;

* avoid race condition

Signed-off-by: myan &lt;myan@redhat.com&gt;

* wording

Signed-off-by: myan &lt;myan@redhat.com&gt;

---------

Signed-off-by: myan &lt;myan@redhat.com&gt;
diff --git a/test/integration/kafka_sarama/kafka_test.go b/test/integration/kafka_sarama/kafka_test.go
@@ -35,15 +35,11 @@ func TestSendEvent(t *testing.T) {
 		clienttest.SendReceive(t, func() interface{} {
 			return protocolFactory(t)
 		}, eventIn, func(eventOut event.Event) {
-			eventOut = test.ConvertEventExtensionsToString(t, eventOut)
-
-			require.Equal(t, TopicName, eventOut.Extensions()[KAFKA_TOPIC])
-			require.NotNil(t, eventOut.Extensions()[KAFKA_PARTITION])
-			require.NotNil(t, eventOut.Extensions()[KAFKA_OFFSET])
-
 			test.AllOf(
 				test.HasExactlyAttributesEqualTo(eventIn.Context),
 				test.HasData(eventIn.Data()),
+				test.HasExtensionKeys([]string{KAFKA_OFFSET, KAFKA_PARTITION}),
+				test.HasExtension(KAFKA_TOPIC, TopicName),
 			)
 		})
 	})
diff --git a/test/integration/kafka_sarama_binding/kafka_test.go b/test/integration/kafka_sarama_binding/kafka_test.go
@@ -56,17 +56,12 @@ func TestSendBinaryMessageToBinary(t *testing.T) {
 
 		in := MustCreateMockBinaryMessage(eventIn)
 		test.SendReceive(t, binding.WithPreferredEventEncoding(context.TODO(), binding.EncodingBinary), in, s, r, func(out binding.Message) {
-			eventOut := MustToEvent(t, context.Background(), out)
 			assert.Equal(t, binding.EncodingBinary, out.ReadEncoding())
-			eventOut = ConvertEventExtensionsToString(t, eventOut)
-
-			require.Equal(t, topicName, eventOut.Extensions()[KAFKA_TOPIC])
-			require.NotNil(t, eventOut.Extensions()[KAFKA_PARTITION])
-			require.NotNil(t, eventOut.Extensions()[KAFKA_OFFSET])
-
 			AllOf(
 				HasExactlyAttributesEqualTo(eventIn.Context),
 				HasData(eventIn.Data()),
+				HasExtensionKeys([]string{KAFKA_OFFSET, KAFKA_PARTITION}),
+				HasExtension(KAFKA_TOPIC, topicName),
 			)
 		})
 	})
diff --git a/v2/test/event_matchers.go b/v2/test/event_matchers.go
@@ -184,6 +184,18 @@ func HasExtensions(ext map[string]interface{}) EventMatcher {
 	}
 }
 
+// HasExtensionKeys checks if the event contains the provided keys from its extensions
+func HasExtensionKeys(keys []string) EventMatcher {
+	return func(have event.Event) error {
+		for _, k := range keys {
+			if _, ok := have.Extensions()[k]; !ok {
+				return fmt.Errorf("expecting extension key %q", k)
+			}
+		}
+		return nil
+	}
+}
+
 // HasExtension checks if the event contains the provided extension
 func HasExtension(key string, value interface{}) EventMatcher {
 	return HasExtensions(map[string]interface{}{key: value})
@@ -277,7 +289,6 @@ func HasAttributeKind(kind spec.Kind, value interface{}) EventMatcher {
 // LICENSE: MIT License
 
 func isEmpty(object interface{}) bool {
-
 	// get nil case out of the way
 	if object == nil {
 		return true
