def test_set_batch_size(self, mock_argparse):
mock_return_partial_shape(PartialShape([-1, 2, 3, 4]))
main(argparse.ArgumentParser(), fem, 'mock_mo_ngraph_frontend')
stat = get_model_statistic()
# verify that 'set_element_type' was called
# 2 is because mock model has 2 inputs
assert stat.get_partial_shape == 2
assert stat.set_partial_shape == 2
assert stat.lastArgPartialShape == PartialShape([123, 2, 3, 4])
def test_error_batch(self, mock_argparse):
# First dimension doesn't look like a batch,
# so MO shall not convert anything and produce specified error
mock_return_partial_shape(PartialShape([122, 2, 3, 4]))
with self.assertLogs() as logger:
main(argparse.ArgumentParser(), fem, 'mock_mo_ngraph_frontend')
stat = get_model_statistic()
assert [s for s in logger.output if 'question=39' in s]
# verify that 'get_element_type' was called
assert stat.get_partial_shape == 1
# verify that 'set_element_type' was not called
assert stat.set_partial_shape == 0
