tests/__main__.py: wrap TestProgram

Instead of parsing the arguments twice, wrap unittest.main
and inject our custom argument into the existing parser.
This commit is contained in:
Michele Guerini Rocco 2025-03-21 13:29:11 +01:00
parent 20a011414a
commit 09724908d1
Signed by: rnhmjoj
GPG Key ID: BFBAF4C975F76450
4 changed files with 49 additions and 42 deletions

View File

@ -6,7 +6,7 @@ general astigmatic beam
'''
from .. import TestCase, GrayTest
from .. import options, load_table
from .. import load_table
import numpy as np
@ -128,7 +128,7 @@ class Test(GrayTest, TestCase):
S_true = eikonal_true(k0, W, K, φ_w, φ_k, s0, data)
S_gray = eikonal_gray(k0, W, K, φ_w, φ_k, data)
if options.visual:
if self.options.visual:
import matplotlib.pyplot as plt
plt.suptitle(self.__module__ + '.test_eikonal')

View File

@ -2,7 +2,7 @@
Combine EC profiles from three independent beams
'''
from .. import TestCase, options, load_table, run_gray
from .. import TestCase, load_table, run_gray
from pathlib import Path
@ -31,13 +31,13 @@ class Test(TestCase):
cls.candidate = Path(cls._tempdir)
# replace reference with candidate
if options.update:
if cls.options.update:
print()
print('Setting new reference for ' + cls.__module__)
cls.candidate = cls.reference
# run gray to generate the candidate outputs
proc = run_gray(cls.inputs, cls.candidate, binary=options.binary,
proc = run_gray(cls.inputs, cls.candidate, binary=cls.options.binary,
options=['-s', cls.inputs / 'filelist.txt'])
assert proc.returncode == 0, \
f"gray failed with exit code {proc.returncode}"
@ -52,7 +52,7 @@ class Test(TestCase):
Clean up after all tests
'''
# remove temporary directory
if cls._passed or not options.keep_failed:
if cls._passed or not cls.options.keep_failed:
shutil.rmtree(cls._tempdir)
else:
print()
@ -95,7 +95,7 @@ class Test(TestCase):
ref = load_table(self.reference / 'sum-ec-profiles.txt')
cand = load_table(self.candidate / 'sum-ec-profiles.txt')
if options.visual:
if self.options.visual:
import matplotlib.pyplot as plt
fig, axes = plt.subplots(3, 1, sharex=True)

View File

@ -20,6 +20,9 @@ class GrayTest:
reference: Path = None # directory of the reference outputs
candidate: Path = None # directory of the candidate outputs
# command line options
options: argparse.Namespace = argparse.Namespace()
# output tables to save
tables: list[int] = [4, 7, 8, 9, 48, 56, 33, 70, 71]
@ -44,14 +47,14 @@ class GrayTest:
cls.candidate = Path(cls._tempdir)
# replace reference with candidate
if options.update:
if cls.options.update:
print()
print('Setting new reference for ' + cls.__module__)
cls.candidate = cls.reference
# run gray to generate the candidate outputs
proc = run_gray(cls.inputs, cls.candidate, params=cls.gray_params,
binary=options.binary, tables=cls.tables)
binary=cls.options.binary, tables=cls.tables)
# 0: all good, 1: input errors, >1: simulation errors
assert proc.returncode != 1, 'gray failed with exit code 1'
@ -66,7 +69,7 @@ class GrayTest:
Clean up after all tests
'''
# remove temporary directory
if cls._passed or not options.keep_failed:
if cls._passed or not cls.options.keep_failed:
shutil.rmtree(cls._tempdir)
else:
print()
@ -134,7 +137,7 @@ class GrayTest:
dist = emd(ref_beam['ρ_t'], cand_beam['ρ_t'], y1, y2)
self.assertLess(dist, 0.001, f'{val} profile changed')
if options.visual:
if self.options.visual:
for index_rt in beams:
ref_beam = ref[ref['index_rt'] == index_rt]
cand_beam = cand[cand['index_rt'] == index_rt]
@ -190,7 +193,7 @@ class GrayTest:
self.assertAlmostEqual(ref_val, cand_val, prec[col],
msg=f"{col} at line {line} changed")
if options.visual:
if self.options.visual:
fig, axes = plt.subplots(4, 3, tight_layout=True)
fig.suptitle(self.__module__ + '.test_flux_averages')
@ -241,7 +244,7 @@ class GrayTest:
except FileNotFoundError:
raise unittest.SkipTest("Beam shape info not available")
if options.visual:
if self.options.visual:
plt.subplot(aspect='equal')
plt.title(self.__module__ + '.test_beam_shape')
plt.xlabel('$x$ / cm')
@ -274,7 +277,7 @@ class GrayTest:
if data.size < 2:
self.skipTest("There is no plasma")
if options.visual:
if self.options.visual:
left = plt.subplot()
plt.title(self.__module__ + '.test_error_biased')
left.set_xlabel('$s$ / cm')
@ -315,10 +318,6 @@ class GrayTest:
self.assertGreater(χ2, 1)
# Command line options
options = argparse.Namespace()
def get_basedir(module: str) -> Path:
"""
Given a module name (es. tests.03-TCV) returns its

View File

@ -1,4 +1,3 @@
import argparse
import unittest
import tests
@ -18,27 +17,36 @@ def get_cases(suite):
return {f'{mod}.{cls}'}
if __name__ == '__main__':
# handle custom CLI arguments
parser = argparse.ArgumentParser(prog='tests.main')
parser.add_argument('--keep-failed', action='store_true',
help='whether to keep output files of failed tests')
parser.add_argument('--visual', action='store_true',
help='whether to show plots, helps with debugging')
parser.add_argument('--update', action='store_true',
help='update the reference results with the'
'current ones')
parser.add_argument('--binary', type=str, default='build/bin/gray',
help='the gray binary to be tested')
parser.add_argument('--list-cases', action='store_true',
help='only list the test cases')
options, other_args = parser.parse_known_args()
tests.options = options
class GrayTestProgram(unittest.TestProgram):
def _initArgParsers(self):
# add custom CLI arguments
super()._initArgParsers()
p = self._main_parser
p.add_argument('--keep-failed', action='store_true',
help='whether to keep output files of failed tests')
p.add_argument('--visual', action='store_true',
help='whether to show plots, helps with debugging')
p.add_argument('--update', action='store_true',
help='update the reference results with the'
'current ones')
p.add_argument('--binary', type=str, default='build/bin/gray',
help='the gray binary to be tested')
p.add_argument('--list-cases', action='store_true',
help='only list the test cases')
if options.list_cases:
# list all test cases
suite = unittest.TestLoader().discover('.')
print(*sorted(get_cases(suite)), sep='\n')
else:
# start the test runner
unittest.main(module=None, argv=["tests.main"] + other_args)
def parseArgs(self, argv):
super().parseArgs(argv)
# store args in the test classes
tests.GrayTest.options = self
tests.TestCase.options = self
# handle custom harguments
if self.list_cases:
print(*sorted(get_cases(self.test)), sep='\n')
exit()
if __name__ == '__main__':
# start the test runner
GrayTestProgram(module=None)