Merge pull request #35 from random-element/direct_calls

freedomofkeima · web-flow · commit c6aae27eca36 · 2018-10-23T11:34:20.000+09:00
Support for direct calls to execute lambda functions and return their output (for py.tests)
diff --git a/.gitignore b/.gitignore
@@ -58,3 +58,6 @@ target/
 
 # Python environment
 .python-version
+
+# Vitual Environments
+venv/
diff --git a/lambda_local/environment_variables.py b/lambda_local/environment_variables.py
@@ -2,6 +2,11 @@
 import os
 
 
+def export_variables(environment_variables):
+    for env_name, env_value in environment_variables.items():
+        os.environ[str(env_name)] = str(env_value)
+
+
 def set_environment_variables(json_file_path):
     """
     Read and set environment variables from a flat json file.
@@ -25,5 +30,4 @@ def set_environment_variables(json_file_path):
         with open(json_file_path) as json_file:
             env_vars = json.loads(json_file.read())
 
-            for env_name, env_value in env_vars.items():
-                os.environ[str(env_name)] = str(env_value)
+            export_variables(env_vars)
diff --git a/lambda_local/main.py b/lambda_local/main.py
@@ -15,7 +15,7 @@
 
 from . import event
 from . import context
-from .environment_variables import set_environment_variables
+from .environment_variables import set_environment_variables, export_variables
 from .timeout import time_limit
 from .timeout import TimeoutException
 
@@ -31,6 +31,17 @@
 EXITCODE_ERR = 1
 
 
+def call(func, event, timeout, environment_variables={}, arn_string="", version_name="", library=None):
+    export_variables(environment_variables)
+    e = json.loads(event)
+    c = context.Context(timeout, arn_string, version_name)
+    if library is not None:
+        load_lib(library)
+    request_id = uuid.uuid4()
+
+    return _runner(request_id, e, c, func)
+
+
 def run(args):
     # set env vars if path to json file was given
     set_environment_variables(args.environment_variables)
@@ -41,16 +52,23 @@ def run(args):
         load_lib(args.library)
     request_id = uuid.uuid4()
     func = load(request_id, args.file, args.function)
+    
+    (result, err_type) = _runner(request_id, e, c, func)
+
+    if err_type is not None:
+        sys.exit(EXITCODE_ERR)
+
 
+def _runner(request_id, event, context, func):
     logger = logging.getLogger()
     result = None
 
-    logger.info("Event: {}".format(e))
+    logger.info("Event: {}".format(event))
 
     logger.info("START RequestId: {}".format(request_id))
 
     start_time = timeit.default_timer()
-    result, err_type = execute(func, e, c)
+    result, err_type = execute(func, event, context)
     end_time = timeit.default_timer()
 
     logger.info("END RequestId: {}".format(request_id))
@@ -64,8 +82,7 @@ def run(args):
     logger.info("REPORT RequestId: {}\tDuration: {}".format(
         request_id, duration))
 
-    if err_type is not None:
-        sys.exit(EXITCODE_ERR)
+    return (result, err_type)
 
 
 def load_lib(path):
diff --git a/tests/test_direct_invocations.py b/tests/test_direct_invocations.py
@@ -0,0 +1,52 @@
+'''
+python-lambda-local: Test Direct Inovactions 
+(command-line and direct).
+
+Meant for use with py.test.
+
+Copyright 2015 HDE, Inc.
+Licensed under MIT
+'''
+import json
+import argparse
+from multiprocessing import Process
+import os
+from lambda_local.main import run as lambda_run
+from lambda_local.main import call as lambda_call
+
+
+def my_lambda_function(event, context):
+    print("Hello World from My Lambda Function!")
+    return 42
+
+def test_function_call_for_pytest():
+    request = json.dumps({})
+    (result, error_type) = lambda_call(func=my_lambda_function, event=request, timeout=1)
+
+    assert error_type is None
+
+    assert result == 42
+
+
+def test_check_command_line():
+    request = json.dumps({})
+    request_file = 'check_command_line_event.json'
+    with open (request_file, "w") as f:
+        f.write(request)
+
+    args = argparse.Namespace(event=request_file,
+                              file='tests/test_direct_invocations.py',
+                              function='my_lambda_function',
+                              timeout=1,
+                              environment_variables='',
+                              library=None,
+                              version_name='',
+                              arn_string=''
+                              )
+    p = Process(target=lambda_run, args=(args,))
+    p.start()
+    p.join()
+
+    os.remove(request_file)
+    assert p.exitcode == 0
+    
