Merge pull request #24 from DataDog/darcy.rayner/improve-testability

Improve ability to run package code in tests

Author: DarcyRaynerDD

ddlambda.go

@@ -10,6 +10,7 @@ package ddlambda
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"os"
@@ -111,14 +112,33 @@ func Metric(metric string, value float64, tags ...string) {
 
 // MetricWithTimestamp sends a distribution metric to DataDog with a custom timestamp
 func MetricWithTimestamp(metric string, value float64, timestamp time.Time, tags ...string) {
-	listener := metrics.GetListener(GetContext())
+	ctx := GetContext()
+
+	if ctx == nil {
+		logger.Debug("no context available, did you wrap your handler?")
+		return
+	}
+
+	listener := metrics.GetListener(ctx)
+
 	if listener == nil {
 		logger.Error(fmt.Errorf("couldn't get metrics listener from current context"))
 		return
 	}
 	listener.AddDistributionMetric(metric, value, timestamp, tags...)
 }
 
+// InvokeDryRun is a utility to easily run your lambda for testing
+func InvokeDryRun(callback func(ctx context.Context), cfg *Config) (interface{}, error) {
+	wrapped := WrapHandler(callback, cfg)
+	// Convert the wrapped handler to it's underlying raw handler type
+	handler, ok := wrapped.(func(ctx context.Context, msg json.RawMessage) (interface{}, error))
+	if !ok {
+		logger.Debug("Could not unwrap lambda during dry run")
+	}
+	return handler(context.Background(), json.RawMessage("{}"))
+}
+
 func (cfg *Config) toMetricsConfig() metrics.Config {
 
 	mc := metrics.Config{
@@ -130,6 +150,7 @@ func (cfg *Config) toMetricsConfig() metrics.Config {
 		mc.ShouldRetryOnFailure = cfg.ShouldRetryOnFailure
 		mc.APIKey = cfg.APIKey
 		mc.KMSAPIKey = cfg.KMSAPIKey
+		mc.Site = cfg.Site
 		mc.ShouldUseLogForwarder = cfg.ShouldUseLogForwarder
 	}
 
@@ -139,7 +160,11 @@ func (cfg *Config) toMetricsConfig() metrics.Config {
 	if mc.Site == "" {
 		mc.Site = DefaultSite
 	}
-	mc.Site = fmt.Sprintf("https://api.%s/api/v1", mc.Site)
+	if strings.HasPrefix(mc.Site, "https://") || strings.HasPrefix(mc.Site, "http://") {
+		mc.Site = fmt.Sprintf("%s/api/v1", mc.Site)
+	} else {
+		mc.Site = fmt.Sprintf("https://api.%s/api/v1", mc.Site)
+	}
 
 	if !mc.ShouldUseLogForwarder {
 		shouldUseLogForwarder := os.Getenv(DatadogShouldUseLogForwarderEnvVar)

ddlambda_test.go

@@ -0,0 +1,48 @@
+/*
+ * 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.
+ */
+package ddlambda
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInvokeDryRun(t *testing.T) {
+	called := false
+	InvokeDryRun(func(ctx context.Context) {
+		called = true
+		globalCtx := GetContext()
+		assert.Equal(t, globalCtx, ctx)
+	}, nil)
+	assert.True(t, called)
+}
+
+func TestMetricsSilentFailWithoutWrapper(t *testing.T) {
+	Metric("my-metric", 100, "my:tag")
+}
+
+func TestMetricsSubmitWithWrapper(t *testing.T) {
+	called := false
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		called = true
+		w.WriteHeader(http.StatusCreated)
+	}))
+	defer server.Close()
+
+	InvokeDryRun(func(ctx context.Context) {
+		Metric("my-metric", 100, "my:tag")
+	}, &Config{
+		APIKey: "abc-123",
+		Site:   server.URL,
+	})
+	assert.True(t, called)
+}

internal/wrapper/wrap_handler.go

@@ -51,6 +51,7 @@ func WrapHandlerWithListeners(handler interface{}, listeners ...HandlerListener)
 		for _, listener := range listeners {
 			listener.HandlerFinished(ctx)
 		}
+		CurrentContext = nil
 		return result, err
 	}
 }