add test for predict method

mhcflurry/predict.py

@@ -17,11 +17,19 @@ from collections import OrderedDict
 import pandas as pd
 
 from .class1_binding_predictor import Class1BindingPredictor
+from .common import normalize_allele_name
 
 def predict(alleles, peptides):
-    allele_results = OrderedDict([])
+    result_dict = OrderedDict([
+        ("allele", []),
+        ("peptide", []),
+        ("ic50", []),
+    ])
     for allele in alleles:
+        allele = normalize_allele_name(allele)
         model = Class1BindingPredictor.from_allele_name(allele)
-        result_dictionary = model.predict_peptides(peptides)
-        allele_results.append(result_dictionary)
-    return pd.concat(allele_results)
+        for i, ic50 in enumerate(model.predict_peptides_ic50(peptides)):
+            result_dict["allele"].append(allele)
+            result_dict["peptide"].append(peptides[i])
+            result_dict["ic50"].append(ic50)
+    return pd.DataFrame(result_dict)

test/test_predict.py

+from mhcflurry import predict
+
+def test_predict_A1_Titin_epitope():
+    result_df = predict(alleles=["HLA-A*01:01"], peptides=["ESDPIVAQY"])
+    assert len(result_df) == 1
+    row = result_df.ix[0]
+    ic50 = row["ic50"]
+    assert ic50 <= 500