def test__consume(self, mocker, redis_client):
# pylint: disable=W0613
queue = 'multiplex'
storage = DummyStorage()
consumer = consumers.MesmerConsumer(redis_client, storage, queue)
empty_data = {'input_file_name': 'file.tiff'}
output_shape = (1, 256, 256, 2)
mock_app = Bunch(
predict=lambda *x, **y: np.random.randint(1, 5, size=output_shape),
model_mpp=1,
model=Bunch(get_batch_size=lambda *x: 1,
input_shape=(1, 32, 32, 1)))
mocker.patch.object(consumer, 'get_grpc_app', lambda *x, **_: mock_app)
mocker.patch.object(consumer, 'get_image_scale', lambda *x, **_: 1)
mocker.patch.object(consumer, 'validate_model_input',
lambda *x, **_: x[0])
mocker.patch.object(consumer, 'detect_dimension_order',
lambda *x, **_: 'YXC')
test_hash = 'some hash'
redis_client.hmset(test_hash, empty_data)
result = consumer._consume(test_hash)
assert result == consumer.final_status
result = redis_client.hget(test_hash, 'status')
assert result == consumer.final_status
def test__consume(self, mocker, redis_client):
queue = 'track'
storage = DummyStorage()
test_hash = 0
dummy_results = {
'y_tracked': np.zeros((32, 32, 1)),
'tracks': []
}
mock_model = Bunch(
get_batch_size=lambda *x: 1,
input_shape=(1, 32, 32, 1)
)
mock_app = Bunch(
predict=lambda *x, **y: dummy_results,
track=lambda *x, **y: dummy_results,
model_mpp=1,
model=mock_model,
)
consumer = consumers.CalibanConsumer(redis_client, storage, queue)
mocker.patch.object(consumer, 'get_grpc_app',
lambda *x, **y: mock_app)
# mock get_model_wrapper for neighborhood encoder
mocker.patch.object(consumer, 'get_model_wrapper',
lambda *x, **y: mock_model)
frames = 3
dummy_data = {
'X': np.array([_get_image(21, 21) for _ in range(frames)]),
'y': np.random.randint(0, 9, size=(frames, 21, 21)),
}
mocker.patch.object(consumer, '_load_data', lambda *x: dummy_data)
# test finished statuses are returned
for status in (consumer.failed_status, consumer.final_status):
test_hash += 1
data = {'input_file_name': 'file.tiff', 'status': status}
redis_client.hmset(test_hash, data)
result = consumer._consume(test_hash)
assert result == status
# test new key is processed
test_hash += 1
data = {'input_file_name': 'file.tiff', 'status': 'new'}
redis_client.hmset(test_hash, data)
result = consumer._consume(test_hash)
assert result == consumer.final_status
assert redis_client.hget(test_hash, 'status') == consumer.final_status
def test__consume_finished_status(self, redis_client):
queue = 'q'
storage = DummyStorage()
consumer = consumers.MesmerConsumer(redis_client, storage, queue)
empty_data = {'input_file_name': 'file.tiff'}
test_hash = 0
# test finished statuses are returned
for status in (consumer.failed_status, consumer.final_status):
test_hash += 1
data = empty_data.copy()
data['status'] = status
redis_client.hmset(test_hash, data)
result = consumer._consume(test_hash)
assert result == status
result = redis_client.hget(test_hash, 'status')
assert result == status
test_hash += 1