Run black on moto & test directories.

This commit is contained in:
Asher Foa 2019-10-31 08:44:26 -07:00
commit 96e5b1993d
507 changed files with 52541 additions and 47814 deletions

View file

@ -3,5 +3,5 @@ from .models import xray_backends
from ..core.models import base_decorator
from .mock_client import mock_xray_client, XRaySegment # noqa
xray_backend = xray_backends['us-east-1']
xray_backend = xray_backends["us-east-1"]
mock_xray = base_decorator(xray_backends)

View file

@ -11,11 +11,14 @@ class AWSError(Exception):
self.status = status if status is not None else self.STATUS
def response(self):
return json.dumps({'__type': self.code, 'message': self.message}), dict(status=self.status)
return (
json.dumps({"__type": self.code, "message": self.message}),
dict(status=self.status),
)
class InvalidRequestException(AWSError):
CODE = 'InvalidRequestException'
CODE = "InvalidRequestException"
class BadSegmentException(Exception):
@ -25,15 +28,15 @@ class BadSegmentException(Exception):
self.message = message
def __repr__(self):
return '<BadSegment {0}>'.format('-'.join([self.id, self.code, self.message]))
return "<BadSegment {0}>".format("-".join([self.id, self.code, self.message]))
def to_dict(self):
result = {}
if self.id is not None:
result['Id'] = self.id
result["Id"] = self.id
if self.code is not None:
result['ErrorCode'] = self.code
result["ErrorCode"] = self.code
if self.message is not None:
result['Message'] = self.message
result["Message"] = self.message
return result

View file

@ -10,8 +10,11 @@ class MockEmitter(UDPEmitter):
"""
Replaces the code that sends UDP to local X-Ray daemon
"""
def __init__(self, daemon_address='127.0.0.1:2000'):
address = os.getenv('AWS_XRAY_DAEMON_ADDRESS_YEAH_NOT_TODAY_MATE', daemon_address)
def __init__(self, daemon_address="127.0.0.1:2000"):
address = os.getenv(
"AWS_XRAY_DAEMON_ADDRESS_YEAH_NOT_TODAY_MATE", daemon_address
)
self._ip, self._port = self._parse_address(address)
def _xray_backend(self, region):
@ -26,7 +29,7 @@ class MockEmitter(UDPEmitter):
pass
def _send_data(self, data):
raise RuntimeError('Should not be running this')
raise RuntimeError("Should not be running this")
def mock_xray_client(f):
@ -39,12 +42,13 @@ def mock_xray_client(f):
We also patch the Emitter by subclassing the UDPEmitter class replacing its methods and pushing
that itno the recorder instance.
"""
@wraps(f)
def _wrapped(*args, **kwargs):
print("Starting X-Ray Patch")
old_xray_context_var = os.environ.get('AWS_XRAY_CONTEXT_MISSING')
os.environ['AWS_XRAY_CONTEXT_MISSING'] = 'LOG_ERROR'
old_xray_context_var = os.environ.get("AWS_XRAY_CONTEXT_MISSING")
os.environ["AWS_XRAY_CONTEXT_MISSING"] = "LOG_ERROR"
old_xray_context = aws_xray_sdk.core.xray_recorder._context
old_xray_emitter = aws_xray_sdk.core.xray_recorder._emitter
aws_xray_sdk.core.xray_recorder._context = AWSContext()
@ -55,9 +59,9 @@ def mock_xray_client(f):
finally:
if old_xray_context_var is None:
del os.environ['AWS_XRAY_CONTEXT_MISSING']
del os.environ["AWS_XRAY_CONTEXT_MISSING"]
else:
os.environ['AWS_XRAY_CONTEXT_MISSING'] = old_xray_context_var
os.environ["AWS_XRAY_CONTEXT_MISSING"] = old_xray_context_var
aws_xray_sdk.core.xray_recorder._emitter = old_xray_emitter
aws_xray_sdk.core.xray_recorder._context = old_xray_context
@ -74,8 +78,11 @@ class XRaySegment(object):
During testing we're going to have to control the start and end of a segment via context managers.
"""
def __enter__(self):
aws_xray_sdk.core.xray_recorder.begin_segment(name='moto_mock', traceid=None, parent_id=None, sampling=1)
aws_xray_sdk.core.xray_recorder.begin_segment(
name="moto_mock", traceid=None, parent_id=None, sampling=1
)
return self

View file

