@@ -36,14 +36,14 @@ func (h testWrapperHandler) Invoke(ctx context.Context, payload []byte) ([]byte,
3636var _ Handler = (testWrapperHandler )(nil )
3737
3838func TestInvoke (t * testing.T ) {
39- srv := & Function { handler : testWrapperHandler (
39+ srv := NewFunction ( testWrapperHandler (
4040 func (ctx context.Context , input []byte ) (interface {}, error ) {
4141 if deadline , ok := ctx .Deadline (); ok {
4242 return deadline .UnixNano (), nil
4343 }
4444 return nil , errors .New ("!?!?!?!?!" )
4545 },
46- )}
46+ ))
4747 deadline := time .Now ()
4848 var response messages.InvokeResponse
4949 err := srv .Invoke (& messages.InvokeRequest {
@@ -59,15 +59,17 @@ func TestInvoke(t *testing.T) {
5959
6060func TestInvokeWithContext (t * testing.T ) {
6161 key := struct {}{}
62- srv := NewFunction (testWrapperHandler (
63- func (ctx context.Context , input []byte ) (interface {}, error ) {
64- assert .Equal (t , "dummy" , ctx .Value (key ))
65- if deadline , ok := ctx .Deadline (); ok {
66- return deadline .UnixNano (), nil
67- }
68- return nil , errors .New ("!?!?!?!?!" )
69- }))
70- srv = srv .withContext (context .WithValue (context .Background (), key , "dummy" ))
62+ srv := NewFunction (& handlerOptions {
63+ Handler : testWrapperHandler (
64+ func (ctx context.Context , input []byte ) (interface {}, error ) {
65+ assert .Equal (t , "dummy" , ctx .Value (key ))
66+ if deadline , ok := ctx .Deadline (); ok {
67+ return deadline .UnixNano (), nil
68+ }
69+ return nil , errors .New ("!?!?!?!?!" )
70+ }),
71+ baseContext : context .WithValue (context .Background (), key , "dummy" ),
72+ })
7173 deadline := time .Now ()
7274 var response messages.InvokeResponse
7375 err := srv .Invoke (& messages.InvokeRequest {
@@ -86,12 +88,11 @@ type CustomError struct{}
8688func (e CustomError ) Error () string { return "Something bad happened!" }
8789
8890func TestCustomError (t * testing.T ) {
89-
90- srv := & Function {handler : testWrapperHandler (
91+ srv := NewFunction (testWrapperHandler (
9192 func (ctx context.Context , input []byte ) (interface {}, error ) {
9293 return nil , CustomError {}
9394 },
94- )}
95+ ))
9596 var response messages.InvokeResponse
9697 err := srv .Invoke (& messages.InvokeRequest {}, & response )
9798 assert .NoError (t , err )
@@ -106,11 +107,11 @@ func (e *CustomError2) Error() string { return "Something bad happened!" }
106107
107108func TestCustomErrorRef (t * testing.T ) {
108109
109- srv := & Function { handler : testWrapperHandler (
110+ srv := NewFunction ( testWrapperHandler (
110111 func (ctx context.Context , input []byte ) (interface {}, error ) {
111112 return nil , & CustomError2 {}
112113 },
113- )}
114+ ))
114115 var response messages.InvokeResponse
115116 err := srv .Invoke (& messages.InvokeRequest {}, & response )
116117 assert .NoError (t , err )
@@ -120,12 +121,12 @@ func TestCustomErrorRef(t *testing.T) {
120121}
121122
122123func TestContextPlumbing (t * testing.T ) {
123- srv := & Function { handler : testWrapperHandler (
124+ srv := NewFunction ( testWrapperHandler (
124125 func (ctx context.Context , input []byte ) (interface {}, error ) {
125126 lc , _ := lambdacontext .FromContext (ctx )
126127 return lc , nil
127128 },
128- )}
129+ ))
129130 var response messages.InvokeResponse
130131 err := srv .Invoke (& messages.InvokeRequest {
131132 CognitoIdentityId : "dummyident" ,
@@ -172,14 +173,14 @@ func TestXAmznTraceID(t *testing.T) {
172173 Ctx string
173174 }
174175
175- srv := & Function { handler : testWrapperHandler (
176+ srv := NewFunction ( testWrapperHandler (
176177 func (ctx context.Context , input []byte ) (interface {}, error ) {
177178 return & XRayResponse {
178179 Env : os .Getenv ("_X_AMZN_TRACE_ID" ),
179180 Ctx : ctx .Value ("x-amzn-trace-id" ).(string ),
180181 }, nil
181182 },
182- )}
183+ ))
183184
184185 sequence := []struct {
185186 Input string
0 commit comments