diff --git a/ddlambda.go b/ddlambda.go index 9ec5ca36..bde7939d 100644 --- a/ddlambda.go +++ b/ddlambda.go @@ -10,6 +10,7 @@ package ddlambda import ( "context" + "encoding/json" "fmt" "net/http" "os" @@ -111,7 +112,15 @@ 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 @@ -119,6 +128,16 @@ func MetricWithTimestamp(metric string, value float64, timestamp time.Time, tags 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) + 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 +149,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 +159,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) diff --git a/ddlambda_test.go b/ddlambda_test.go new file mode 100644 index 00000000..bc65f236 --- /dev/null +++ b/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) +} diff --git a/internal/wrapper/wrap_handler.go b/internal/wrapper/wrap_handler.go index dd421abc..8a4802b3 100644 --- a/internal/wrapper/wrap_handler.go +++ b/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 } }