@ -18,18 +18,36 @@ class TelemetryRecords(BaseModel):
@classmethod
def from_json(cls, json):
instance_id = json.get('EC2InstanceId', None)
hostname = json.get('Hostname')
resource_arn = json.get('ResourceARN')
telemetry_records = json['TelemetryRecords']
instance_id = json.get("EC2InstanceId", None)
hostname = json.get("Hostname")
resource_arn = json.get("ResourceARN")
telemetry_records = json["TelemetryRecords"]
return cls(instance_id, hostname, resource_arn, telemetry_records)
# https://docs.aws.amazon.com/xray/latest/devguide/xray-api-segmentdocuments.html
class TraceSegment(BaseModel):
def __init__(self, name, segment_id, trace_id, start_time, raw, end_time=None, in_progress=False, service=None, user=None,
origin=None, parent_id=None, http=None, aws=None, metadata=None, annotations=None, subsegments=None, **kwargs):
def __init__(
self,
name,
segment_id,
trace_id,
start_time,
raw,
end_time=None,
in_progress=False,
service=None,
user=None,
origin=None,
parent_id=None,
http=None,
aws=None,
metadata=None,
annotations=None,
subsegments=None,
**kwargs
):
self.name = name
self.id = segment_id
self.trace_id = trace_id
@ -61,14 +79,16 @@ class TraceSegment(BaseModel):
@property
def trace_version(self):
if self._trace_version is None:
self._trace_version = int(self.trace_id.split('-', 1)[0])
self._trace_version = int(self.trace_id.split("-", 1)[0])
return self._trace_version
@property
def request_start_date(self):
if self._original_request_start_time is None:
start_time = int(self.trace_id.split('-')[1], 16)
self._original_request_start_time = datetime.datetime.fromtimestamp(start_time)
start_time = int(self.trace_id.split("-")[1], 16)
self._original_request_start_time = datetime.datetime.fromtimestamp(
start_time
)
return self._original_request_start_time
@property
@ -86,19 +106,27 @@ class TraceSegment(BaseModel):
@classmethod
def from_dict(cls, data, raw):
# Check manditory args
if 'id' not in data:
raise BadSegmentException(code='MissingParam', message='Missing segment ID')
seg_id = data['id']
data['segment_id'] = seg_id # Just adding this key for future convenience
if "id" not in data:
raise BadSegmentException(code="MissingParam", message="Missing segment ID")
seg_id = data["id"]
data["segment_id"] = seg_id # Just adding this key for future convenience
for arg in ('name', 'trace_id', 'start_time'):
for arg in ("name", "trace_id", "start_time"):
if arg not in data:
raise BadSegmentException(seg_id=seg_id, code='MissingParam', message='Missing segment ID')
raise BadSegmentException(
seg_id=seg_id, code="MissingParam", message="Missing segment ID"
)
if 'end_time' not in data and 'in_progress' not in data:
raise BadSegmentException(seg_id=seg_id, code='MissingParam', message='Missing end_time or in_progress')
if 'end_time' not in data and data['in_progress'] == 'false':
raise BadSegmentException(seg_id=seg_id, code='MissingParam', message='Missing end_time')
if "end_time" not in data and "in_progress" not in data:
raise BadSegmentException(
seg_id=seg_id,
code="MissingParam",
message="Missing end_time or in_progress",
)
if "end_time" not in data and data["in_progress"] == "false":
raise BadSegmentException(
seg_id=seg_id, code="MissingParam", message="Missing end_time"
)
return cls(raw=raw, **data)
@ -110,65 +138,79 @@ class SegmentCollection(object):
@staticmethod
def _new_trace_item():
return {
'start_date': datetime.datetime(1970, 1, 1),
'end_date': datetime.datetime(1970, 1, 1),
'finished': False,
'trace_id': None,
'segments': []
"start_date": datetime.datetime(1970, 1, 1),
"end_date": datetime.datetime(1970, 1, 1),
"finished": False,
"trace_id": None,
"segments": [],
}
def put_segment(self, segment):
# insert into a sorted list
bisect.insort_left(self._traces[segment.trace_id]['segments'], segment)
bisect.insort_left(self._traces[segment.trace_id]["segments"], segment)
# Get the last segment (takes into account incorrect ordering)
# and if its the last one, mark trace as complete
if self._traces[segment.trace_id]['segments'][-1].end_time is not None:
self._traces[segment.trace_id]['finished'] = True
if self._traces[segment.trace_id]["segments"][-1].end_time is not None:
self._traces[segment.trace_id]["finished"] = True
start_time = self._traces[segment.trace_id]['segments'][0].start_date
end_time = self._traces[segment.trace_id]['segments'][-1].end_date
self._traces[segment.trace_id]['start_date'] = start_time
self._traces[segment.trace_id]['end_date'] = end_time
self._traces[segment.trace_id]['trace_id'] = segment.trace_id
start_time = self._traces[segment.trace_id]["segments"][0].start_date
end_time = self._traces[segment.trace_id]["segments"][-1].end_date
self._traces[segment.trace_id]["start_date"] = start_time
self._traces[segment.trace_id]["end_date"] = end_time
self._traces[segment.trace_id]["trace_id"] = segment.trace_id
# Todo consolidate trace segments into a trace.
# not enough working knowledge of xray to do this
def summary(self, start_time, end_time, filter_expression=None, sampling=False):
# This beast https://docs.aws.amazon.com/xray/latest/api/API_GetTraceSummaries.html#API_GetTraceSummaries_ResponseSyntax
if filter_expression is not None:
raise AWSError('Not implemented yet - moto', code='InternalFailure', status=500)
raise AWSError(
"Not implemented yet - moto", code="InternalFailure", status=500
)
summaries = []
for tid, trace in self._traces.items():
if trace['finished'] and start_time < trace['start_date'] and trace['end_date'] < end_time:
duration = int((trace['end_date'] - trace['start_date']).total_seconds())
if (
trace["finished"]
and start_time < trace["start_date"]
and trace["end_date"] < end_time
):
duration = int(
(trace["end_date"] - trace["start_date"]).total_seconds()
)
# this stuff is mostly guesses, refer to TODO above
has_error = any(['error' in seg.misc for seg in trace['segments']])
has_fault = any(['fault' in seg.misc for seg in trace['segments']])
has_throttle = any(['throttle' in seg.misc for seg in trace['segments']])
has_error = any(["error" in seg.misc for seg in trace["segments"]])
has_fault = any(["fault" in seg.misc for seg in trace["segments"]])
has_throttle = any(
["throttle" in seg.misc for seg in trace["segments"]]
)
# Apparently all of these options are optional
summary_part = {
'Annotations': {}, # Not implemented yet
'Duration': duration,
'HasError': has_error,
'HasFault': has_fault,
'HasThrottle': has_throttle,
'Http': {}, # Not implemented yet
'Id': tid,
'IsParital': False, # needs lots more work to work on partials
'ResponseTime': 1, # definitely 1ms resposnetime
'ServiceIds': [], # Not implemented yet
'Users': {} # Not implemented yet
"Annotations": {}, # Not implemented yet
"Duration": duration,
"HasError": has_error,
"HasFault": has_fault,
"HasThrottle": has_throttle,
"Http": {}, # Not implemented yet
"Id": tid,
"IsParital": False, # needs lots more work to work on partials
"ResponseTime": 1, # definitely 1ms resposnetime
"ServiceIds": [], # Not implemented yet
"Users": {}, # Not implemented yet
}
summaries.append(summary_part)
result = {
"ApproximateTime": int((datetime.datetime.now() - datetime.datetime(1970, 1, 1)).total_seconds()),
"ApproximateTime": int(
(
datetime.datetime.now() - datetime.datetime(1970, 1, 1)
).total_seconds()
),
"TracesProcessedCount": len(summaries),
"TraceSummaries": summaries
"TraceSummaries": summaries,
}
return result
@ -189,59 +231,57 @@ class SegmentCollection(object):
class XRayBackend(BaseBackend):
def __init__(self):
self._telemetry_records = []
self._segment_collection = SegmentCollection()
def add_telemetry_records(self, json):
self._telemetry_records.append(
TelemetryRecords.from_json(json)
)
self._telemetry_records.append(TelemetryRecords.from_json(json))
def process_segment(self, doc):
try:
data = json.loads(doc)
except ValueError:
raise BadSegmentException(code='JSONFormatError', message='Bad JSON data')
raise BadSegmentException(code="JSONFormatError", message="Bad JSON data")
try:
# Get Segment Object
segment = TraceSegment.from_dict(data, raw=doc)
except ValueError:
raise BadSegmentException(code='JSONFormatError', message='Bad JSON data')
raise BadSegmentException(code="JSONFormatError", message="Bad JSON data")
try:
# Store Segment Object
self._segment_collection.put_segment(segment)
except Exception as err:
raise BadSegmentException(seg_id=segment.id, code='InternalFailure', message=str(err))
raise BadSegmentException(
seg_id=segment.id, code="InternalFailure", message=str(err)
)
def get_trace_summary(self, start_time, end_time, filter_expression, summaries):
return self._segment_collection.summary(start_time, end_time, filter_expression, summaries)
return self._segment_collection.summary(
start_time, end_time, filter_expression, summaries
)
def get_trace_ids(self, trace_ids, next_token):
traces, unprocessed_ids = self._segment_collection.get_trace_ids(trace_ids)
result = {
'Traces': [],
'UnprocessedTraceIds': unprocessed_ids
}
result = {"Traces": [], "UnprocessedTraceIds": unprocessed_ids}
for trace in traces:
segments = []
for segment in trace['segments']:
segments.append({
'Id': segment.id,
'Document': segment.raw
})
for segment in trace["segments"]:
segments.append({"Id": segment.id, "Document": segment.raw})
result['Traces'].append({
'Duration': int((trace['end_date'] - trace['start_date']).total_seconds()),
'Id': trace['trace_id'],
'Segments': segments
})
result["Traces"].append(
{
"Duration": int(
(trace["end_date"] - trace["start_date"]).total_seconds()
),
"Id": trace["trace_id"],
"Segments": segments,
}
)
return result

