Skip to content

Commit

Permalink
adding extra exception handling
Browse files Browse the repository at this point in the history
  • Loading branch information
@meyertst-aws
meyertst-aws committed Jun 25, 2024
1 parent 13eb06a commit 23c6969
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 5 deletions.
27 changes: 27 additions & 0 deletions python/example_code/ssm/document.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
45 changes: 40 additions & 5 deletions python/example_code/ssm/test/test_document.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down

0 comments on commit 23c6969

Please sign in to comment.