aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/rapi/middleware/middleware_test.go at 1435a9b64dd03ec7122d36c85dbbd69ec307548e · aws/aws-lambda-runtime-interface-emulator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package middleware

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/appctx"
	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/rapi/model"
)

type mockHandler struct{}

func (h *mockHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {}

func TestRuntimeReleaseMiddleware(t *testing.T) {
	appCtx := appctx.NewApplicationContext()
	router := chi.NewRouter()
	handler := &mockHandler{}
	router.Use(RuntimeReleaseMiddleware())
	router.Get("/", handler.ServeHTTP)

	userAgent := "foobar"

	responseRecorder := httptest.NewRecorder()
	responseBody := make([]byte, 100)
	request := httptest.NewRequest("GET", "/", bytes.NewReader(responseBody))
	request.Header.Set("User-Agent", userAgent)
	router.ServeHTTP(responseRecorder, appctx.RequestWithAppCtx(request, appCtx))

	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	ctxRuntimeRelease, ok := appCtx.Load(appctx.AppCtxRuntimeReleaseKey)
	assert.True(t, ok)
	assert.Equal(t, userAgent, ctxRuntimeRelease)
}

func TestAgentUniqueIdentifierHeaderValidatorForbidden(t *testing.T) {
	router := chi.NewRouter()
	mockHandler := &mockHandler{}
	router.Get("/", AgentUniqueIdentifierHeaderValidator(mockHandler).ServeHTTP)
	responseBody := make([]byte, 100)
	var errorResponse model.ErrorResponse

	request := httptest.NewRequest("GET", "/", bytes.NewReader(responseBody))

	responseRecorder := httptest.NewRecorder()
	router.ServeHTTP(responseRecorder, request)
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
	respBody, _ := io.ReadAll(responseRecorder.Body)
	require.NoError(t, json.Unmarshal(respBody, &errorResponse))
	assert.Equal(t, model.ErrAgentIdentifierMissing, errorResponse.ErrorType)

	responseRecorder = httptest.NewRecorder()
	request.Header.Set(model.LambdaAgentIdentifier, "invalid-unique-identifier")
	router.ServeHTTP(responseRecorder, request)
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
	respBody, _ = io.ReadAll(responseRecorder.Body)
	require.NoError(t, json.Unmarshal(respBody, &errorResponse))
	assert.Equal(t, model.ErrAgentIdentifierInvalid, errorResponse.ErrorType)
}

func TestAgentUniqueIdentifierHeaderValidatorSuccess(t *testing.T) {
	router := chi.NewRouter()
	mockHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		val, ok := r.Context().Value(model.AgentIDCtxKey).(uuid.UUID)
		if !ok {
			assert.FailNow(t, "expected key not in request context")
		}
		assert.Equal(t, "85083764-ff1e-476f-ada1-d51f26e4f6be", val.String())
	})
	router.Get("/", AgentUniqueIdentifierHeaderValidator(mockHandler).ServeHTTP)
	responseBody := make([]byte, 100)
	request := httptest.NewRequest("GET", "/", bytes.NewReader(responseBody))
	ctx := context.Background()
	request = request.WithContext(ctx)

	responseRecorder := httptest.NewRecorder()
	responseRecorder.Code = http.StatusOK
	request.Header.Set(model.LambdaAgentIdentifier, "85083764-ff1e-476f-ada1-d51f26e4f6be")
	router.ServeHTTP(responseRecorder, request)
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
}