From 27b5c05900c8200d41e5d71ab372bb89919bcf93 Mon Sep 17 00:00:00 2001 From: slopqueue Date: Sat, 8 Aug 2026 19:45:30 +0000 Subject: [PATCH] Parameterize data-based tests and make data path configurable Add pytest command-line options --data-path/--code-path so the shared data and code roots can be passed on the command line with fallback to the DATA/CODE environment variables. Parametrize the EQDSK and efit_to_boozer tests over the standard reference data files so each is a distinct test case, and remove the module-level invocation from the standalone NEO-RT benchmark util test so it runs once under pytest. --- .../NEO-RT/benchmark_with_NEO_2/test_util.py | 2 - tests/conftest.py | 40 ++++++-- tests/libneo/python/test_efit_to_boozer.py | 92 +++++++++---------- tests/libneo/python/test_eqdsk.py | 34 ++++--- 4 files changed, 103 insertions(+), 65 deletions(-) diff --git a/tests/NEO-RT/benchmark_with_NEO_2/test_util.py b/tests/NEO-RT/benchmark_with_NEO_2/test_util.py index 6763dfe..0426d7a 100644 --- a/tests/NEO-RT/benchmark_with_NEO_2/test_util.py +++ b/tests/NEO-RT/benchmark_with_NEO_2/test_util.py @@ -7,5 +7,3 @@ def test_replace_template(): } text = fill_template(template, vars) assert(text == 'bla 1 blu bli') - -test_replace_template() diff --git a/tests/conftest.py b/tests/conftest.py index 035d0df..99e633f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,12 +1,40 @@ -import pytest - import os from pathlib import Path +import pytest + + +def pytest_addoption(parser): + parser.addoption( + "--data-path", + action="store", + default=None, + help="Path to the shared data directory (defaults to the DATA environment variable)", + ) + parser.addoption( + "--code-path", + action="store", + default=None, + help="Path to the CODE workspace root (defaults to the CODE environment variable)", + ) + + +def _resolve_path(option_value, env_var, option_name): + if option_value: + return Path(option_value) + if env_var in os.environ: + return Path(os.environ[env_var]) + raise pytest.UsageError( + f"Neither --{option_name} nor the {env_var} environment variable is set. " + f"Pass --{option_name} on the command line or export {env_var}." + ) + + @pytest.fixture -def code_path(): - return Path(os.environ["CODE"]) +def code_path(request): + return _resolve_path(request.config.getoption("--code-path"), "CODE", "code-path") + @pytest.fixture -def data_path(): - return Path(os.environ["DATA"]) +def data_path(request): + return _resolve_path(request.config.getoption("--data-path"), "DATA", "data-path") diff --git a/tests/libneo/python/test_efit_to_boozer.py b/tests/libneo/python/test_efit_to_boozer.py index bfd2c05..ce656b5 100644 --- a/tests/libneo/python/test_efit_to_boozer.py +++ b/tests/libneo/python/test_efit_to_boozer.py @@ -46,52 +46,52 @@ def test_files(code_path, data_path): @pytest.mark.slow -def test_q_profile_eqdsk(test_files): - for key in ["DEMO CHEASE", "MASTU"]: - # TODO: local, AUG, MASTU, PROCESS, standardized - test_file = test_files[key] - print(f"Testing {key} EQDSK file: {test_file}") - tmp_path = init_run_path(test_file) - print(f"Running in {tmp_path}") - os.chdir(tmp_path) - eqdsk_data = read_eqdsk(str(test_file)) - - # The data in EQDSK is writting in poloidal flux label, - # but efit_to_boozer uses toroidal flux label. Therefore, we need a mapping - q_profile_eqdsk = eqdsk_data["qprof"] - spol_eqdsk = np.linspace(0.0, 1.0, q_profile_eqdsk.shape[0]) - converter = FluxConverter(q_profile_eqdsk) - stor_eqdsk = converter.spol2stor(spol_eqdsk) - - # The safety factor can be calculated alternatively using field line integration - # this is done as part of the symmetry flux transformation and given as output here - - q_profile_field_line_integration = [] - - # inp_label = 1 makes it that efit_to_boozer.magdata_in_symfluxcoord_ext uses the - # flux label si instead of the actual flux psi to determine which - # flux surface one is on. Therefore psi is just set as a dummy variable here. - inp_label = 1 - psi = np.array(0.0) - theta = np.array(0.0) - - efit_to_boozer.efit_to_boozer.init() - for si in stor_eqdsk: - (q, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _) = ( - efit_to_boozer.magdata_in_symfluxcoord_ext(inp_label, si, psi, theta) - ) - q_profile_field_line_integration.append(q) - - q_profile_field_line_integration = np.array(q_profile_field_line_integration) - - plt.figure() - plt.plot(spol_eqdsk, q_profile_eqdsk, "-r", label=r"q profile from EQDSK") - plt.plot(spol_eqdsk, q_profile_field_line_integration, "--b", label=r"q profile from field line integration") - plt.title(f"{key}") - plt.legend() - - assert_allclose(q_profile_field_line_integration, q_profile_eqdsk, rtol=1e-2) - print("Alternative safety factor calculation agrees with EQDSK file within 1%") +@pytest.mark.parametrize("key", ["DEMO CHEASE", "MASTU"]) +def test_q_profile_eqdsk(key, test_files): + # TODO: local, AUG, PROCESS, standardized + test_file = test_files[key] + print(f"Testing {key} EQDSK file: {test_file}") + tmp_path = init_run_path(test_file) + print(f"Running in {tmp_path}") + os.chdir(tmp_path) + eqdsk_data = read_eqdsk(str(test_file)) + + # The data in EQDSK is writting in poloidal flux label, + # but efit_to_boozer uses toroidal flux label. Therefore, we need a mapping + q_profile_eqdsk = eqdsk_data["qprof"] + spol_eqdsk = np.linspace(0.0, 1.0, q_profile_eqdsk.shape[0]) + converter = FluxConverter(q_profile_eqdsk) + stor_eqdsk = converter.spol2stor(spol_eqdsk) + + # The safety factor can be calculated alternatively using field line integration + # this is done as part of the symmetry flux transformation and given as output here + + q_profile_field_line_integration = [] + + # inp_label = 1 makes it that efit_to_boozer.magdata_in_symfluxcoord_ext uses the + # flux label si instead of the actual flux psi to determine which + # flux surface one is on. Therefore psi is just set as a dummy variable here. + inp_label = 1 + psi = np.array(0.0) + theta = np.array(0.0) + + efit_to_boozer.efit_to_boozer.init() + for si in stor_eqdsk: + (q, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _) = ( + efit_to_boozer.magdata_in_symfluxcoord_ext(inp_label, si, psi, theta) + ) + q_profile_field_line_integration.append(q) + + q_profile_field_line_integration = np.array(q_profile_field_line_integration) + + plt.figure() + plt.plot(spol_eqdsk, q_profile_eqdsk, "-r", label=r"q profile from EQDSK") + plt.plot(spol_eqdsk, q_profile_field_line_integration, "--b", label=r"q profile from field line integration") + plt.title(f"{key}") + plt.legend() + + assert_allclose(q_profile_field_line_integration, q_profile_eqdsk, rtol=1e-2) + print("Alternative safety factor calculation agrees with EQDSK file within 1%") def init_run_path(gfile): diff --git a/tests/libneo/python/test_eqdsk.py b/tests/libneo/python/test_eqdsk.py index 8272d83..e344d30 100644 --- a/tests/libneo/python/test_eqdsk.py +++ b/tests/libneo/python/test_eqdsk.py @@ -9,6 +9,16 @@ from libneo import eqdsk +# The set of standard reference files each test is run against. +STANDARD_TEST_FILES = [ + "local", + "PROCESS", + "standardized", + "AUG", + # TODO: "MASTU", + # TODO: "CHEASE", +] + @pytest.fixture def test_files(code_path, data_path): @@ -22,24 +32,26 @@ def test_files(code_path, data_path): } -def test_eqdsk_read(test_files): - for key, test_file in test_files.items(): - print(f"Testing {key} EQDSK file: {test_file}") - _ = eqdsk.eqdsk_file(test_file) +@pytest.mark.parametrize("key", STANDARD_TEST_FILES) +def test_eqdsk_read(key, test_files): + test_file = test_files[key] + print(f"Testing {key} EQDSK file: {test_file}") + _ = eqdsk.eqdsk_file(test_file) @pytest.mark.slow -def test_eqdsk_golden_records(data_path, test_files): +@pytest.mark.parametrize("key", STANDARD_TEST_FILES) +def test_eqdsk_golden_records(key, data_path, test_files): golden_record_path = data_path / "TESTS/libneo/eqdsk" - store_golden_records(test_files.values(), golden_record_path) + test_file = test_files[key] + store_golden_records([test_file], golden_record_path) - for test_file in test_files.values(): - eqdsk_object = eqdsk.eqdsk_file(test_file) + eqdsk_object = eqdsk.eqdsk_file(test_file) - data = eqdsk_object.__dict__ - replace_array_members_by_lists(data) + data = eqdsk_object.__dict__ + replace_array_members_by_lists(data) - assert are_dicts_equal(data, get_golden_record(test_file, golden_record_path)) + assert are_dicts_equal(data, get_golden_record(test_file, golden_record_path)) def get_golden_record(file_path, storage_path):