update tests for newrunner

jdebacker · Jun 21, 2021 · dd0f7a9 · dd0f7a9
1 parent 9d3252f
commit dd0f7a9
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 26 deletions.
diff --git a/ogusa/tests/test_basic.py b/ogusa/tests/test_basic.py
@@ -33,10 +33,11 @@ def test_run_small(time_path, dask_client):
     TPI.ENFORCE_SOLUTION_CHECKS = False
     SS.MINIMIZER_TOL = 1e-6
     TPI.MINIMIZER_TOL = 1e-6
-    og_spec = TEST_PARAM_DICT
-    runner(output_base=OUTPUT_DIR, baseline_dir=OUTPUT_DIR,
-           time_path=time_path, baseline=True, og_spec=og_spec,
-           client=dask_client, num_workers=NUM_WORKERS)
+    p = Specifications(baseline=True, client=dask_client,
+                       num_workers=NUM_WORKERS, baseline_dir=OUTPUT_DIR,
+                       output_base=OUTPUT_DIR)
+    p.update_specifications(TEST_PARAM_DICT)
+    runner(p, time_path=time_path, client=dask_client)
 
 
 @pytest.mark.local

diff --git a/ogusa/tests/test_run_ogusa.py b/ogusa/tests/test_run_ogusa.py
@@ -5,6 +5,7 @@
 import time
 import os
 from ogusa.execute import runner
+from ogusa.parameters import Specifications
 from ogusa.utils import safe_read_pickle
 import ogusa.output_tables as ot
 SS.ENFORCE_SOLUTION_CHECKS = False
@@ -40,25 +41,22 @@ def run_micro_macro(og_spec, guid, client):
         Run baseline
     ------------------------------------------------------------------------
     '''
-    output_base = BASELINE_DIR
-    kwargs = {'output_base': output_base, 'baseline_dir': BASELINE_DIR,
-              'time_path': True, 'baseline': True,
-              'og_spec': og_spec, 'guid': guid, 'client': client,
-              'num_workers': NUM_WORKERS}
-    runner(**kwargs)
+    p = Specifications(baseline=True, client=client,
+                       num_workers=NUM_WORKERS, baseline_dir=BASELINE_DIR,
+                       output_base=BASELINE_DIR)
+    p.update_specifications(og_spec)
+    runner(p, time_path=True, client=client)
 
     '''
     ------------------------------------------------------------------------
         Run reform
     ------------------------------------------------------------------------
     '''
-
-    output_base = REFORM_DIR
-    kwargs = {'output_base': output_base, 'baseline_dir': BASELINE_DIR,
-              'time_path': True, 'baseline': False,
-              'og_spec': og_spec, 'guid': guid, 'client': client,
-              'num_workers': NUM_WORKERS}
-    runner(**kwargs)
+    p = Specifications(baseline=False, client=client,
+                       num_workers=NUM_WORKERS, baseline_dir=BASELINE_DIR,
+                       output_base=REFORM_DIR)
+    p.update_specifications(og_spec)
+    runner(p, time_path=True, client=client)
     time.sleep(0.5)
     base_tpi = safe_read_pickle(
         os.path.join(BASELINE_DIR, 'TPI', 'TPI_vars.pkl'))

diff --git a/ogusa/tests/test_user_inputs.py b/ogusa/tests/test_user_inputs.py
@@ -3,6 +3,7 @@
 import pytest
 import os
 from ogusa.execute import runner
+from ogusa.parameters import Specifications
 NUM_WORKERS = min(multiprocessing.cpu_count(), 7)
 CUR_PATH = os.path.abspath(os.path.dirname(__file__))
 TAX_FUNC_PATH = os.path.join(
@@ -25,9 +26,11 @@ def dask_client():
                          ids=['Frisch 0.32', 'Frisch 0.4', 'Frisch 0.6'])
 def test_frisch(frisch, dask_client):
     og_spec = {'frisch': frisch, 'debt_ratio_ss': 1.0}
-    runner(output_base=OUTPUT_DIR, baseline_dir=OUTPUT_DIR,
-           time_path=False, baseline=True, og_spec=og_spec,
-           client=dask_client, num_workers=NUM_WORKERS)
+    p = Specifications(baseline=True, client=dask_client,
+                       num_workers=NUM_WORKERS, baseline_dir=OUTPUT_DIR,
+                       output_base=OUTPUT_DIR)
+    p.update_specifications(og_spec)
+    runner(p, time_path=False, client=dask_client)
 
 
 @pytest.mark.local
@@ -36,9 +39,11 @@ def test_frisch(frisch, dask_client):
 def test_gy(g_y_annual, dask_client):
     og_spec = {'frisch': 0.41, 'debt_ratio_ss': 1.0,
                'g_y_annual': g_y_annual}
-    runner(output_base=OUTPUT_DIR, baseline_dir=OUTPUT_DIR,
-           time_path=False, baseline=True, og_spec=og_spec,
-           client=dask_client, num_workers=NUM_WORKERS)
+    p = Specifications(baseline=True, client=dask_client,
+                       num_workers=NUM_WORKERS, baseline_dir=OUTPUT_DIR,
+                       output_base=OUTPUT_DIR)
+    p.update_specifications(og_spec)
+    runner(p, time_path=False, client=dask_client)
 
 
 @pytest.mark.local
@@ -47,6 +52,8 @@ def test_gy(g_y_annual, dask_client):
                               'sigma=1.9'])
 def test_sigma(sigma, dask_client):
     og_spec = {'frisch': 0.41, 'debt_ratio_ss': 1.0, 'sigma': sigma}
-    runner(output_base=OUTPUT_DIR, baseline_dir=OUTPUT_DIR,
-           time_path=False, baseline=True, og_spec=og_spec,
-           client=dask_client, num_workers=NUM_WORKERS)
+    p = Specifications(baseline=True, client=dask_client,
+                       num_workers=NUM_WORKERS, baseline_dir=OUTPUT_DIR,
+                       output_base=OUTPUT_DIR)
+    p.update_specifications(og_spec)
+    runner(p, time_path=False, client=dask_client)