From 05def7d224ced9be5b38ac9f9261c4b9aa23a1f0 Mon Sep 17 00:00:00 2001
From: Tim O'Donnell <timodonnell@gmail.com>
Date: Mon, 9 Sep 2019 12:07:07 -0400
Subject: [PATCH] Add teardown to tests to clear keras session

---
 .travis.yml                                      | 3 +--
 mhcflurry/testing_utils.py                       | 5 +++++
 test/test_calibrate_percentile_ranks_command.py  | 3 +++
 test/test_changing_allele_representations.py     | 3 +++
 test/test_class1_affinity_predictor.py           | 5 +++++
 test/test_class1_neural_network.py               | 3 +++
 test/test_class1_pan.py                          | 3 +++
 test/test_custom_loss.py                         | 3 +++
 test/test_download_models_class1.py              | 3 +++
 test/test_multi_output.py                        | 3 +++
 test/test_network_merging.py                     | 3 +++
 test/test_predict_command.py                     | 3 +++
 test/test_released_predictors_on_hpv_dataset.py  | 3 +++
 test/test_released_predictors_well_correlated.py | 3 +++
 test/test_speed.py                               | 3 +++
 test/test_train_and_related_commands.py          | 3 +++
 test/test_train_pan_allele_models_command.py     | 3 +++
 17 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 mhcflurry/testing_utils.py

diff --git a/.travis.yml b/.travis.yml
index c13fc1db..11d0d831 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -44,5 +44,4 @@ script:
   # download data and models, then run tests
   - mhcflurry-downloads fetch data_curated models_class1 models_class1_pan allele_sequences
   - mhcflurry-downloads info  # just to test this command works
-  #- travis_wait 30 nosetests --with-timer -sv
-  - bash -c 'for i in $(ls test/test_*.py) ; do echo "Invoking test: $i" ; nosetests -sv $i ; done'
+  - nosetests --with-timer -sv test
diff --git a/mhcflurry/testing_utils.py b/mhcflurry/testing_utils.py
new file mode 100644
index 00000000..862369c6
--- /dev/null
+++ b/mhcflurry/testing_utils.py
@@ -0,0 +1,5 @@
+
+
+def module_cleanup():
+    import keras.backend as K
+    K.clear_session()
diff --git a/test/test_calibrate_percentile_ranks_command.py b/test/test_calibrate_percentile_ranks_command.py
index 2a891cb4..4c97e80c 100644
--- a/test/test_calibrate_percentile_ranks_command.py
+++ b/test/test_calibrate_percentile_ranks_command.py
@@ -14,6 +14,9 @@ from mhcflurry.downloads import get_path
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 
 def run_and_check(n_jobs=0, delete=True, additional_args=[]):
     source_models_dir = get_path("models_class1_pan", "models.with_mass_spec")
diff --git a/test/test_changing_allele_representations.py b/test/test_changing_allele_representations.py
index 84492fad..7c70048d 100644
--- a/test/test_changing_allele_representations.py
+++ b/test/test_changing_allele_representations.py
@@ -8,6 +8,9 @@ from mhcflurry.downloads import get_path
 
 from numpy.testing import assert_equal
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 ALLELE_TO_SEQUENCE = pandas.read_csv(
     get_path(
         "allele_sequences", "allele_sequences.csv"),
diff --git a/test/test_class1_affinity_predictor.py b/test/test_class1_affinity_predictor.py
index 024d4a6a..277e78cb 100644
--- a/test/test_class1_affinity_predictor.py
+++ b/test/test_class1_affinity_predictor.py
@@ -16,6 +16,10 @@ from numpy import testing
 
 from mhcflurry.downloads import get_path
 
+import mhcflurry.testing_utils
+teardown = mhcflurry.testing_utils.module_cleanup
+
+
 DOWNLOADED_PREDICTOR = Class1AffinityPredictor.load()
 
 logging.basicConfig(level=logging.DEBUG)
@@ -252,3 +256,4 @@ def test_predict_implementations_equivalent():
                 peptides=peptides,
                 centrality_measure=centrality_measure).prediction.values
             testing.assert_almost_equal(pred1, pred2, decimal=2)
+
diff --git a/test/test_class1_neural_network.py b/test/test_class1_neural_network.py
index 8baaded0..e614d602 100644
--- a/test/test_class1_neural_network.py
+++ b/test/test_class1_neural_network.py
@@ -13,6 +13,9 @@ from mhcflurry.class1_neural_network import Class1NeuralNetwork
 from mhcflurry.downloads import get_path
 from mhcflurry.common import random_peptides
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 
 def test_class1_neural_network_a0205_training_accuracy():
     # Memorize the dataset.
