From 6f5c80962b905af96fbcacefcfa93f7d12d00548 Mon Sep 17 00:00:00 2001 From: Jorge Tapicha Date: Wed, 29 May 2024 23:43:29 -0500 Subject: [PATCH 1/5] chaosaws msk ready Signed-off-by: Jorge Tapicha Signed-off-by: Jorge Tapicha --- chaosaws/msk/__init__.py | 0 chaosaws/msk/actions.py | 48 +++++++++++++++++ chaosaws/msk/probes.py | 43 +++++++++++++++ tests/msk/test_msk_actions.py | 91 ++++++++++++++++++++++++++++++++ tests/msk/test_msk_probes.py | 98 +++++++++++++++++++++++++++++++++++ 5 files changed, 280 insertions(+) create mode 100644 chaosaws/msk/__init__.py create mode 100644 chaosaws/msk/actions.py create mode 100644 chaosaws/msk/probes.py create mode 100644 tests/msk/test_msk_actions.py create mode 100644 tests/msk/test_msk_probes.py diff --git a/chaosaws/msk/__init__.py b/chaosaws/msk/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/chaosaws/msk/actions.py b/chaosaws/msk/actions.py new file mode 100644 index 00000000..2443fbfd --- /dev/null +++ b/chaosaws/msk/actions.py @@ -0,0 +1,48 @@ +from typing import List + +from chaoslib.types import Configuration, Secrets +from chaoslib.exceptions import FailedActivity +from chaosaws import aws_client, get_logger +from chaosaws.types import AWSResponse + +__all__ = ["reboot_msk_broker", "delete_cluster"] + +logger = get_logger() + + +def reboot_msk_broker( + cluster_arn: str, + broker_ids: List[str], + configuration: Configuration = None, + secrets: Secrets = None, +) -> AWSResponse: + """ + Reboot the specified brokers in an MSK cluster. + """ + client = aws_client("kafka", configuration, secrets) + logger.debug( + f"Rebooting MSK brokers: {broker_ids} in cluster {cluster_arn}" + ) + try: + return client.reboot_broker( + ClusterArn=cluster_arn, + BrokerIds=broker_ids + ) + except client.exceptions.NotFoundException as e: + raise FailedActivity("The specified cluster was not found") from e + + +def delete_cluster( + cluster_arn: str, + configuration: Configuration = None, + secrets: Secrets = None, +) -> AWSResponse: + """ + Delete the given MSK cluster. + """ + client = aws_client("kafka", configuration, secrets) + logger.debug(f"Deleting MSK cluster: {cluster_arn}") + try: + return client.delete_cluster(ClusterArn=cluster_arn) + except client.exceptions.NotFoundException as e: + raise FailedActivity("The specified cluster was not found") from e \ No newline at end of file diff --git a/chaosaws/msk/probes.py b/chaosaws/msk/probes.py new file mode 100644 index 00000000..407bc73b --- /dev/null +++ b/chaosaws/msk/probes.py @@ -0,0 +1,43 @@ +from typing import List + +from chaoslib.types import Configuration, Secrets +from chaoslib.exceptions import FailedActivity +from chaosaws import aws_client, get_logger +from chaosaws.types import AWSResponse + +__all__ = ["describe_msk_cluster", "get_bootstrap_servers"] + +logger = get_logger() + + +def describe_msk_cluster( + cluster_arn: str, + configuration: Configuration = None, + secrets: Secrets = None, +) -> AWSResponse: + """ + Describe an MSK cluster. + """ + client = aws_client("kafka", configuration, secrets) + logger.debug(f"Describing MSK cluster: {cluster_arn}") + try: + return client.describe_cluster(ClusterArn=cluster_arn) + except client.exceptions.NotFoundException as e: + raise FailedActivity("The specified cluster was not found") from e + + +def get_bootstrap_servers( + cluster_arn: str, + configuration: Configuration = None, + secrets: Secrets = None, +) -> List[str]: + """ + Get the bootstrap servers for an MSK cluster. + """ + client = aws_client("kafka", configuration, secrets) + logger.debug(f"Getting bootstrap servers for MSK cluster: {cluster_arn}") + try: + response = client.get_bootstrap_brokers(ClusterArn=cluster_arn) + return response["BootstrapBrokerString"].split(",") if response else [] + except client.exceptions.NotFoundException as e: + raise FailedActivity("The specified cluster was not found") from e diff --git a/tests/msk/test_msk_actions.py b/tests/msk/test_msk_actions.py new file mode 100644 index 00000000..ffdf1fa7 --- /dev/null +++ b/tests/msk/test_msk_actions.py @@ -0,0 +1,91 @@ +from unittest.mock import MagicMock, patch +import pytest +from chaosaws.msk.actions import reboot_msk_broker, delete_cluster +from chaoslib.exceptions import FailedActivity + + +class NotFoundException(Exception): + def __init__(self, message="Cluster not found"): + super().__init__(f"{message}") + + +@patch("chaosaws.msk.actions.aws_client", autospec=True) +def test_reboot_msk_broker_success(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + broker_ids = ["1"] + client.reboot_broker.return_value = { + "ResponseMetadata": { + "HTTPStatusCode": 200 + } + } + + response = reboot_msk_broker(cluster_arn=cluster_arn, + broker_ids=broker_ids) + + client.reboot_broker.assert_called_with(ClusterArn=cluster_arn, + BrokerIds=broker_ids) + assert response["ResponseMetadata"]["HTTPStatusCode"] == 200 + + +@patch("chaosaws.msk.actions.aws_client", autospec=True) +def test_reboot_msk_broker_not_found(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + broker_ids = ["1"] + + client.exceptions = MagicMock() + client.exceptions.NotFoundException = NotFoundException + client.reboot_broker.side_effect = NotFoundException( + "Cluster not found" + ) + + expected_error_message = "The specified cluster was not found" + + with pytest.raises(FailedActivity) as exc_info: + reboot_msk_broker(cluster_arn=cluster_arn, broker_ids=broker_ids) + + assert expected_error_message in str( + exc_info.value + ) + + +@patch("chaosaws.msk.actions.aws_client", autospec=True) +def test_delete_cluster_success(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + client.delete_cluster.return_value = { + "ResponseMetadata": { + "HTTPStatusCode": 200 + } + } + + response = delete_cluster(cluster_arn=cluster_arn) + + client.delete_cluster.assert_called_with(ClusterArn=cluster_arn) + assert response["ResponseMetadata"]["HTTPStatusCode"] == 200 + + +@patch("chaosaws.msk.actions.aws_client", autospec=True) +def test_delete_cluster_not_found(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + + client.exceptions = MagicMock() + client.exceptions.NotFoundException = NotFoundException + client.delete_cluster.side_effect = NotFoundException( + "Cluster not found" + ) + + expected_error_message = "The specified cluster was not found" + + with pytest.raises(FailedActivity) as exc_info: + delete_cluster(cluster_arn=cluster_arn) + + assert expected_error_message in str( + exc_info.value + ) \ No newline at end of file diff --git a/tests/msk/test_msk_probes.py b/tests/msk/test_msk_probes.py new file mode 100644 index 00000000..e56c3d70 --- /dev/null +++ b/tests/msk/test_msk_probes.py @@ -0,0 +1,98 @@ +from unittest.mock import MagicMock, patch +import pytest +from chaosaws.msk.probes import describe_msk_cluster, get_bootstrap_servers +from chaoslib.exceptions import FailedActivity + + +class NotFoundException(Exception): + def __init__(self, message="Cluster not found"): + super().__init__(f"{message}") + + +@patch("chaosaws.msk.probes.aws_client", autospec=True) +def test_describe_msk_cluster_success(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + expected_response = {"ClusterInfo": {"ClusterArn": cluster_arn}} + + client.describe_cluster.return_value = expected_response + + response = describe_msk_cluster(cluster_arn=cluster_arn) + + client.describe_cluster.assert_called_with(ClusterArn=cluster_arn) + assert response == expected_response + + +@patch("chaosaws.msk.probes.aws_client", autospec=True) +def test_describe_msk_cluster_not_found(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + + client.exceptions = MagicMock() + client.exceptions.NotFoundException = NotFoundException + client.describe_cluster.side_effect = NotFoundException( + "Cluster not found" + ) + + expected_error_message = "The specified cluster was not found" + + with pytest.raises(FailedActivity) as exc_info: + describe_msk_cluster(cluster_arn=cluster_arn) + + assert expected_error_message in str( + exc_info.value + ) + + +@patch("chaosaws.msk.probes.aws_client", autospec=True) +def test_get_bootstrap_servers_success(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + bootstrap_servers = "broker1,broker2,broker3" + expected_response = {"BootstrapBrokerString": bootstrap_servers} + + client.get_bootstrap_brokers.return_value = expected_response + + response = get_bootstrap_servers(cluster_arn=cluster_arn) + + client.get_bootstrap_brokers.assert_called_with(ClusterArn=cluster_arn) + assert response == bootstrap_servers.split(",") + + +@patch("chaosaws.msk.probes.aws_client", autospec=True) +def test_get_bootstrap_servers_empty_response(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + + client.get_bootstrap_brokers.return_value = None + + response = get_bootstrap_servers(cluster_arn=cluster_arn) + + client.get_bootstrap_brokers.assert_called_with(ClusterArn=cluster_arn) + assert response == [] + + +@patch("chaosaws.msk.probes.aws_client", autospec=True) +def test_get_bootstrap_server_cluster_not_found(aws_client): + client = MagicMock() + aws_client.return_value = client + cluster_arn = "arn_msk_cluster" + + client.exceptions = MagicMock() + client.exceptions.NotFoundException = NotFoundException + client.get_bootstrap_brokers.side_effect = NotFoundException( + "Cluster not found" + ) + + expected_error_message = "The specified cluster was not found" + + with pytest.raises(FailedActivity) as exc_info: + get_bootstrap_servers(cluster_arn=cluster_arn) + + assert expected_error_message in str( + exc_info.value + ) From 780f01f8ab03f5a43dccc3e3c61f2ab97cc190bc Mon Sep 17 00:00:00 2001 From: Jorge Tapicha Date: Thu, 30 May 2024 00:10:19 -0500 Subject: [PATCH 2/5] add changelog and activities extends Signed-off-by: Jorge Tapicha Signed-off-by: Jorge Tapicha --- CHANGELOG.md | 5 +++++ chaosaws/__init__.py | 2 ++ 2 files changed, 7 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 049e6264..a4160f6b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # Changelog +### Added + +* MSK probes `chaosaws.msk.probes.describe_msk_cluster` `chaosaws.msk.probes.get_bootstrap_servers` +* MSK actions `chaosaws.msk.actions.reboot_msk_broker` `chaosaws.msk.actions.delete_cluster` + ## [Unreleased][] [Unreleased]: https://github.com/chaostoolkit-incubator/chaostoolkit-aws/compare/0.33.0...HEAD diff --git a/chaosaws/__init__.py b/chaosaws/__init__.py index c57e82d7..3e13d6e9 100644 --- a/chaosaws/__init__.py +++ b/chaosaws/__init__.py @@ -307,6 +307,8 @@ def load_exported_activities() -> List[DiscoveredActivities]: activities.extend(discover_actions("chaosaws.fis.actions")) activities.extend(discover_probes("chaosaws.s3.probes")) activities.extend(discover_actions("chaosaws.s3.actions")) + activities.extend(discover_probes("chaosaws.msk.probes")) + activities.extend(discover_actions("chaosaws.msk.actions")) activities.extend( discover_activities("chaosaws.s3.controls.upload", "control") ) From cae858c0c4a0cb47e3490ad6e400a9845a8cb581 Mon Sep 17 00:00:00 2001 From: Jorge Tapicha Date: Thu, 30 May 2024 08:58:29 -0500 Subject: [PATCH 3/5] remove unnecesary empty response test Signed-off-by: Jorge Tapicha --- tests/msk/test_msk_probes.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tests/msk/test_msk_probes.py b/tests/msk/test_msk_probes.py index e56c3d70..8ed3e4c6 100644 --- a/tests/msk/test_msk_probes.py +++ b/tests/msk/test_msk_probes.py @@ -62,20 +62,6 @@ def test_get_bootstrap_servers_success(aws_client): assert response == bootstrap_servers.split(",") -@patch("chaosaws.msk.probes.aws_client", autospec=True) -def test_get_bootstrap_servers_empty_response(aws_client): - client = MagicMock() - aws_client.return_value = client - cluster_arn = "arn_msk_cluster" - - client.get_bootstrap_brokers.return_value = None - - response = get_bootstrap_servers(cluster_arn=cluster_arn) - - client.get_bootstrap_brokers.assert_called_with(ClusterArn=cluster_arn) - assert response == [] - - @patch("chaosaws.msk.probes.aws_client", autospec=True) def test_get_bootstrap_server_cluster_not_found(aws_client): client = MagicMock() From 9b460dd0559200d56b632252ef34ac99315fb0ed Mon Sep 17 00:00:00 2001 From: Jorge Tapicha Date: Fri, 31 May 2024 11:49:04 -0500 Subject: [PATCH 4/5] first corrections Pull Request Signed-off-by: Jorge Tapicha --- CHANGELOG.md | 4 ++-- chaosaws/msk/actions.py | 9 +++++---- chaosaws/msk/probes.py | 9 +++++---- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a4160f6b..f459bdb2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,12 +1,12 @@ # Changelog +## [Unreleased][] + ### Added * MSK probes `chaosaws.msk.probes.describe_msk_cluster` `chaosaws.msk.probes.get_bootstrap_servers` * MSK actions `chaosaws.msk.actions.reboot_msk_broker` `chaosaws.msk.actions.delete_cluster` -## [Unreleased][] - [Unreleased]: https://github.com/chaostoolkit-incubator/chaostoolkit-aws/compare/0.33.0...HEAD ## [0.33.0][] - 2024-02-26 diff --git a/chaosaws/msk/actions.py b/chaosaws/msk/actions.py index 2443fbfd..a6432b56 100644 --- a/chaosaws/msk/actions.py +++ b/chaosaws/msk/actions.py @@ -2,6 +2,7 @@ from chaoslib.types import Configuration, Secrets from chaoslib.exceptions import FailedActivity + from chaosaws import aws_client, get_logger from chaosaws.types import AWSResponse @@ -28,8 +29,8 @@ def reboot_msk_broker( ClusterArn=cluster_arn, BrokerIds=broker_ids ) - except client.exceptions.NotFoundException as e: - raise FailedActivity("The specified cluster was not found") from e + except client.exceptions.NotFoundException: + raise FailedActivity(f"The specified cluster was not found" ) def delete_cluster( @@ -44,5 +45,5 @@ def delete_cluster( logger.debug(f"Deleting MSK cluster: {cluster_arn}") try: return client.delete_cluster(ClusterArn=cluster_arn) - except client.exceptions.NotFoundException as e: - raise FailedActivity("The specified cluster was not found") from e \ No newline at end of file + except client.exceptions.NotFoundException: + raise FailedActivity(f"The specified cluster was not found") \ No newline at end of file diff --git a/chaosaws/msk/probes.py b/chaosaws/msk/probes.py index 407bc73b..d1b3d7dd 100644 --- a/chaosaws/msk/probes.py +++ b/chaosaws/msk/probes.py @@ -2,6 +2,7 @@ from chaoslib.types import Configuration, Secrets from chaoslib.exceptions import FailedActivity + from chaosaws import aws_client, get_logger from chaosaws.types import AWSResponse @@ -22,8 +23,8 @@ def describe_msk_cluster( logger.debug(f"Describing MSK cluster: {cluster_arn}") try: return client.describe_cluster(ClusterArn=cluster_arn) - except client.exceptions.NotFoundException as e: - raise FailedActivity("The specified cluster was not found") from e + except client.exceptions.NotFoundException: + raise FailedActivity("The specified cluster was not found") def get_bootstrap_servers( @@ -39,5 +40,5 @@ def get_bootstrap_servers( try: response = client.get_bootstrap_brokers(ClusterArn=cluster_arn) return response["BootstrapBrokerString"].split(",") if response else [] - except client.exceptions.NotFoundException as e: - raise FailedActivity("The specified cluster was not found") from e + except client.exceptions.NotFoundException: + raise FailedActivity("The specified cluster was not found") From 70eae9d3a2f9c89ff5474432c1b1b393b6d400a1 Mon Sep 17 00:00:00 2001 From: Jorge Tapicha Date: Tue, 4 Jun 2024 18:46:43 -0500 Subject: [PATCH 5/5] fix f-string without any placeholders Signed-off-by: Jorge Tapicha --- chaosaws/msk/actions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/chaosaws/msk/actions.py b/chaosaws/msk/actions.py index a6432b56..a3647818 100644 --- a/chaosaws/msk/actions.py +++ b/chaosaws/msk/actions.py @@ -30,7 +30,7 @@ def reboot_msk_broker( BrokerIds=broker_ids ) except client.exceptions.NotFoundException: - raise FailedActivity(f"The specified cluster was not found" ) + raise FailedActivity("The specified cluster was not found" ) def delete_cluster( @@ -46,4 +46,4 @@ def delete_cluster( try: return client.delete_cluster(ClusterArn=cluster_arn) except client.exceptions.NotFoundException: - raise FailedActivity(f"The specified cluster was not found") \ No newline at end of file + raise FailedActivity("The specified cluster was not found") \ No newline at end of file