def test_add_tag():
key = "my_key"
value = "my_value"
SpansContainer.get_span().add_tag(key, value)
assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [
{"key": key, "value": value}
]
def test_get_span_by_id():
container = SpansContainer.get_span()
container.add_span({"id": 1, "extra": "a"})
container.add_span({"id": 2, "extra": "b"})
container.add_span({"id": 3, "extra": "c"})
assert SpansContainer.get_span().get_span_by_id(2)["extra"] == "b"
assert SpansContainer.get_span().get_span_by_id(5) is None
def test_timeout_mechanism_too_short_time(monkeypatch, context):
monkeypatch.setattr(Configuration, "timeout_timer", True)
monkeypatch.setattr(context, "get_remaining_time_in_millis", lambda: 1000)
SpansContainer.create_span()
SpansContainer.get_span().start(context=context)
assert not TimeoutMechanism.is_activated()
def update_event_response(host: Optional[str], status_code: int, headers: dict,
body: bytes) -> None:
"""
:param host: If None, use the host from the last span, otherwise this is the first chuck and we can empty
the aggregated response body
This function assumes synchronous execution - we update the last http event.
"""
if is_lumigo_edge(host):
return
last_event = SpansContainer.get_span().pop_last_span()
if last_event:
http_info = last_event.get("info", {}).get("httpInfo", {})
if not host:
host = http_info.get("host", "unknown")
else:
HttpState.previous_response_body = b""
has_error = is_error_code(status_code)
max_size = Configuration.get_max_entry_size(has_error)
headers = {k.lower(): v for k, v in headers.items()} if headers else {}
parser = get_parser(host, headers)() # type: ignore
if len(HttpState.previous_response_body) < max_size:
HttpState.previous_response_body += body
if has_error:
_update_request_data_increased_size_limit(http_info, max_size)
update = parser.parse_response( # type: ignore
host,
status_code,
headers,
HttpState.previous_response_body # type: ignore
)
SpansContainer.get_span().add_span(
recursive_json_join(update, last_event))
def started(self, event):
with lumigo_safe_execute("pymongo started"):
span_id = str(uuid.uuid4())
LumigoMongoMonitoring.request_to_span_id[
event.request_id] = span_id
SpansContainer.get_span().add_span({
"id":
span_id,
"type":
self.MONGO_SPAN,
"started":
get_current_ms_time(),
"databaseName":
event.database_name,
"commandName":
event.command_name,
"request":
lumigo_dumps(event.command),
"mongoRequestId":
event.request_id,
"mongoOperationId":
event.operation_id,
"mongoConnectionId":
event.connection_id,
})
def add_unparsed_request(parse_params: HttpRequest):
"""
This function handle the case where we got a request the is not fully formatted as we expected,
I.e. there isn't '\r\n' in the request data that <i>logically</i> splits the headers from the body.
In that case, we will consider it as a continuance of the previous request if they got the same url,
and we didn't get any answer yet.
"""
if is_lumigo_edge(parse_params.host):
return
last_event = SpansContainer.get_span().get_last_span()
if last_event:
if last_event and last_event.get(
"type") == HTTP_TYPE and HttpState.previous_request:
if last_event.get("info",
{}).get("httpInfo",
{}).get("host") == parse_params.host:
if "response" not in last_event["info"]["httpInfo"]:
SpansContainer.get_span().pop_last_span()
body = (HttpState.previous_request.body +
parse_params.body)[:get_size_upper_bound()]
add_request_event(
HttpState.previous_request.clone(body=body))
return
add_request_event(parse_params.clone(headers=None))
def test_lambda_wrapper_exception(exc, context):
@lumigo_tracer(token=TOKEN)
def lambda_test_function(event, context):
a = "A" # noqa
raise exc
try:
lambda_test_function({}, context)
except ValueError:
pass
else:
assert False
function_span = SpansContainer.get_span().function_span
assert not SpansContainer.get_span().spans
assert function_span.get("error", {}).get("type") == "ValueError"
# Make sure no lumigo_tracer
assert len(function_span["error"]["frames"]) == 1
assert function_span["error"]["frames"][0].pop("lineno") > 0
assert function_span["error"]["frames"][0] == {
"function": "lambda_test_function",
"fileName": __file__,
"variables": {
"a": '"A"',
"context": f'"{str(context)}"',
"event": "{}",
"exc": f'"{str(exc)}"',
},
}
assert not function_span["id"].endswith("_started")
assert "reporter_rtt" in function_span
assert "maxFinishTime" not in function_span
# Test that we can create an output message out of this span
assert _create_request_body([function_span], prune_size_flag=False)
def test_catch_file_like_object_sent_on_http(context, token):
class A:
def seek(self, where):
pass
def tell(self):
return 1
def read(self, amount=None):
return b"body"
@lumigo_tracer.lumigo_tracer(token=token)
def lambda_test_function(event, context):
try:
http.client.HTTPConnection("www.github.com").send(A())
except Exception:
# We don't care about errors
pass
lambda_test_function({}, context)
http_events = list(SpansContainer.get_span().spans.values())
assert len(http_events) == 1
span = list(SpansContainer.get_span().spans.values())[0]
assert span["info"]["httpInfo"]["request"]["body"] == '"body"'
assert span["info"]["httpInfo"]["request"].get("instance_id") is not None
def test_spans_container_not_send_start_span_on_send_only_on_errors_mode(
monkeypatch):
Configuration.send_only_if_error = True
SpansContainer.create_span()
SpansContainer.get_span().start()
assert _is_start_span_sent() is False
def _read_stream_wrapper_generator(stream_generator, instance):
for partial_response in stream_generator:
with lumigo_safe_execute("parse response.read_chunked"):
SpansContainer.get_span().update_event_response(
None, instance.status, instance.headers, partial_response
)
yield partial_response
def add_execution_tag(key: str,
value: str,
should_log_errors: bool = True) -> bool:
"""
Use this function to add an execution_tag to your function with a dynamic value.
This value can be searched within the Lumigo platform.
The maximum number of tags is 50.
:param key: Length should be between 1 and 50.
:param value: Length should be between 1 and 70.
:param should_log_errors: Should a log message be printed in case the tag can't be added.
"""
try:
key = str(key)
value = str(value)
tags_len = SpansContainer.get_span().get_tags_len()
if validate_tag(key, value, tags_len, should_log_errors):
SpansContainer.get_span().add_tag(key, value)
else:
return False
except Exception:
if should_log_errors:
warn_client(ADD_TAG_ERROR_MSG_PREFIX)
return False
return True
def add_unparsed_request(span_id: Optional[str],
parse_params: HttpRequest) -> Optional[Dict]:
"""
This function handle the case where we got a request the is not fully formatted as we expected,
I.e. there isn't '\r\n' in the request data that <i>logically</i> splits the headers from the body.
In that case, we will consider it as a continuance of the previous request if they got the same url,
and we didn't get any answer yet.
"""
if is_lumigo_edge(parse_params.host):
return None
last_event = SpansContainer.get_span().get_span_by_id(span_id)
if last_event:
if last_event and last_event.get("type") == HTTP_TYPE:
http_info = last_event.get("info", {}).get("httpInfo", {})
if http_info.get("host") == parse_params.host:
if "response" not in http_info:
SpansContainer.get_span().get_span_by_id(span_id)
http_info["request"]["body"] = concat_old_body_to_new(
http_info.get("request", {}).get("body"),
parse_params.body)
if HttpState.previous_span_id == span_id and HttpState.previous_request:
HttpState.previous_request.body += parse_params.body
return last_event
return add_request_event(span_id, parse_params)
def _putheader_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function that called after that the http request was sent.
Note that we don't examine the response data because it may change the original behaviour (ret_val.peek()).
"""
if SpansContainer.get_span().can_path_root():
kwargs["headers"]["X-Amzn-Trace-Id"] = SpansContainer.get_span(
).get_patched_root()
ret_val = func(*args, **kwargs)
return ret_val
def add_request_event(parse_params: HttpRequest):
"""
This function parses an request event and add it to the span.
"""
if is_lumigo_edge(parse_params.host):
return
parser = get_parser(parse_params.host)()
msg = parser.parse_request(parse_params)
HttpState.previous_request = parse_params
SpansContainer.get_span().add_span(msg)
def test_spans_container_end_function_send_only_on_errors_mode_false_not_effecting(
monkeypatch, dummy_http_request):
SpansContainer.create_span()
SpansContainer.get_span().start()
SpansContainer.get_span().add_request_event(dummy_http_request)
reported_ttl = SpansContainer.get_span().end({})
assert reported_ttl is not None
def _read_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function that can be called only after `getresponse` was called.
"""
ret_val = func(*args, **kwargs)
if ret_val:
with lumigo_safe_execute("parse response.read"):
SpansContainer.get_span().update_event_response(
None, instance.code, instance.headers, ret_val
)
return ret_val
def test_malformed_txid(monkeypatch, context):
monkeypatch.setenv(
"_X_AMZN_TRACE_ID", f"Root=1-5fd891b8-{MALFORMED_TXID};Parent=0a885f800de045d4;Sampled=0"
)
SpansContainer.create_span({}, context)
assert SpansContainer.get_span().transaction_id != MALFORMED_TXID
assert SpansContainer.get_span().function_span["isMalformedTransactionId"]
result = SpansContainer.get_span().get_patched_root()
output_trace_id = result.split(";")[0].split("=")[1].split("-")[2]
assert output_trace_id == SpansContainer.get_span().transaction_id
def _requests_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function `requests.request`.
This function is being wrapped specifically because it initializes the connection by itself and parses the response,
which creates a gap from the traditional http.client wrapping.
"""
start_time = datetime.now()
ret_val = func(*args, **kwargs)
with lumigo_safe_execute("requests wrapper time updates"):
SpansContainer.get_span().update_event_times(start_time=start_time)
return ret_val
def _response_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function that can be called only after that the http request was sent.
Note that we don't examine the response data because it may change the original behaviour (ret_val.peek()).
"""
ret_val = func(*args, **kwargs)
with lumigo_safe_execute("parse response"):
headers = ret_val.headers
status_code = ret_val.code
SpansContainer.get_span().update_event_response(instance.host, status_code, headers, b"")
return ret_val
def test_spans_container_end_function_not_send_spans_on_send_only_on_errors_mode(
monkeypatch, dummy_http_request):
Configuration.send_only_if_error = True
SpansContainer.create_span()
SpansContainer.get_span().start()
SpansContainer.get_span().add_request_event(dummy_http_request)
reported_ttl = SpansContainer.get_span().end({})
assert reported_ttl is None
def test_timeout_mechanism_timeout_occurred_send_new_spans(monkeypatch, context, dummy_span):
SpansContainer.create_span()
SpansContainer.get_span().start(context=context)
SpansContainer.get_span().add_span(dummy_span)
SpansContainer.get_span().handle_timeout()
SpansContainer.get_span().add_span(dummy_span)
assert SpansContainer.get_span().span_ids_to_send
def only_if_error(dummy_span, monkeypatch, tmpdir):
extension_dir = tmpdir.mkdir("tmp")
monkeypatch.setenv("LUMIGO_EXTENSION_SPANS_DIR_KEY", extension_dir)
monkeypatch.setattr(uuid, "uuid4", lambda *args, **kwargs: "span_name")
Configuration.send_only_if_error = True
SpansContainer.create_span()
SpansContainer.get_span().start()
SpansContainer.get_span().add_span(dummy_span)
reported_ttl = SpansContainer.get_span().end({})
stop_path_path = f"{lumigo_utils.get_extension_dir()}/span_name_stop"
return reported_ttl, stop_path_path
def test_lambda_wrapper_http():
@lumigo_tracer(token="123")
def lambda_test_function():
time.sleep(0.01)
http.client.HTTPConnection("www.google.com").request("POST", "/")
lambda_test_function()
http_spans = SpansContainer.get_span().http_spans
assert http_spans
assert http_spans[0].get("info", {}).get("httpInfo", {}).get("host") == "www.google.com"
assert "started" in http_spans[0]
assert http_spans[0]["started"] > SpansContainer.get_span().function_span["started"]
assert "ended" in http_spans[0]
assert "Content-Length" in http_spans[0]["info"]["httpInfo"]["request"]["headers"]
def test_omitting_keys():
@lumigo_tracer()
def lambda_test_function(event, context):
d = {"a": "b", "myPassword": "******"}
conn = http.client.HTTPConnection("www.google.com")
conn.request("POST", "/", json.dumps(d))
return {"secret_password": "******"}
lambda_test_function({"key": "24"}, None)
span = SpansContainer.get_span()
assert span.function_span["return_value"] == '{"secret_password": "******"}'
assert span.function_span["event"] == '{"key": "****"}'
assert SpansContainer.get_span().http_spans[0]["info"]["httpInfo"]["request"][
"body"
] == json.dumps({"a": "b", "myPassword": "******"})
def test_lambda_wrapper_http(context, token):
@lumigo_tracer.lumigo_tracer(token=token)
def lambda_test_function(event, context):
time.sleep(0.01)
http.client.HTTPConnection("www.google.com").request("POST", "/")
lambda_test_function({}, context)
http_spans = list(SpansContainer.get_span().spans.values())
assert http_spans
assert http_spans[0].get("info", {}).get("httpInfo", {}).get("host") == "www.google.com"
assert "started" in http_spans[0]
assert http_spans[0]["started"] > SpansContainer.get_span().function_span["started"]
assert "ended" in http_spans[0]
assert "content-length" in http_spans[0]["info"]["httpInfo"]["request"]["headers"]
assert http_spans[0]["info"]["httpInfo"]["request"].get("instance_id") is not None
def test_wrapping_urlib_stream_get():
"""
This is the same case as the one of `requests.get`.
"""
@lumigo_tracer()
def lambda_test_function(event, context):
r = urllib3.PoolManager().urlopen("GET", "https://www.google.com", preload_content=False)
return b"".join(r.stream(32))
lambda_test_function({}, None)
assert len(SpansContainer.get_span().http_spans) == 1
event = SpansContainer.get_span().http_spans[0]
assert event["info"]["httpInfo"]["response"]["body"]
assert event["info"]["httpInfo"]["response"]["statusCode"] == 200
assert event["info"]["httpInfo"]["host"] == "www.google.com"
def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
with lumigo_safe_execute("handle sqlalchemy after execute"):
span = SpansContainer.get_span().get_last_span()
if not span:
get_logger().warning("Redis span ended without a record on its start")
return
span.update({"ended": get_current_ms_time(), "response": ""})
def test_report_error_without_enhance_print(capsys):
Configuration.enhanced_print = False
SpansContainer.get_span().function_span["id"] = "123"
msg = "oh no - an error"
report_error(msg)
captured = capsys.readouterr()
assert captured.out == f"RequestId: 123 {LUMIGO_REPORT_ERROR_STRING} {msg}\n"
def test_kinesis(kinesis_resource, region):
@lumigo_tracer(token=TOKEN)
def lambda_test_function():
client = boto3.client("kinesis")
client.put_record(StreamName=kinesis_resource, Data=b"my data", PartitionKey="1")
client.put_records(
StreamName=kinesis_resource,
Records=[
{"Data": "First", "PartitionKey": "1"},
{"Data": "Second", "PartitionKey": "1"},
],
)
lambda_test_function()
events = SpansContainer.get_span().spans
assert len(events) == 2
# Single message.
assert events[0]["info"]["httpInfo"]["host"] == f"kinesis.{region}.amazonaws.com"
assert events[0]["info"]["resourceName"] == kinesis_resource
assert events[0]["info"]["messageId"]
# No scrubbing for PartitionKey
assert json.loads(events[0]["info"]["httpInfo"]["request"]["body"])["PartitionKey"] == "1"
# Batch messages.
assert events[1]["info"]["httpInfo"]["host"] == f"kinesis.{region}.amazonaws.com"
assert events[1]["info"]["resourceName"] == kinesis_resource
assert events[1]["info"]["messageId"]
def test_configuration_handler_auto_tag_non_string(value, expected):
Configuration.auto_tag = ["key1"]
ConfigurationHandler.auto_tag({"key1": value})
tags = SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY]
assert {"key": "key1", "value": expected} in tags
