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