Fix tests for Python 2

astrofrog · Oct 31, 2014 · 110965b · 110965b
1 parent 9ed1f7f
commit 110965b
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/sedfitter/tests/test_full_pipeline.py b/sedfitter/tests/test_full_pipeline.py
@@ -116,7 +116,13 @@ class BasePipelineTest(object):
     def setup_class(self):
 
         self.tmpdir = tempfile.mkdtemp()
-        self.model_generator(self.tmpdir, aperture_dependent=self.aperture_dependent)
+
+        if self.model_version == 1:
+            model_generator = generate_random_models_1
+        else:
+            model_generator = generate_random_models_2
+
+        model_generator(self.tmpdir, aperture_dependent=self.aperture_dependent)
 
         from ..extinction import Extinction
 
@@ -286,25 +292,21 @@ class TestApertureIndependentPipeline1(BasePipelineTest):
 
     aperture_dependent = False
     model_version = 1
-    model_generator = generate_random_models_1
 
 
 class TestApertureIndependentPipeline2(BasePipelineTest):
 
     aperture_dependent = False
     model_version = 2
-    model_generator = generate_random_models_2
 
 
 class TestApertureDependentPipeline1(BasePipelineTest):
 
     aperture_dependent = True
     model_version = 1
-    model_generator = generate_random_models_1
 
 
 class TestApertureDependentPipeline2(BasePipelineTest):
 
     aperture_dependent = True
     model_version = 2
-    model_generator = generate_random_models_2