Add lambda functionality

CHANGELOG.md

@@ -14,6 +14,7 @@
 -   support for elbv2 with basic probes and action
 -   support for asg with basic probes
 -   fix asg probe to support pagination
+-   support for lambda with basic probes and action
 
 ### Changed
 

chaosaws/__init__.py

@@ -181,5 +181,7 @@ def load_exported_activities() -> List[DiscoveredActivities]:
     activities.extend(discover_actions("chaosaws.elbv2.actions"))
     activities.extend(discover_probes("chaosaws.elbv2.probes"))
     activities.extend(discover_probes("chaosaws.asg.probes"))
+    activities.extend(discover_actions("chaosaws.awslambda.actions"))
+    activities.extend(discover_probes("chaosaws.awslambda.probes"))
 
     return activities

chaosaws/awslambda/actions.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+from chaoslib.exceptions import FailedActivity
+from chaoslib.types import Configuration, Secrets
+
+from chaosaws import aws_client
+from chaosaws.types import AWSResponse
+
+__all__ = ["put_function_concurrency", "delete_function_concurrency"]
+
+
+def put_function_concurrency(function_name: str,
+                             concurrent_executions: int,
+                             configuration: Configuration = None,
+                             secrets: Secrets = None) -> AWSResponse:
+    """
+    Throttles Lambda by setting reserved concurrency amount.
+    """
+    client = aws_client("lambda", configuration, secrets)
+    if not function_name:
+        raise FailedActivity(
+            "you must specify the lambda function name"
+        )
+    try:
+        return client.put_function_concurrency(
+            FunctionName=function_name,
+            ReservedConcurrentExecutions=concurrent_executions
+        )
+    except Exception as x:
+        raise FailedActivity(
+            "failed throttling lambda function '{}': '{}'".format(
+                function_name, str(x)))
+
+
+def delete_function_concurrency(function_name: str,
+                                configuration: Configuration = None,
+                                secrets: Secrets = None) -> AWSResponse:
+    """
+    Removes concurrency limit applied to the specified Lambda
+    """
+    client = aws_client("lambda", configuration, secrets)
+    return client.delete_function_concurrency(FunctionName=function_name)

chaosaws/awslambda/probes.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+from chaoslib.types import Configuration, Secrets
+from chaosaws import aws_client
+
+__all__ = ["get_function_concurrency"]
+
+
+def get_function_concurrency(function_name: str, configuration:
+                             Configuration = None,
+                             secrets: Secrets = None) -> bool:
+    """
+    Get configuration information of lambda by its function name
+    """
+    client = aws_client("lambda", configuration, secrets)
+    result = client.get_function(FunctionName=function_name)
+    return result["Concurrency"]["ReservedConcurrentExecutions"]

requirements-dev.txt

@@ -1,5 +1,5 @@
 coverage
-moto
+moto==1.3.4
 pycodestyle
 pytest>=2.8
 pytest-cov

tests/awslambda/test_awslambda_actions.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+from unittest.mock import MagicMock, patch
+
+import pytest
+from chaoslib.exceptions import FailedActivity
+
+from chaosaws.awslambda.actions import put_function_concurrency
+from chaosaws.awslambda.actions import delete_function_concurrency
+
+
+@patch('chaosaws.awslambda.actions.aws_client', autospec=True)
+def test_aws_lambda_put_function_concurrency(aws_client):
+    client = MagicMock()
+    aws_client.return_value = client
+    lambda_function_name = 'my-lambda-function'
+    concurrency = 0
+    put_function_concurrency(lambda_function_name, concurrency)
+    client.put_function_concurrency.assert_called_with(FunctionName=lambda_function_name,
+                                                       ReservedConcurrentExecutions=concurrency)
+
+@patch('chaosaws.awslambda.actions.aws_client', autospec=True)
+def test_aws_lambda_put_function_concurrency_empty_string(aws_client):
+    client = MagicMock()
+    aws_client.return_value = client
+    lambda_function_name = ''
+    concurrency = 0
+    with pytest.raises(FailedActivity):
+        put_function_concurrency(lambda_function_name, concurrency)
+
+
+@patch('chaosaws.awslambda.actions.aws_client', autospec=True)
+def test_aws_lambda_put_function_concurrency_failedactivity(aws_client):
+    client = MagicMock()
+    aws_client.return_value = client
+    lambda_function_name = 'my-lambda-function'
+    concurrency = 0
+
+    with patch.object(client, 'put_function_concurrency', FailedActivity):
+        with pytest.raises(Exception):
+            put_function_concurrency(lambda_function_name, concurrency)
+
+
+@patch('chaosaws.awslambda.actions.aws_client', autospec=True)
+def test_aws_lambda_delete_function_concurrency(aws_client):
+    client = MagicMock()
+    aws_client.return_value = client
+    lambda_function_name = 'my-lambda-function'
+    delete_function_concurrency(lambda_function_name)
+    client.delete_function_concurrency.assert_called_with(FunctionName=lambda_function_name)

tests/awslambda/test_awslambda_probes.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+from unittest.mock import MagicMock, patch
+from chaosaws.awslambda.probes import get_function_concurrency
+
+
+@patch('chaosaws.awslambda.probes.aws_client', autospec=True)
+def test_aws_lambda_get_function_concurrency(aws_client):
+    client = MagicMock()
+    aws_client.return_value = client
+    lambda_function_name = 'my-lambda-function'
+    get_function_concurrency(lambda_function_name)
+    client.get_function.assert_called_with(FunctionName=lambda_function_name)