@@ -14,6 +14,7 @@ import (
1414 "testing"
1515 "unicode/utf8"
1616
17+ "github.com/aws/aws-lambda-go/lambdacontext"
1718 "github.com/stretchr/testify/assert"
1819)
1920
@@ -37,7 +38,7 @@ func TestRuntimeAPILoop(t *testing.T) {
3738 defer ts .Close ()
3839
3940 n := 0
40- handler := NewHandler (func () (string , error ) {
41+ handler := NewHandler (func (ctx context. Context ) (string , error ) {
4142 n += 1
4243 if n % 3 == 0 {
4344 return "" , errors .New ("error time!" )
@@ -51,6 +52,42 @@ func TestRuntimeAPILoop(t *testing.T) {
5152 assert .Equal (t , nInvokes , record .nPosts )
5253}
5354
55+ func TestRuntimeAPIContextPlumbing (t * testing.T ) {
56+ handler := NewHandler (func (ctx context.Context ) (interface {}, error ) {
57+ lc , _ := lambdacontext .FromContext (ctx )
58+ return lc , nil
59+ })
60+
61+ ts , record := runtimeAPIServer (`` , 1 )
62+ defer ts .Close ()
63+
64+ endpoint := strings .Split (ts .URL , "://" )[1 ]
65+ expectedError := fmt .Sprintf ("failed to GET http://%s/2018-06-01/runtime/invocation/next: got unexpected status code: 410" , endpoint )
66+ assert .EqualError (t , startRuntimeAPILoop (context .Background (), endpoint , handler ), expectedError )
67+
68+ expected := `
69+ {
70+ "AwsRequestID": "dummyid",
71+ "InvokedFunctionArn": "dummyarn",
72+ "Identity": {
73+ "CognitoIdentityID": "dummyident",
74+ "CognitoIdentityPoolID": "dummypool"
75+ },
76+ "ClientContext": {
77+ "Client": {
78+ "installation_id": "dummyinstallid",
79+ "app_title": "dummytitle",
80+ "app_version_code": "dummycode",
81+ "app_package_name": "dummyname"
82+ },
83+ "env": null,
84+ "custom": null
85+ }
86+ }
87+ `
88+ assert .JSONEq (t , expected , string (record .responses [0 ]))
89+ }
90+
5491func TestReadPayload (t * testing.T ) {
5592 ts , record := runtimeAPIServer (`{"message": "I am craving tacos"}` , 1 )
5693 defer ts .Close ()
@@ -87,9 +124,21 @@ func runtimeAPIServer(eventPayload string, failAfter int) (*httptest.Server, *re
87124 w .WriteHeader (http .StatusGone )
88125 _ , _ = w .Write ([]byte ("END THE TEST!" ))
89126 }
90- w .Header ().Add (string (headerAWSRequestID ), "dummy-request-id " )
127+ w .Header ().Add (string (headerAWSRequestID ), "dummyid " )
91128 w .Header ().Add (string (headerDeadlineMS ), "22" )
92- w .Header ().Add (string (headerInvokedFunctionARN ), "anarn" )
129+ w .Header ().Add (string (headerInvokedFunctionARN ), "dummyarn" )
130+ w .Header ().Add (string (headerClientContext ), `{
131+ "Client": {
132+ "app_title": "dummytitle",
133+ "installation_id": "dummyinstallid",
134+ "app_version_code": "dummycode",
135+ "app_package_name": "dummyname"
136+ }
137+ }` )
138+ w .Header ().Add (string (headerCognitoIdentity ), `{
139+ "cognitoIdentityId": "dummyident",
140+ "cognitoIdentityPoolId": "dummypool"
141+ }` )
93142 w .WriteHeader (http .StatusOK )
94143 _ , _ = w .Write ([]byte (eventPayload ))
95144 case http .MethodPost :
0 commit comments