diff --git a/test/test_class1_pan.py b/test/test_class1_pan.py
index 76a7a929..cdb483cf 100644
--- a/test/test_class1_pan.py
+++ b/test/test_class1_pan.py
@@ -11,6 +11,9 @@ from mhcflurry import Class1AffinityPredictor,Class1NeuralNetwork
 from mhcflurry.allele_encoding import AlleleEncoding
 from mhcflurry.downloads import get_path
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 
 HYPERPARAMETERS = {
     'activation': 'tanh',
diff --git a/test/test_custom_loss.py b/test/test_custom_loss.py
index 55bfa7b5..4c177666 100644
--- a/test/test_custom_loss.py
+++ b/test/test_custom_loss.py
@@ -11,6 +11,9 @@ import keras.backend as K
 
 from mhcflurry.custom_loss import CUSTOM_LOSSES
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 
 def evaluate_loss(loss, y_true, y_pred):
     y_true = numpy.array(y_true)
diff --git a/test/test_download_models_class1.py b/test/test_download_models_class1.py
index be36e44f..23b25c3d 100644
--- a/test/test_download_models_class1.py
+++ b/test/test_download_models_class1.py
@@ -5,6 +5,9 @@ from numpy.testing import assert_equal
 
 from mhcflurry import Class1AffinityPredictor, Class1NeuralNetwork
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 
 DOWNLOADED_PREDICTOR = Class1AffinityPredictor.load()
 
diff --git a/test/test_multi_output.py b/test/test_multi_output.py
index 2e126154..efbcccf0 100644
--- a/test/test_multi_output.py
+++ b/test/test_multi_output.py
@@ -12,6 +12,9 @@ logging.getLogger('tensorflow').disabled = True
 from mhcflurry.class1_neural_network import Class1NeuralNetwork
 from mhcflurry.common import random_peptides
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 
 def test_multi_output():
     hyperparameters = dict(
diff --git a/test/test_network_merging.py b/test/test_network_merging.py
index 6021652b..a180df49 100644
--- a/test/test_network_merging.py
+++ b/test/test_network_merging.py
@@ -5,6 +5,9 @@ from mhcflurry import Class1AffinityPredictor, Class1NeuralNetwork
 from mhcflurry.common import random_peptides
 from mhcflurry.downloads import get_path
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 logging.getLogger('tensorflow').disabled = True
 
 PAN_ALLELE_PREDICTOR = Class1AffinityPredictor.load(
diff --git a/test/test_predict_command.py b/test/test_predict_command.py
index b85e7202..7d4d9340 100644
--- a/test/test_predict_command.py
+++ b/test/test_predict_command.py
@@ -6,6 +6,9 @@ from numpy.testing import assert_equal
 
 from mhcflurry import predict_command
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 TEST_CSV = '''
 Allele,Peptide,Experiment
 HLA-A0201,SYNFEKKL,17
diff --git a/test/test_released_predictors_on_hpv_dataset.py b/test/test_released_predictors_on_hpv_dataset.py
index 525f06c2..10f0fc23 100644
--- a/test/test_released_predictors_on_hpv_dataset.py
+++ b/test/test_released_predictors_on_hpv_dataset.py
@@ -13,6 +13,9 @@ from nose.tools import eq_, assert_less, assert_greater, assert_almost_equal
 from mhcflurry import Class1AffinityPredictor
 from mhcflurry.downloads import get_path
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 
 def data_path(name):
     '''
diff --git a/test/test_released_predictors_well_correlated.py b/test/test_released_predictors_well_correlated.py
index 395a6e80..6b8a6a7e 100644
--- a/test/test_released_predictors_well_correlated.py
+++ b/test/test_released_predictors_well_correlated.py
@@ -15,6 +15,9 @@ from mhcflurry.encodable_sequences import EncodableSequences
 from mhcflurry.downloads import get_path
 from mhcflurry.common import random_peptides
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 
 PREDICTORS = {
     'allele-specific': Class1AffinityPredictor.load(
diff --git a/test/test_speed.py b/test/test_speed.py
index e9c342a6..b35327fe 100644
--- a/test/test_speed.py
+++ b/test/test_speed.py
@@ -18,6 +18,9 @@ from mhcflurry.encodable_sequences import EncodableSequences
 from mhcflurry.common import random_peptides
 from mhcflurry.downloads import get_path
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 ALLELE_SPECIFIC_PREDICTOR = Class1AffinityPredictor.load(
     get_path("models_class1", "models"))
 
diff --git a/test/test_train_and_related_commands.py b/test/test_train_and_related_commands.py
index f6abfd26..09ae60e5 100644
--- a/test/test_train_and_related_commands.py
+++ b/test/test_train_and_related_commands.py
@@ -14,6 +14,9 @@ from numpy.testing import assert_array_less, assert_equal
 from mhcflurry import Class1AffinityPredictor
 from mhcflurry.downloads import get_path
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
 HYPERPARAMETERS = [
diff --git a/test/test_train_pan_allele_models_command.py b/test/test_train_pan_allele_models_command.py
index ea5261ec..dd69e810 100644
--- a/test/test_train_pan_allele_models_command.py
+++ b/test/test_train_pan_allele_models_command.py
@@ -15,6 +15,9 @@ from numpy.testing import assert_equal, assert_array_less
 from mhcflurry import Class1AffinityPredictor,Class1NeuralNetwork
 from mhcflurry.downloads import get_path
 
+from mhcflurry.testing_utils import module_cleanup
+teardown = module_cleanup
+
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
 
-- 
GitLab