def test_get_model_name_digits():
"""Assert that a list of models is returned if using digits."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.successive_halving(["LR", "ET", "RF", "LGB"])
assert atom._get_model_name("2") == ["LR2", "ET2"]
def test_get_model_name_exact_name():
"""Assert that a single model is returned if the name matches exactly."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.run(["LR", "LR2"])
assert atom._get_model_name("lr") == ["LR"]
def test_get_model_name_multiple_models():
"""Assert that a list of models is returned when starting the same."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.run(["LR1", "LR2"])
assert atom._get_model_name("lr") == ["LR1", "LR2"]
def test_get_model_name_winner():
"""Assert that the winner is returned when used as name."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.run(["LR", "LGB"])
assert atom._get_model_name("winner") == ["LGB"]