From 6f5ddf291a1acda7e0ab5b32e142f0a6175c3132 Mon Sep 17 00:00:00 2001 From: Kyle Solie Date: Wed, 22 Aug 2018 13:13:52 -0500 Subject: [PATCH] Add lambda functionality Signed-off-by: Kyle Solie --- CHANGELOG.md | 1 + chaosaws/__init__.py | 2 + chaosaws/awslambda/__init__.py | 0 chaosaws/awslambda/actions.py | 41 +++++++++++++++++++ chaosaws/awslambda/probes.py | 16 ++++++++ requirements-dev.txt | 2 +- tests/awslambda/test_awslambda_actions.py | 49 +++++++++++++++++++++++ tests/awslambda/test_awslambda_probes.py | 12 ++++++ 8 files changed, 122 insertions(+), 1 deletion(-) create mode 100644 chaosaws/awslambda/__init__.py create mode 100644 chaosaws/awslambda/actions.py create mode 100644 chaosaws/awslambda/probes.py create mode 100644 tests/awslambda/test_awslambda_actions.py create mode 100644 tests/awslambda/test_awslambda_probes.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 472d3c5..a63539d 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/chaosaws/__init__.py b/chaosaws/__init__.py index 57994f9..81c8688 100644 --- a/chaosaws/__init__.py +++ b/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 diff --git a/chaosaws/awslambda/__init__.py b/chaosaws/awslambda/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/chaosaws/awslambda/actions.py b/chaosaws/awslambda/actions.py new file mode 100644 index 0000000..d7d5570 --- /dev/null +++ b/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) diff --git a/chaosaws/awslambda/probes.py b/chaosaws/awslambda/probes.py new file mode 100644 index 0000000..98e4265 --- /dev/null +++ b/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"] diff --git a/requirements-dev.txt b/requirements-dev.txt index a2c1f64..143f154 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,5 +1,5 @@ coverage -moto +moto==1.3.4 pycodestyle pytest>=2.8 pytest-cov diff --git a/tests/awslambda/test_awslambda_actions.py b/tests/awslambda/test_awslambda_actions.py new file mode 100644 index 0000000..89cfb47 --- /dev/null +++ b/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) diff --git a/tests/awslambda/test_awslambda_probes.py b/tests/awslambda/test_awslambda_probes.py new file mode 100644 index 0000000..4fa4b0d --- /dev/null +++ b/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)