def _FormatResultAsJson(self, result, format_mode=None):
if result is None:
return dict(status="OK")
if format_mode == JsonMode.PROTO3_JSON_MODE:
return json.loads(
json_format.MessageToJson(result.AsPrimitiveProto()))
elif format_mode == JsonMode.GRR_ROOT_TYPES_STRIPPED_JSON_MODE:
result_dict = {}
for field, value in result.ListSetFields():
if isinstance(
field,
(rdf_structs.ProtoDynamicEmbedded,
rdf_structs.ProtoEmbedded, rdf_structs.ProtoList)):
result_dict[field.name] = api_value_renderers.RenderValue(
value)
else:
result_dict[field.name] = api_value_renderers.RenderValue(
value)["value"]
return result_dict
elif format_mode == JsonMode.GRR_TYPE_STRIPPED_JSON_MODE:
rendered_data = api_value_renderers.RenderValue(result)
return api_value_renderers.StripTypeInfo(rendered_data)
elif format_mode == JsonMode.GRR_JSON_MODE:
return api_value_renderers.RenderValue(result)
else:
raise ValueError("Invalid format_mode: %s" % format_mode)
def HandleCheck(self, method_metadata, args=None, replace=None):
"""Does regression check for given method, args and a replace function."""
if not replace:
raise ValueError("replace can't be None")
if self.__class__.api_version == 1:
request, prepped_request = self._PrepareV1Request(
method_metadata.name, args=args)
elif self.__class__.api_version == 2:
request, prepped_request = self._PrepareV2Request(
method_metadata.name, args=args)
else:
raise ValueError("api_version may be only 1 or 2, not %d" %
flags.FLAGS.api_version)
session = requests.Session()
response = session.send(prepped_request)
check_result = {
"url": replace(prepped_request.path_url),
"method": request.method
}
if request.data:
if compatibility.PY2:
data = request.data.decode("utf-8")
else:
data = request.data
request_payload = self._ParseJSON(replace(data))
if request_payload:
check_result["request_payload"] = request_payload
if (method_metadata.result_type == api_call_router.
RouterMethodMetadata.BINARY_STREAM_RESULT_TYPE):
check_result["response"] = replace(
utils.SmartUnicode(response.content))
else:
content = response.content.decode("utf-8")
check_result["response"] = self._ParseJSON(replace(content))
if self.__class__.api_version == 1:
stripped_response = api_value_renderers.StripTypeInfo(
check_result["response"])
if stripped_response != check_result["response"]:
check_result["type_stripped_response"] = stripped_response
return check_result
def _PrepareV1Request(self, method, args=None):
"""Prepares API v1 request for a given method and args."""
args_proto = None
if args:
args_proto = args.AsPrimitiveProto()
request = self.connector.BuildRequest(method, args_proto)
request.url = request.url.replace("/api/v2/", "/api/")
if args and request.data:
body_proto = args.__class__().AsPrimitiveProto()
json_format.Parse(request.data, body_proto)
body_args = args.__class__()
body_args.ParseFromString(body_proto.SerializeToString())
request.data = json.dumps(
api_value_renderers.StripTypeInfo(
api_value_renderers.RenderValue(body_args)),
cls=http_api.JSONEncoderWithRDFPrimitivesSupport)
prepped_request = request.prepare()
return request, prepped_request