View file

@ -10,9 +10,8 @@ from .exceptions import AWSError, BadSegmentException
class XRayResponse(BaseResponse):
def _error(self, code, message):
return json.dumps({'__type': code, 'message': message}), dict(status=400)
return json.dumps({"__type": code, "message": message}), dict(status=400)
@property
def xray_backend(self):
@ -32,7 +31,7 @@ class XRayResponse(BaseResponse):
# Amazon is just calling urls like /TelemetryRecords etc...
# This uses the value after / as the camalcase action, which then
# gets converted in call_action to find the following methods
return urlsplit(self.uri).path.lstrip('/')
return urlsplit(self.uri).path.lstrip("/")
# PutTelemetryRecords
def telemetry_records(self):
@ -41,15 +40,18 @@ class XRayResponse(BaseResponse):
except AWSError as err:
return err.response()
return ''
return ""
# PutTraceSegments
def trace_segments(self):
docs = self._get_param('TraceSegmentDocuments')
docs = self._get_param("TraceSegmentDocuments")
if docs is None:
msg = 'Parameter TraceSegmentDocuments is missing'
return json.dumps({'__type': 'MissingParameter', 'message': msg}), dict(status=400)
msg = "Parameter TraceSegmentDocuments is missing"
return (
json.dumps({"__type": "MissingParameter", "message": msg}),
dict(status=400),
)
# Raises an exception that contains info about a bad segment,
# the object also has a to_dict() method
@ -60,91 +62,120 @@ class XRayResponse(BaseResponse):
except BadSegmentException as bad_seg:
bad_segments.append(bad_seg)
except Exception as err:
return json.dumps({'__type': 'InternalFailure', 'message': str(err)}), dict(status=500)
return (
json.dumps({"__type": "InternalFailure", "message": str(err)}),
dict(status=500),
)
result = {'UnprocessedTraceSegments': [x.to_dict() for x in bad_segments]}
result = {"UnprocessedTraceSegments": [x.to_dict() for x in bad_segments]}
return json.dumps(result)
# GetTraceSummaries
def trace_summaries(self):
start_time = self._get_param('StartTime')
end_time = self._get_param('EndTime')
start_time = self._get_param("StartTime")
end_time = self._get_param("EndTime")
if start_time is None:
msg = 'Parameter StartTime is missing'
return json.dumps({'__type': 'MissingParameter', 'message': msg}), dict(status=400)
msg = "Parameter StartTime is missing"
return (
json.dumps({"__type": "MissingParameter", "message": msg}),
dict(status=400),
)
if end_time is None:
msg = 'Parameter EndTime is missing'
return json.dumps({'__type': 'MissingParameter', 'message': msg}), dict(status=400)
msg = "Parameter EndTime is missing"
return (
json.dumps({"__type": "MissingParameter", "message": msg}),
dict(status=400),
)
filter_expression = self._get_param('FilterExpression')
sampling = self._get_param('Sampling', 'false') == 'true'
filter_expression = self._get_param("FilterExpression")
sampling = self._get_param("Sampling", "false") == "true"
try:
start_time = datetime.datetime.fromtimestamp(int(start_time))
end_time = datetime.datetime.fromtimestamp(int(end_time))
except ValueError:
msg = 'start_time and end_time are not integers'
return json.dumps({'__type': 'InvalidParameterValue', 'message': msg}), dict(status=400)
msg = "start_time and end_time are not integers"
return (
json.dumps({"__type": "InvalidParameterValue", "message": msg}),
dict(status=400),
)
except Exception as err:
return json.dumps({'__type': 'InternalFailure', 'message': str(err)}), dict(status=500)
return (
json.dumps({"__type": "InternalFailure", "message": str(err)}),
dict(status=500),
)
try:
result = self.xray_backend.get_trace_summary(start_time, end_time, filter_expression, sampling)
result = self.xray_backend.get_trace_summary(
start_time, end_time, filter_expression, sampling
)
except AWSError as err:
return err.response()
except Exception as err:
return json.dumps({'__type': 'InternalFailure', 'message': str(err)}), dict(status=500)
return (
json.dumps({"__type": "InternalFailure", "message": str(err)}),
dict(status=500),
)
return json.dumps(result)
# BatchGetTraces
def traces(self):
trace_ids = self._get_param('TraceIds')
next_token = self._get_param('NextToken') # not implemented yet
trace_ids = self._get_param("TraceIds")
next_token = self._get_param("NextToken") # not implemented yet
if trace_ids is None:
msg = 'Parameter TraceIds is missing'
return json.dumps({'__type': 'MissingParameter', 'message': msg}), dict(status=400)
msg = "Parameter TraceIds is missing"
return (
json.dumps({"__type": "MissingParameter", "message": msg}),
dict(status=400),
)
try:
result = self.xray_backend.get_trace_ids(trace_ids, next_token)
except AWSError as err:
return err.response()
except Exception as err:
return json.dumps({'__type': 'InternalFailure', 'message': str(err)}), dict(status=500)
return (
json.dumps({"__type": "InternalFailure", "message": str(err)}),
dict(status=500),
)
return json.dumps(result)
# GetServiceGraph - just a dummy response for now
def service_graph(self):
start_time = self._get_param('StartTime')
end_time = self._get_param('EndTime')
start_time = self._get_param("StartTime")
end_time = self._get_param("EndTime")
# next_token = self._get_param('NextToken') # not implemented yet
if start_time is None:
msg = 'Parameter StartTime is missing'
return json.dumps({'__type': 'MissingParameter', 'message': msg}), dict(status=400)
msg = "Parameter StartTime is missing"
return (
json.dumps({"__type": "MissingParameter", "message": msg}),
dict(status=400),
)
if end_time is None:
msg = 'Parameter EndTime is missing'
return json.dumps({'__type': 'MissingParameter', 'message': msg}), dict(status=400)
msg = "Parameter EndTime is missing"
return (
json.dumps({"__type": "MissingParameter", "message": msg}),
dict(status=400),
)
result = {
'StartTime': start_time,
'EndTime': end_time,
'Services': []
}
result = {"StartTime": start_time, "EndTime": end_time, "Services": []}
return json.dumps(result)
# GetTraceGraph - just a dummy response for now
def trace_graph(self):
trace_ids = self._get_param('TraceIds')
trace_ids = self._get_param("TraceIds")
# next_token = self._get_param('NextToken') # not implemented yet
if trace_ids is None:
msg = 'Parameter TraceIds is missing'
return json.dumps({'__type': 'MissingParameter', 'message': msg}), dict(status=400)
msg = "Parameter TraceIds is missing"
return (
json.dumps({"__type": "MissingParameter", "message": msg}),
dict(status=400),
)
result = {
'Services': []
}
result = {"Services": []}
return json.dumps(result)

View file

@ -1,15 +1,13 @@
from __future__ import unicode_literals
from .responses import XRayResponse
url_bases = [
"https?://xray.(.+).amazonaws.com",
]
url_bases = ["https?://xray.(.+).amazonaws.com"]
url_paths = {
'{0}/TelemetryRecords$': XRayResponse.dispatch,
'{0}/TraceSegments$': XRayResponse.dispatch,
'{0}/Traces$': XRayResponse.dispatch,
'{0}/ServiceGraph$': XRayResponse.dispatch,
'{0}/TraceGraph$': XRayResponse.dispatch,
'{0}/TraceSummaries$': XRayResponse.dispatch,
"{0}/TelemetryRecords$": XRayResponse.dispatch,
"{0}/TraceSegments$": XRayResponse.dispatch,
"{0}/Traces$": XRayResponse.dispatch,
"{0}/ServiceGraph$": XRayResponse.dispatch,
"{0}/TraceGraph$": XRayResponse.dispatch,
"{0}/TraceSummaries$": XRayResponse.dispatch,
}