@pytest.fixture def dummy_feats_and_labels(): feats = np.array([[0.7057, -5.4981, 8.3368, -2.8715], [2.4391, 6.4417, -0.80743, -0.69139], [-0.2062, 9.2207, -3.7044, -6.8103], [4.2586, 11.2962, -4.0943, -4.3457], [-2.343, 12.9516, 3.3285, -5.9426], [-2.0545, -10.8679, 9.4926, -1.4116], [2.2279, 4.0951, -4.8037, -2.1112], [-6.1632, 8.7096, -0.21621, -3.6345], [0.52374, 3.644, -4.0746, -1.9909], [1.5077, 1.9596, -3.0584, -0.12243] ]) labels = np.array([1, 1, 1, 1, 1, 0, 0, 0, 0, 0]) return feats, labels def test_dt_overfit(dummy_feats_and_labels): feats, labels = dummy_feats_and_labels dt = DecisionTree() dt.fit(feats, labels) pred = np.round(dt.predict(feats)) assert np.array_equal(labels, pred), 'DecisionTree should fit data perfectly and prediction should == labels.'