Skip to content

Commit

Permalink
Fix lint
Browse files Browse the repository at this point in the history
  • Loading branch information
@NathanielRN
NathanielRN committed Nov 7, 2020
1 parent 591a21f commit 27fbdaa
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 52 deletions.
32 changes: 20 additions & 12 deletions ...telemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,14 +105,10 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
span.set_attribute("aws.region", instance.meta.region_name)
span.set_attribute("aws.service", service_name)
if "QueueUrl" in api_params:
span.set_attribute(
"aws.queue_url",
api_params["QueueUrl"]
)
span.set_attribute("aws.queue_url", api_params["QueueUrl"])
if "TableName" in api_params:
span.set_attribute(
"aws.table_name",
api_params["TableName"]
"aws.table_name", api_params["TableName"]
)

try:
Expand All @@ -125,15 +121,27 @@ def _patched_api_call(self, original_func, instance, args, kwargs):

if span.is_recording():
if "ResponseMetadata" in result:
if "RequestId" in result["ResponseMetadata"]:
metadata = result["ResponseMetadata"]
req_id = None
if "RequestId" in metadata:
req_id = metadata["RequestId"]
elif "HTTPHeaders" in metadata:
headers = metadata["HTTPHeaders"]
if "x-amzn-RequestId" in headers:
req_id = headers["x-amzn-RequestId"]
elif "x-amz-request-id" in headers:
req_id = headers["x-amz-request-id"]
elif "x-amz-id-2" in headers:
req_id = headers["x-amz-id-2"]

if req_id:
span.set_attribute(
"aws.request_id",
result["ResponseMetadata"]["RequestId"],
"aws.request_id", req_id,
)
if "HTTPStatusCode" in result["ResponseMetadata"]:

if "HTTPStatusCode" in metadata:
span.set_attribute(
"http.status_code",
result["ResponseMetadata"]["HTTPStatusCode"],
"http.status_code", metadata["HTTPStatusCode"],
)

if error:
Expand Down
65 changes: 25 additions & 40 deletions ...rumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,8 @@
import botocore.session
from botocore.exceptions import ParamValidationError
from moto import ( # pylint: disable=import-error
mock_ec2,
mock_dynamodb2,
mock_ec2,
mock_kinesis,
mock_kms,
mock_lambda,
Expand Down Expand Up @@ -203,22 +203,21 @@ def test_sqs_client(self):
def test_sqs_send_message(self):
sqs = self.session.create_client("sqs", region_name="us-east-1")

TEST_QUEUE_NAME = 'test_queue_name'
TEST_QUEUE_NAME = "test_queue_name"

response = sqs.create_queue(
QueueName=TEST_QUEUE_NAME
)
response = sqs.create_queue(QueueName=TEST_QUEUE_NAME)

sqs.send_message(
QueueUrl=response["QueueUrl"],
MessageBody="Test SQS MESSAGE!"
QueueUrl=response["QueueUrl"], MessageBody="Test SQS MESSAGE!"
)

spans = self.memory_exporter.get_finished_spans()
assert spans
self.assertEqual(len(spans), 2)
create_queue_attributes = dict(spans[0].attributes)
self.assertRegex(create_queue_attributes["aws.request_id"], r"[A-Z0-9]{52}")
self.assertRegex(
create_queue_attributes["aws.request_id"], r"[A-Z0-9]{52}"
)
del create_queue_attributes["aws.request_id"]
self.assertEqual(
create_queue_attributes,
Expand All @@ -230,7 +229,9 @@ def test_sqs_send_message(self):
},
)
send_msg_attributes = dict(spans[1].attributes)
self.assertRegex(send_msg_attributes["aws.request_id"], r"[A-Z0-9]{52}")
self.assertRegex(
send_msg_attributes["aws.request_id"], r"[A-Z0-9]{52}"
)
del send_msg_attributes["aws.request_id"]
self.assertEqual(
send_msg_attributes,
Expand Down Expand Up @@ -365,47 +366,27 @@ def test_dynamodb_client(self):

ddb.create_table(
AttributeDefinitions=[
{
'AttributeName': 'id',
'AttributeType': 'S',
},
],
KeySchema=[
{
'AttributeName': 'id',
'KeyType': 'HASH',
},
{"AttributeName": "id", "AttributeType": "S"},
],
KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}],
ProvisionedThroughput={
'ReadCapacityUnits': 5,
'WriteCapacityUnits': 5,
"ReadCapacityUnits": 5,
"WriteCapacityUnits": 5,
},
TableName=TEST_TABLE_NAME,
)

ddb.put_item(
TableName=TEST_TABLE_NAME,
Item={
'id': {
'S': 'test_key'
}
}
)
ddb.put_item(TableName=TEST_TABLE_NAME, Item={"id": {"S": "test_key"}})

ddb.get_item(
TableName=TEST_TABLE_NAME,
Key={
'id': {
'S': 'test_key'
}
}
)
ddb.get_item(TableName=TEST_TABLE_NAME, Key={"id": {"S": "test_key"}})

spans = self.memory_exporter.get_finished_spans()
assert spans
self.assertEqual(len(spans), 3)
create_table_attributes = dict(spans[0].attributes)
self.assertRegex(create_table_attributes["aws.request_id"], r"[A-Z0-9]{52}")
self.assertRegex(
create_table_attributes["aws.request_id"], r"[A-Z0-9]{52}"
)
del create_table_attributes["aws.request_id"]
self.assertEqual(
create_table_attributes,
Expand All @@ -418,7 +399,9 @@ def test_dynamodb_client(self):
},
)
put_item_attributes = dict(spans[1].attributes)
self.assertRegex(put_item_attributes["aws.request_id"], r"[A-Z0-9]{52}")
self.assertRegex(
put_item_attributes["aws.request_id"], r"[A-Z0-9]{52}"
)
del put_item_attributes["aws.request_id"]
self.assertEqual(
put_item_attributes,
Expand All @@ -431,7 +414,9 @@ def test_dynamodb_client(self):
},
)
get_item_attributes = dict(spans[2].attributes)
self.assertRegex(get_item_attributes["aws.request_id"], r"[A-Z0-9]{52}")
self.assertRegex(
get_item_attributes["aws.request_id"], r"[A-Z0-9]{52}"
)
del get_item_attributes["aws.request_id"]
self.assertEqual(
get_item_attributes,
Expand Down

0 comments on commit 27fbdaa

Please sign in to comment.