def test_branch_setter_change():
"""Assert that we can change to an old branch."""
atom = ATOMClassifier(X10_nan, y10, random_state=1)
atom.branch = "branch_2"
atom.clean()
atom.branch = "master"
assert atom.pipeline.empty # Has no clean estimator
def test_getitem():
"""Assert that atom is subscriptable."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
atom.impute()
assert atom[1].__class__.__name__ == "Imputer"
assert atom["mean radius"].equals(atom.dataset["mean radius"])
def test_export_pipeline():
"""Assert that we can export the pipeline."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
pytest.raises(RuntimeError, atom.export_pipeline)
atom.clean()
atom.run(["GNB", "LGB"])
assert len(atom.export_pipeline("GNB")) == 2 # Without scaler
assert isinstance(atom.export_pipeline("LGB")[1], Scaler) # With scaler
def test_vote_branch_transformation(pipeline):
"""Assert that the branches transform every estimator only once."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
atom.impute()
atom.branch = "branch_2"
atom.encode()
atom.run(models=["Tree", "LGB"])
atom.voting()
assert isinstance(atom.vote.predict(X_bin, pipeline=pipeline), np.ndarray)
def test_vote_prediction_methods():
"""Assert that the prediction methods work as intended."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
atom.run(models=["Tree"])
atom.branch = "branch_2"
atom.impute(strat_num="mean", strat_cat="most_frequent")
atom.run(["LGB"])
atom.voting(models=["Tree", "LGB"])
pytest.raises(AttributeError, atom.vote.decision_function, X_bin)
assert isinstance(atom.vote.predict(X_bin), np.ndarray)
assert isinstance(atom.vote.predict_proba(X_bin), np.ndarray)
assert isinstance(atom.vote.score(X_bin, y_bin), np.float64)
def test_iter():
"""Assert that we can iterate over atom's pipeline."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
atom.impute()
assert [item for item in atom][1] == atom.pipeline[1]
def test_clean():
"""Assert that the clean method cleans the dataset."""
atom = ATOMClassifier(X10, y10_sn, random_state=1)
atom.clean()
assert len(atom.dataset) == 9
assert atom.mapping == {"n": 0, "y": 1}
def test_verbose_raises_when_invalid():
"""Assert an error is raised for an invalid value of verbose."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
pytest.raises(ValueError, atom.transform, X_bin, verbose=3)
def test_branch_setter_new():
"""Assert that we can create a new branch."""
atom = ATOMClassifier(X10_nan, y10, random_state=1)
atom.clean()
atom.branch = "branch_2"
assert list(atom._branches.keys()) == ["master", "branch_2"]
def test_len():
"""Assert that the length of atom is the length of the pipeline."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.clean()
assert len(atom) == 1