def test_parameters(self):
input_img = build_img((16, 16))
reference_img = build_img((36, 36))
class ReferenceImgOperator(Operator):
def run(self, img: Image, ctx: PipelineContext) -> Image:
return reference_img
reference_operator = ReferenceImgOperator()
class TestHook(PipelineHook):
def before_operator(self, operator: Operator, img: Image,
ctx: PipelineContext):
assert img is input_img
assert operator is reference_operator
def after_operator(self, operator: Operator, img: Image,
ctx: PipelineContext):
assert img is reference_img
assert operator is reference_operator
def after_pipeline(self, img: Image, ctx: PipelineContext):
assert img is reference_img
hook = TestHook()
pipeline = CompVizPipeline()
pipeline.add_operator('test', reference_operator)
_ = pipeline.run(input_img, hooks=[hook])
def test_pipeline_execution_order(order_asserter_operator):
pipeline = CompVizPipeline()
pipeline.add_operator('op1', order_asserter_operator(0))
pipeline.add_operator('op2', order_asserter_operator(1))
pipeline.add_operator('op3', order_asserter_operator(2))
pipeline.run(build_img((16, 16)))
def test_add_info_works(self):
info_name = 'info_name'
info_value = object()
class TestCtxAddInfo(Operator):
def run(self, img: Image, ctx: PipelineContext) -> Image:
ctx.add_info(info_name, info_value)
return img
pipeline = CompVizPipeline()
op_name1 = 'test_op1'
pipeline.add_operator(op_name1, TestCtxAddInfo())
op_name2 = 'test_op2'
pipeline.add_operator(op_name2, TestCtxAddInfo())
_, returned_ctx = pipeline.run(build_img((16, 16)))
expected_info = {
op_name1: {
info_name: info_value
},
op_name2: {
info_name: info_value
}
}
actual_info = returned_ctx.info
assert actual_info == expected_info
def test_hooks_order(self, pipeline):
call_count = 0
class TestHook(PipelineHook):
def before_pipeline(self, ctx: PipelineContext):
nonlocal call_count
assert call_count == 0
call_count += 1
def before_operator(self, operator: Operator, img: Image,
ctx: PipelineContext):
nonlocal call_count
assert call_count in (1, 3)
call_count += 1
def after_operator(self, operator: Operator, img: Image,
ctx: PipelineContext):
nonlocal call_count
assert call_count in (2, 4)
call_count += 1
def after_pipeline(self, img: Image, ctx: PipelineContext):
nonlocal call_count
assert call_count == 5
call_count += 1
img = build_img((16, 16))
_ = pipeline.run(img, hooks=[TestHook()])
assert call_count == 6
def test_rename_keeps_operators_order(self, order_asserter_operator):
pipeline = CompVizPipeline()
pipeline.add_operator('op1', order_asserter_operator(0))
pipeline.add_operator('op2', order_asserter_operator(1))
pipeline.add_operator('op3', order_asserter_operator(2))
pipeline.rename_operator('op2', 'new_name')
pipeline.run(build_img((16, 16)))
def test_operator_fail():
class FailingOperator(Operator):
def run(self, img: Image, ctx: PipelineContext) -> Image:
raise ValueError('Failed')
pipeline = CompVizPipeline()
pipeline.add_operator('op', FailingOperator())
with pytest.raises(OperatorFailedError):
pipeline.run(build_img((16, 16)))
def test_check_op_return(self, return_value):
class TestWrongReturnOperator(Operator):
def run(self, img: np.ndarray, ctx: PipelineContext) -> Any:
return return_value
pipeline = CompVizPipeline()
pipeline.add_operator('test_op', TestWrongReturnOperator())
with pytest.raises(BadImageError) as e:
pipeline.run(build_img((16, 16)))
assert_terms_in_exception(e, ['return', 'invalid'])
def test_gray_only_flag(self):
@settings.GRAY_ONLY(True)
class TestOnlyGrayFlagOperator(Operator):
def run(self, img: Image, ctx: PipelineContext) -> Image:
return img
pipeline = CompVizPipeline()
pipeline.add_operator('test_op', TestOnlyGrayFlagOperator())
with pytest.raises(OperatorFailedError) as e:
pipeline.run(build_img((16, 16), rgb=True))
assert_terms_in_exception(e, ["expect", "gray"])
def test_operator_cant_alter_original_img(self):
class TestCtxOriginalImg(Operator):
def run(self, img: Image, ctx: PipelineContext) -> Image:
original_img = ctx.original_img
original_img[10, ...] = 255
return img
pipeline = CompVizPipeline()
pipeline.add_operator('test_op', TestCtxOriginalImg())
test_img = build_img((128, 128), kind='black')
with pytest.raises(OperatorFailedError) as e:
pipeline.run(test_img)
assert_terms_in_exception(e, ['read-only'])
def test_pipeline_run_empty_info(self, pipeline):
img = build_img((16, 16))
out, ctx = pipeline.run(img)
assert ctx.info == {'op1': {}, 'op2': {}}
def original_img():
return build_img((16, 16))