@@ -13,9 +13,10 @@ def __init__(self, handler, base_path, timeout=None):
1313        self .timeout  =  timeout 
1414
1515    def  get_handler (self ):
16-         """ Load handler function. 
16+         """ 
17+         Load handler function. 
1718
18-              :returns function: Lambda handler function 
19+         :returns function: Lambda handler function 
1920        """ 
2021        * path , func  =  self .handler .split ('.' )
2122        name  =  '.' .join (path )
@@ -32,20 +33,43 @@ def get_handler(self):
3233        except  AttributeError :
3334            raise  ValueError (f"Handler '{ func }  ' missing on module '{ name }  '" )
3435
36+     def  get_httpMethod (self , event ):
37+         """ 
38+         Helper to get httpMethod from v1 or v2 events. 
39+         """ 
40+         if  event .get ('version' ) ==  '2.0' :
41+             return  event ['requestContext' ]['http' ]['method' ]
42+         elif  event .get ('version' ) ==  '1.0' :
43+             return  event ['httpMethod' ]
44+         raise  ValueError (  # pragma: no cover 
45+             f"Unknown API Gateway payload version: { event .get ('version' )}  " )
46+ 
47+     def  get_path (self , event ):
48+         """ 
49+         Helper to get path from v1 or v2 events. 
50+         """ 
51+         if  event .get ('version' ) ==  '2.0' :
52+             return  event ['rawPath' ]
53+         elif  event .get ('version' ) ==  '1.0' :
54+             return  event ['path' ]
55+         raise  ValueError (  # pragma: no cover 
56+             f"Unknown API Gateway payload version: { event .get ('version' )}  " )
57+ 
3558    def  invoke (self , event ):
3659        with  lambda_context .start (self .timeout ) as  context :
3760            logger .info ('Invoking "%s"' , self .handler )
3861            return  asyncio .run (self .invoke_async_with_timeout (event , context ))
3962
4063    async  def  invoke_async (self , event , context = None ):
41-         """ Wrapper to invoke the Lambda handler asynchronously. 
64+         """ 
65+         Wrapper to invoke the Lambda handler asynchronously. 
4266
43-              :param dict event: Lambda event object 
44-              :param Context context: Mock Lambda context 
45-              :returns dict: Lamnda invocation result 
67+         :param dict event: Lambda event object 
68+         :param Context context: Mock Lambda context 
69+         :returns dict: Lamnda invocation result 
4670        """ 
47-         httpMethod  =  event [ 'httpMethod' ] 
48-         path  =  event [ 'path' ] 
71+         httpMethod  =  self . get_httpMethod ( event ) 
72+         path  =  self . get_path ( event ) 
4973
5074        # Reject request if not starting at base path 
5175        if  not  path .startswith (self .base_path ):
@@ -64,27 +88,29 @@ async def invoke_async(self, event, context=None):
6488            return  self .jsonify (httpMethod , 502 , message = message )
6589
6690    async  def  invoke_async_with_timeout (self , event , context = None ):
67-         """ Wrapper to invoke the Lambda handler with a timeout. 
91+         """ 
92+         Wrapper to invoke the Lambda handler with a timeout. 
6893
69-              :param dict event: Lambda event object 
70-              :param Context context: Mock Lambda context 
71-              :returns dict: Lamnda invocation result or 408 TIMEOUT 
94+         :param dict event: Lambda event object 
95+         :param Context context: Mock Lambda context 
96+         :returns dict: Lamnda invocation result or 408 TIMEOUT 
7297        """ 
7398        try :
7499            coroutine  =  self .invoke_async (event , context )
75100            return  await  asyncio .wait_for (coroutine , self .timeout )
76101        except  asyncio .TimeoutError :
77-             httpMethod  =  event [ 'httpMethod' ] 
102+             httpMethod  =  self . get_httpMethod ( event ) 
78103            message  =  'Endpoint request timed out' 
79104            return  self .jsonify (httpMethod , 504 , message = message )
80105
81106    @staticmethod  
82107    def  jsonify (httpMethod , statusCode , ** kwargs ):
83-         """ Convert dict into API Gateway response object. 
108+         """ 
109+         Convert dict into API Gateway response object. 
84110
85-              :params str httpMethod: HTTP request method 
86-              :params int statusCode: Response status code 
87-              :params dict kwargs: Response object 
111+         :params str httpMethod: HTTP request method 
112+         :params int statusCode: Response status code 
113+         :params dict kwargs: Response object 
88114        """ 
89115        body  =  ''  if  httpMethod  in  ['HEAD' ] else  json .dumps (kwargs )
90116        return  {
0 commit comments