TFLite (tools): added test for xxd_output_to_bytes; fixed styling

This commit is contained in:
danielyou0230 2020-09-01 11:43:48 -07:00
parent bf0764d491
commit c7d3d6eee2
2 changed files with 30 additions and 1 deletions

View File

@ -159,5 +159,34 @@ class RandomizeWeightsTest(test_util.TensorFlowTestCase):
self.assertNotEqual(initial_buffer.data[j], final_buffer.data[j]) self.assertNotEqual(initial_buffer.data[j], final_buffer.data[j])
class XxdOutputToBytesTest(test_util.TensorFlowTestCase):
def testXxdOutputToBytes(self):
# 1. SETUP
# Define the initial model
initial_model = test_utils.build_mock_model()
initial_bytes = flatbuffer_utils.convert_object_to_bytearray(initial_model)
# Define temporary files
tmp_dir = self.get_temp_dir()
model_filename = os.path.join(tmp_dir, 'model.tflite')
# 2. INVOKE
# Invoke the write_model and read_model functions
flatbuffer_utils.write_model(initial_model, model_filename)
# 3. DUMP WITH xxd
input_cc_file = os.path.join(tmp_dir, 'model.cc')
command = "xxd -i {} > {}".format(model_filename, input_cc_file)
subprocess.call(command, shell=True)
# 4. VALIDATE
final_bytes = flatbuffer_utils.xxd_output_to_bytes(input_cc_file)
# Validate that the initial and final bytearray are the same
self.assertEqual(initial_bytes, final_bytes)
if __name__ == '__main__': if __name__ == '__main__':
test.main() test.main()