tests/test_predict.py

"""
Tests for the predict.py module
"""
import sys
import pathlib
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
from predict import predict, Tox21RandomClassifier


class TestTox21RandomClassifier:
    
    def test_init(self):
        classifier = Tox21RandomClassifier()
        assert len(classifier.target_names) == 12
        expected_targets = [
            "NR-AR", "NR-AR-LBD", "NR-AhR", "NR-Aromatase", 
            "NR-ER", "NR-ER-LBD", "NR-PPAR-gamma", 
            "SR-ARE", "SR-ATAD5", "SR-HSE", "SR-MMP", "SR-p53"
        ]
        assert classifier.target_names == expected_targets
    
    def test_predict_single_smiles(self):
        classifier = Tox21RandomClassifier()
        smiles_list = ["CCO"]
        result = classifier.predict(smiles_list)
        
        assert "CCO" in result
        assert len(result["CCO"]) == 12
        
        for target in classifier.target_names:
            assert target in result["CCO"]
            assert 0 <= result["CCO"][target] <= 1
    
    def test_predict_multiple_smiles(self):
        classifier = Tox21RandomClassifier()
        smiles_list = ["CCO", "CCN", "CCC"]
        result = classifier.predict(smiles_list)
        
        assert len(result) == 3
        for smiles in smiles_list:
            assert smiles in result
            assert len(result[smiles]) == 12
            
            for target in classifier.target_names:
                assert target in result[smiles]
                assert 0 <= result[smiles][target] <= 1
    
    def test_predict_empty_list(self):
        classifier = Tox21RandomClassifier()
        result = classifier.predict([])
        assert result == {}


class TestPredictFunction:
    
    def test_predict_function(self):
        smiles_list = ["CCO", "CCN"]
        result = predict(smiles_list)
        
        assert len(result) == 2
        for smiles in smiles_list:
            assert smiles in result
            assert len(result[smiles]) == 12
    
    def test_predict_function_empty(self):
        result = predict([])
        assert result == {}

#---------------------------------------------------------------------------------------
# Debugging
if __name__ == "__main__":
    test = TestTox21RandomClassifier()
    test.test_predict_multiple_smiles()