fix test

mhcflurry/predict_command.py

@@ -278,7 +278,7 @@ def run(argv=sys.argv[1:]):
             del predictions["best_allele"]
 
     for col in predictions.columns:
-        if col not in ("allele", "peptide", "experiment_name"):
+        if col not in ("allele", "peptide", "experiment_name", "peptide_num"):
             df[args.prediction_column_prefix + col] = predictions[col]
 
     if args.out:

test/test_predict_command.py

@@ -45,6 +45,7 @@ def test_no_csv():
         "--alleles", "HLA-A0201", "H-2-Kb",
         "--peptides", "SIINFEKL", "DENDREKLLL", "PICKLEEE",
         "--prediction-column-prefix", "mhcflurry1_",
+        "--affinity-only",
     ]
 
     deletes = []
@@ -61,7 +62,7 @@ def test_no_csv():
             os.unlink(delete)
 
     print(result)
-    assert_equal(result.shape, (6, 6))
+    assert_equal(len(result), 6)
     sub_result1 = result.loc[result.peptide == "SIINFEKL"].set_index("allele")
     print(sub_result1)
     assert (