diff --git a/python/example_code/ssm/document.py b/python/example_code/ssm/document.py index b1799fa7722..b498d6f992c 100644 --- a/python/example_code/ssm/document.py +++ b/python/example_code/ssm/document.py @@ -42,6 +42,9 @@ def create(self, content, name): Name=name, Content=content, DocumentType="Command" ) self.name = name + except self.ssm_client.exceptions.DocumentAlreadyExists: + logger.warning("Document %s already exists. ", name) + self.name = name except ClientError as err: logger.error( "Couldn't create %s. Here's why: %s: %s", @@ -64,6 +67,14 @@ def delete(self): try: self.ssm_client.delete_document(Name=self.name) self.name = None + except self.ssm_client.exceptions.InvalidDocument as err: + logger.error( + "Document %s is not valid. %s: %s", + self.name, + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + raise except ClientError as err: logger.error( "Couldn't delete %s. Here's why: %s: %s", @@ -88,6 +99,14 @@ def send_command(self, instance_ids): InstanceIds=instance_ids, DocumentName=self.name, TimeoutSeconds=3600 ) return response["Command"]["CommandId"] + except self.ssm_client.exceptions.InvalidDocument as err: + logger.error( + "Document %s is not valid. %s: %s", + self.name, + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + raise except ClientError as err: logger.error( "Couldn't send command to %s. Here's why: %s: %s", @@ -176,6 +195,14 @@ def list_commands(self, instance_id): print( f" The time of command invocation is {command['RequestedDateTime'].strftime(date_format)}" ) + except self.ssm_client.exceptions.InvalidInstanceId as err: + logger.error( + "Instance ID %s is not valid. %s: %s", + instance_id, + err.response["Error"]["Code"], + err.response["Error"]["Message"], + ) + raise except ClientError as err: logger.error( "Couldn't list commands for %s. Here's why: %s: %s", diff --git a/python/example_code/ssm/test/test_document.py b/python/example_code/ssm/test/test_document.py index c6f6471d6c5..4bdc401551c 100644 --- a/python/example_code/ssm/test/test_document.py +++ b/python/example_code/ssm/test/test_document.py @@ -10,9 +10,18 @@ from botocore.exceptions import ClientError from document import DocumentWrapper +from logging.handlers import QueueHandler +import logging +import queue +log_queue = queue.LifoQueue() -@pytest.mark.parametrize("error_code", [None, "TestException"]) +queue_handler = QueueHandler(log_queue) +document_logger = logging.getLogger("document") +document_logger.addHandler(queue_handler) + + +@pytest.mark.parametrize("error_code", [None, "TestException", "DocumentAlreadyExists"]) def test_create(make_stubber, error_code): ssm_client = boto3.client("ssm") ssm_stubber = make_stubber(ssm_client) @@ -48,6 +57,15 @@ def test_create(make_stubber, error_code): name, ) assert document_wrapper.name == name + elif error_code == "DocumentAlreadyExists": + document_wrapper.create( + content, + name, + ) + assert document_wrapper.name == name + assert ( + log_queue.qsize() > 0 and "already exists" in log_queue.get().getMessage() + ) else: with pytest.raises(ClientError) as exc_info: document_wrapper.create( @@ -57,7 +75,7 @@ def test_create(make_stubber, error_code): assert exc_info.value.response["Error"]["Code"] == error_code -@pytest.mark.parametrize("error_code", [None, "TestException"]) +@pytest.mark.parametrize("error_code", [None, "InvalidDocument", "TestException"]) def test_delete(make_stubber, error_code): ssm_client = boto3.client("ssm") ssm_stubber = make_stubber(ssm_client) @@ -73,13 +91,17 @@ def test_delete(make_stubber, error_code): if error_code is None: document_wrapper.delete() assert document_wrapper.name is None + elif error_code == "InvalidDocument": + with pytest.raises(ssm_client.exceptions.InvalidDocument) as exc_info: + document_wrapper.delete() + assert exc_info.value.response["Error"]["Code"] == error_code else: with pytest.raises(ClientError) as exc_info: document_wrapper.delete() assert exc_info.value.response["Error"]["Code"] == error_code -@pytest.mark.parametrize("error_code", [None, "TestException"]) +@pytest.mark.parametrize("error_code", [None, "InvalidDocument", "TestException"]) def test_describe_document(make_stubber, error_code): ssm_client = boto3.client("ssm") ssm_stubber = make_stubber(ssm_client) @@ -95,13 +117,17 @@ def test_describe_document(make_stubber, error_code): if error_code is None: status = document_wrapper.describe() assert status == "Active" + elif error_code == "InvalidDocument": + with pytest.raises(ssm_client.exceptions.InvalidDocument) as exc_info: + document_wrapper.describe() + assert exc_info.value.response["Error"]["Code"] == error_code else: with pytest.raises(ClientError) as exc_info: document_wrapper.describe() assert exc_info.value.response["Error"]["Code"] == error_code -@pytest.mark.parametrize("error_code", [None, "TestException"]) +@pytest.mark.parametrize("error_code", [None, "InvalidInstanceId", "TestException"]) def test_list_commands(make_stubber, error_code): ssm_client = boto3.client("ssm") ssm_stubber = make_stubber(ssm_client) @@ -116,13 +142,18 @@ def test_list_commands(make_stubber, error_code): if error_code is None: document_wrapper.list_commands(instance_id) + elif error_code == "InvalidInstanceId": + with pytest.raises(ssm_client.exceptions.InvalidInstanceId) as exc_info: + document_wrapper.list_commands(instance_id) + + assert exc_info.value.response["Error"]["Code"] == error_code else: with pytest.raises(ClientError) as exc_info: document_wrapper.list_commands(instance_id) assert exc_info.value.response["Error"]["Code"] == error_code -@pytest.mark.parametrize("error_code", [None, "TestException"]) +@pytest.mark.parametrize("error_code", [None, "InvalidDocument", "TestException"]) def test_send_command(make_stubber, error_code): ssm_client = boto3.client("ssm") ssm_stubber = make_stubber(ssm_client) @@ -142,6 +173,10 @@ def test_send_command(make_stubber, error_code): if error_code is None: response = document_wrapper.send_command(instance_ids) assert response == command_id + elif error_code == "InvalidDocument": + with pytest.raises(ssm_client.exceptions.InvalidDocument) as exc_info: + document_wrapper.send_command(instance_ids) + assert exc_info.value.response["Error"]["Code"] == error_code else: with pytest.raises(ClientError) as exc_info: document_wrapper.send_command(instance_ids)