def test_get_models_protobuf(self):
'''
Test for endpoint: /api/v2/Model/
With accept-type "application/x-protobuf"
'''
response = self.client.get('/api/v2/Model/', \
**{'HTTP_ACCEPT':'application/x-protobuf'})
self.assertEqual(response.status_code, 200, \
msg='/api/v2/Model/ endpoint did not return 200 status code.')
response_pb = Model_pb2.ModelsCollection()
response_pb.ParseFromString(response.content)
self.assertEqual(len(response_pb.Models), 1, \
msg='/api/v2/Model/ endpoint did not return correct JSON message.')
self.assertEqual(response_pb.Models[0].CanonicalName, \
GetModelsTestCase.model_family.canonical_name, \
msg='/api/v2/Model/ endpoint did not return correct Phantom model.')
def get_models(request):
'''
Endpoint for API GET Model Collections request.
Returns:
1. ModelsCollection protobuf message if "accept" header is "application/x-protobuf"
2. JSON response otherwise
'''
# Tries to construct collection of model families, returning an error if failure.
try:
# Creates protobuf collection object
response = Model_pb2.ModelsCollection()
# Iterates through all model familys and adds to collection
for model_fam in ModelFamily.objects.all():
# Check if pb file exists.
if model_fam.pb == '':
fname = model_fam.canonical_name
fname.replace(' ', '-')
pb_msg = model_fam.model_to_pb().SerializeToString()
file_io = io.BytesIO(bytes(pb_msg))
model_fam.pb.save(fname + ".pb", File(file_io))
model_fam.save()
mf_file = open(model_fam.pb.path, 'rb')
file_pb = mf_file.read()
mf_file.close()
# Check if fields are filled.
if model_fam.canonical_name == '':
model_fam.pb_to_model(file_pb)
model_fam.save()
mf_pb = Model_pb2.ModelFamily()
mf_pb.ParseFromString(file_pb)
response.Models.append(mf_pb)
return format_and_send_response(request, response)
# On exception, returns invalid request error message pb.
except:
msg = "Invalid request."
details = "The request is not acceptable."
return bad_request_helper(request, msg, details, 406)