def test_predict():
model = Mock()
layer = Mock
criterion = Mock()
optimizer = Mock()
metrics = Mock()
sizes = [(1, 24, 24), (24, 10)]
l1 = layer()
l1.get_input_shape_at.return_value = sizes[0]
l2 = layer()
l2.get_output_shape_at.return_value = sizes[1]
model.layers = [l1, l2]
batch = 32
epoch = 2
kdpm = DeepLearningModel(model, criterion, optimizer, batch, epoch, metrics)
num_data = 30
data = np.array([np.random.rand(24, 24) for i in range(num_data)])
kdpm.predict(data)
kdpm._model.predict.assert_called_once_with(data, batch_size=batch)
def test_wrong_predict():
model = Mock()
layer = Mock
sizes = [(1, 24, 24), (24, 10)]
l1 = layer()
l1.get_input_shape_at.return_value = sizes[0]
l2 = layer()
l2.get_output_shape_at.return_value = sizes[1]
model.layers = [l1, l2]
batch = 32
epoch = 2
kdpm = DeepLearningModel(model, batch, epoch)
num_data = 30
data = np.array([np.random.rand(16, 16) for i in range(num_data)])
with pytest.raises(AssertionError):
kdpm.predict(data)