Merge pull request #14 from FireTail-io/test-and-refactor

Test and refactor

Author: TheTeaCat

agent/http.go

@@ -0,0 +1,30 @@
+package main
+
+import (
+	"context"
+	"firetail-lambda-extension/extensionsapi"
+	"log"
+
+	"github.com/pkg/errors"
+)
+
+// awaitShutdown calls /event/next until a shutdown event is received, or the context is cancelled.
+// It returns a reason, or an error, depending upon the cause of the shutdown.
+func awaitShutdown(extensionClient *extensionsapi.Client, ctx context.Context) (string, error) {
+	for {
+		select {
+		case <-ctx.Done():
+			log.Printf(", returning...")
+			return "context cancelled", nil
+		default:
+			log.Printf("Waiting for event...")
+			res, err := extensionClient.NextEvent(ctx) // This is a blocking call
+			if err != nil {
+				return "", errors.WithMessage(err, "failed to get next event")
+			}
+			if res.EventType == extensionsapi.Shutdown {
+				return "received shutdown event", nil
+			}
+		}
+	}
+}

await_shutdown.go

@@ -0,0 +1,62 @@
+package main
+
+import (
+	"context"
+	"firetail-lambda-extension/extensionsapi"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func getMockExtensionsApiServer() *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintf(w, `{"eventType": "SHUTDOWN"}`)
+	}))
+}
+
+func TestAwaitShutdownContextCancelled(t *testing.T) {
+	extensionClient := extensionsapi.NewClient()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	cancel()
+
+	reason, err := awaitShutdown(extensionClient, ctx)
+
+	assert.Nil(t, err)
+	assert.Equal(t, "context cancelled", reason)
+}
+
+func TestAwaitShutdownNextEventErrs(t *testing.T) {
+	extensionClient := extensionsapi.NewClient()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+
+	reason, err := awaitShutdown(extensionClient, ctx)
+
+	require.NotNil(t, err)
+	assert.Contains(t, err.Error(), "failed to get next event")
+	assert.Equal(t, "", reason)
+
+	cancel()
+}
+
+func TestAwaitShutdownShutdownEvent(t *testing.T) {
+	mockExtensionsApi := getMockExtensionsApiServer()
+	defer mockExtensionsApi.Close()
+
+	t.Setenv("AWS_LAMBDA_RUNTIME_API", strings.Join(strings.Split(mockExtensionsApi.URL, ":")[1:], ":")[2:])
+	extensionClient := extensionsapi.NewClient()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	reason, err := awaitShutdown(extensionClient, ctx)
+	require.Nil(t, err)
+	assert.Equal(t, "received shutdown event", reason)
+}

await_shutdown_test.go

@@ -0,0 +1,22 @@
+package main
+
+import (
+	"context"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+// Returns a context which will be cancelled whenever a SIGTERM or SIGINT signal is received
+func getContext() context.Context {
+	ctx, cancel := context.WithCancel(context.Background())
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGTERM, syscall.SIGINT)
+	go func() {
+		s := <-sigs
+		log.Printf("Received signal '%s'. Exiting...", s.String())
+		cancel()
+	}()
+	return ctx
+}