def test_image_consumer(): mock_messages = MOCK_CONSUMER_MESSAGES hdf5_file = MockH5PYFile() prepare_hdf5_file(hdf5_file, 4, 5, 8) socket = MockSocket(zmq.PULL, to_recv=mock_messages) image_consumer(socket, hdf5_file, 4) assert_equal(hdf5_file['encoded_images'][0], [6, 6, 6]) assert_equal(hdf5_file['encoded_images'][1], [1, 8, 1, 2, 0]) assert_equal(hdf5_file['encoded_images'][2], [1, 9, 7, 9]) assert_equal(hdf5_file['encoded_images'][3], [1, 8, 6, 7]) assert_equal(hdf5_file['filenames'][:4], [[b'foo.jpeg'], [b'bar.jpeg'], [b'baz.jpeg'], [b'bur.jpeg']]) assert_equal(hdf5_file['targets'][:4], [[2], [3], [5], [7]])
def test_images_consumer_randomized(): mock_messages = MOCK_CONSUMER_MESSAGES + [ {'type': 'recv_pyobj', 'flags': zmq.SNDMORE, 'obj': ('jenny.jpeg', 1)}, {'type': 'recv', 'flags': 0, 'data': numpy.cast['uint8']([8, 6, 7, 5, 3, 0, 9])} ] hdf5_file = MockH5PYFile() prepare_hdf5_file(hdf5_file, 4, 5, 8) socket = MockSocket(zmq.PULL, to_recv=mock_messages) image_consumer(socket, hdf5_file, 5, offset=4, shuffle_seed=0) written_data = set(tuple(s) for s in hdf5_file['encoded_images'][4:9]) expected_data = set(tuple(s['data']) for s in mock_messages[1::2]) assert written_data == expected_data written_targets = set(hdf5_file['targets'][4:9].flatten()) expected_targets = set(s['obj'][1] for s in mock_messages[::2]) assert written_targets == expected_targets written_filenames = set(hdf5_file['filenames'][4:9].flatten()) expected_filenames = set(s['obj'][0].encode('ascii') for s in mock_messages[::2]) assert written_filenames == expected_filenames