Merge pull request #6 from DataDog/darcy.rayner/fix-context-only-handler

Fix support for handler single argument context input

Author: DarcyRaynerDD

internal/wrapper/wrap_handler.go

@@ -1,7 +1,7 @@
 /*
  * Unless explicitly stated otherwise all files in this repository are licensed
  * under the Apache License Version 2.0.
- * 
+ *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
  * Copyright 2019 Datadog, Inc.
  */
@@ -94,8 +94,15 @@ func callHandler(ctx context.Context, msg json.RawMessage, handler interface{})
 	args := []reflect.Value{}
 
 	if handlerType.NumIn() == 1 {
-		// When there is only one argument, argument is always the event payload.
-		args = []reflect.Value{ev.Elem()}
+		// When there is only one argument, argument is either the event payload, or the context.
+		contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
+		firstArgType := handlerType.In(0)
+		if firstArgType.Implements(contextType) {
+			args = []reflect.Value{reflect.ValueOf(ctx)}
+		} else {
+			args = []reflect.Value{ev.Elem()}
+
+		}
 	} else if handlerType.NumIn() == 2 {
 		// Or when there are two arguments, context is always first, followed by event payload.
 		args = []reflect.Value{reflect.ValueOf(ctx), ev.Elem()}
@@ -130,6 +137,13 @@ func unmarshalEventForHandler(ev json.RawMessage, handler interface{}) (reflect.
 	}
 
 	messageType := handlerType.In(handlerType.NumIn() - 1)
+	contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
+	firstArgType := handlerType.In(0)
+
+	if handlerType.NumIn() == 1 && firstArgType.Implements(contextType) {
+		return reflect.ValueOf(nil), nil
+	}
+
 	newMessage := reflect.New(messageType)
 	err := json.Unmarshal(ev, newMessage.Interface())
 	if err != nil {

internal/wrapper/wrap_handler_test.go

@@ -1,7 +1,7 @@
 /*
  * Unless explicitly stated otherwise all files in this repository are licensed
  * under the Apache License Version 2.0.
- * 
+ *
  * This product includes software developed at Datadog (https://www.datadoghq.com/).
  * Copyright 2019 Datadog, Inc.
  */
@@ -177,6 +177,20 @@ func TestWrapHandlerEventArgumentOnly(t *testing.T) {
 	assert.Equal(t, 5, response)
 }
 
+func TestWrapHandlerContextArgumentOnly(t *testing.T) {
+	called := true
+	var handler = func(ctx context.Context) (interface{}, error) {
+		return nil, nil
+	}
+
+	mhl := mockHandlerListener{}
+	wrappedHandler := WrapHandlerWithListeners(handler, &mhl).(func(context.Context, json.RawMessage) (interface{}, error))
+
+	_, err := wrappedHandler(context.Background(), nil)
+	assert.NoError(t, err)
+	assert.True(t, called)
+}
+
 func TestWrapHandlerNoArguments(t *testing.T) {
 	called := false
 
@@ -223,6 +237,22 @@ func TestWrapHandlerReturnsError(t *testing.T) {
 	assert.Equal(t, 5, response)
 }
 
+func TestWrapHandlerReturnsErrorOnly(t *testing.T) {
+	called := false
+	defaultErr := errors.New("Some error")
+
+	handler := func(request mockNonProxyEvent) error {
+		called = true
+		return defaultErr
+	}
+
+	_, response, err := runHandlerWithJSON(t, "../testdata/non-proxy-no-metadata.json", handler)
+
+	assert.True(t, called)
+	assert.Equal(t, defaultErr, err)
+	assert.Equal(t, nil, response)
+}
+
 func TestWrapHandlerReturnsOriginalHandlerIfInvalid(t *testing.T) {
 
 	var handler interface{} = func(arg1, arg2, arg3 int) (int, error) {