def test_featurize_data():
"""
Test that the featurize_data model correctly outputs the features of a toy
network on a toy tensor
"""
# Create the checked array
init_array = np.ones((5, 5, 5, 3))
for i in range(5):
init_array[i] = init_array[i] * i
# Check the prediction vs. the saved array
check_array = np.load('tests/data_featurizing_testing/array_testing/check_featurize.npy')
assert np.allclose(featurize_data(MODEL, init_array), check_array, atol=ATOL)
def test_featurize_data_bad_array():
"""Test errors with a badly formatted array"""
error_array = np.ones((5, 5, 10))
with pytest.raises(ValueError):
featurize_data(MODEL, error_array)