def test_check_2d_on_2d():
df = pd.DataFrame({
"x1": [-2, -1, 0, 1, 2],
"x2": [-1, 2, -3, 4, -5],
"y": [2, -2, 0, 4, -10]
})
pd.testing.assert_frame_equal(check_2d(df), df)
a = np.arange(10).reshape(-1, 2)
assert np.array_equal(check_2d(a), a)
def test_check_2d_on_3d():
with pytest.raises(ValueError):
check_2d(np.arange(12).reshape(-1, 2, 2))
def test_check_2d_on_2d():
assert np.array_equal(check_2d(np.arange(10).reshape(-1, 2)),
np.arange(10).reshape(-1, 2))
def test_check_2d_on_1d():
with pytest.raises(ValueError):
check_2d(np.arange(3))
def test_check_2d_on_3d():
with pytest.raises(ValueError):
check_2d(np.arange(12).reshape(-1, 2, 2))
def test_check_2d_on_2d():
assert np.array_equal(check_2d(np.arange(10).reshape(-1, 2)), np.arange(10).reshape(-1, 2))
def test_check_2d_on_1d():
with pytest.raises(ValueError):
check_2d(np.arange(3))