Merge pull request #23 from FireTail-io/hotfix/close-race

Hotfix/close race

Author: rileyfiretail

logsapi/client.go

@@ -53,7 +53,6 @@ func (c *Client) Start(ctx context.Context) error {
 	if err != http.ErrServerClosed {
 		err = errors.WithMessage(err, "Log server closed unexpectedly")
 		c.errCallback(err)
-		c.Shutdown(ctx)
 	} else if err != nil {
 		c.errCallback(err)
 	}

main.go

@@ -55,5 +55,5 @@ func main() {
 
 	// Sleep for 500ms to allow any final logs to be sent to the extension by the Lambda Logs API
 	log.Printf("Sleeping for 500ms to allow final logs to be processed...")
-	time.Sleep(500)
+	time.Sleep(500 * time.Millisecond)
 }

main_test.go

@@ -6,6 +6,9 @@ import (
 	"net/http/httptest"
 	"strings"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestMain(t *testing.T) {
@@ -32,3 +35,19 @@ func TestMainDebug(t *testing.T) {
 
 	main()
 }
+
+func TestMainReturnsInNoLessThan500Milliseconds(t *testing.T) {
+	http.DefaultServeMux = new(http.ServeMux)
+	mockExtensionsApi := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintf(w, `{"eventType": "SHUTDOWN"}`)
+	}))
+	defer mockExtensionsApi.Close()
+
+	t.Setenv("AWS_LAMBDA_RUNTIME_API", strings.Join(strings.Split(mockExtensionsApi.URL, ":")[1:], ":")[2:])
+
+	startTime := time.Now()
+
+	main()
+
+	assert.Greater(t, time.Since(startTime), 500*time.Millisecond)
+}