Replace unittest by pytest

Author: druzm

tests/test_calibrator/calibrator_solutions.py

@@ -30,8 +30,10 @@
 =====================
 """
 # Default model
-solution_default_df = pd.read_csv(cal_sol_default_path, float_precision='round_trip', converters=calibrator_sol_converters)
-solution_custom_df = pd.read_csv(cal_sol_custom_path, float_precision='round_trip', converters=calibrator_sol_converters)
+solution_default_df = pd.read_csv(cal_sol_default_path, float_precision='round_trip',
+                                  converters=calibrator_sol_converters)
+solution_custom_df = pd.read_csv(cal_sol_custom_path, float_precision='round_trip',
+                                 converters=calibrator_sol_converters)
 # v211w model
 solution_v211w_default_df = pd.read_csv(cal_sol_v211w_default_path, float_precision='round_trip',
                                         converters=calibrator_sol_converters)

tests/test_calibrator/test_calibrator.py

@@ -1,7 +1,6 @@
-import unittest
-
 import numpy as np
 import numpy.testing as npt
+import pytest
 from pandas import testing as pdt
 
 from gaiaxpy import calibrate
@@ -12,11 +11,12 @@
 from gaiaxpy.input_reader.required_columns import MANDATORY_INPUT_COLS, CORR_INPUT_COLUMNS
 from gaiaxpy.spectrum.absolute_sampled_spectrum import AbsoluteSampledSpectrum
 from gaiaxpy.spectrum.sampled_basis_functions import SampledBasisFunctions
-from tests.files.paths import mean_spectrum_csv_file, mean_spectrum_avro_file, mean_spectrum_fits_file, \
-    mean_spectrum_xml_file, mean_spectrum_xml_plain_file, mean_spectrum_ecsv_file
-from tests.test_calibrator.calibrator_solutions import solution_default_df, solution_custom_df, \
-    solution_v211w_default_df, solution_v211w_custom_df, sol_custom_sampling_array, sol_v211w_default_sampling_array, \
-    sol_default_sampling_array
+from tests.files.paths import (mean_spectrum_csv_file, mean_spectrum_avro_file, mean_spectrum_fits_file,
+                               mean_spectrum_xml_file, mean_spectrum_xml_plain_file, mean_spectrum_ecsv_file)
+from tests.test_calibrator.calibrator_solutions import (solution_default_df, solution_custom_df,
+                                                        solution_v211w_default_df, solution_v211w_custom_df,
+                                                        sol_custom_sampling_array, sol_v211w_default_sampling_array,
+                                                        sol_default_sampling_array)
 from tests.utils.utils import is_instance_err_message, npt_array_err_message
 
 # Load variables
@@ -32,67 +32,72 @@
                    mean_spectrum_xml_file, mean_spectrum_xml_plain_file]
 
 
-def generate_single_spectrum(mean_spectrum_path):
-    # Read mean Spectrum
-    parser = InternalContinuousParser(MANDATORY_INPUT_COLS['calibrate'] + CORR_INPUT_COLUMNS)
-    parsed_spectrum_file, extension = parser.parse_file(mean_spectrum_path)
-    # Create sampled basis functions
-    sampled_basis_func = {band: SampledBasisFunctions.from_design_matrix(xp_sampling_grid, xp_design_matrices[band])
-                          for band in BANDS}
-    first_row = parsed_spectrum_file.iloc[0]
-    return _create_spectrum(first_row, truncation=False, design_matrix=sampled_basis_func, merge=xp_merge)
-
-
-class TestCalibrator(unittest.TestCase):
-
-    def test_create_spectrum(self):
-        for file in cal_input_files:
-            spectrum = generate_single_spectrum(file)
-            instance = AbsoluteSampledSpectrum
-            self.assertIsInstance(spectrum, instance, msg=is_instance_err_message(file, instance))
-
-    def test_calibrate_both_bands_default_calibration_model(self):
-        # Default sampling and default calibration sampling
-        for file in cal_input_files:
-            spectra_df_csv, positions = calibrate(file, save_file=False)
-            npt.assert_array_equal(positions, sol_default_sampling_array, err_msg=npt_array_err_message(file))
-            # Assert_frame_equal doesn't seem to have a parameter to print an error message with details.
-            pdt.assert_frame_equal(spectra_df_csv, solution_default_df, atol=_atol, rtol=_rtol)
-
-    def test_custom_sampling_default_calibration_model(self):
-        for file in cal_input_files:
-            spectra_df_custom_sampling, positions = calibrate(file, sampling=np.arange(350, 1050, 200), save_file=False)
-            npt.assert_array_equal(positions, sol_custom_sampling_array, err_msg=npt_array_err_message(file))
-            pdt.assert_frame_equal(spectra_df_custom_sampling, solution_custom_df, atol=_atol, rtol=_rtol)
-
-    def test_calibrate_both_bands_v211w_model(self):
-        for file in cal_input_files:
-            spectra_df_csv, positions = _calibrate(file, save_file=False, bp_model=bp_model)
-            npt.assert_array_equal(positions, sol_v211w_default_sampling_array, err_msg=npt_array_err_message(file))
-            pdt.assert_frame_equal(spectra_df_csv, solution_v211w_default_df, atol=_atol, rtol=_rtol)
-
-    def test_custom_sampling_v211w_model(self):
-        for file in cal_input_files:
-            spectra_df_custom_sampling, positions = _calibrate(file, sampling=np.arange(350, 1050, 200),
-                                                               save_file=False, bp_model=bp_model)
-            npt.assert_array_equal(positions, sol_custom_sampling_array, err_msg=npt_array_err_message(file))
-            pdt.assert_frame_equal(spectra_df_custom_sampling, solution_v211w_custom_df, atol=_atol, rtol=_rtol)
-
-
-class TestCalibratorSamplingRange(unittest.TestCase):
-
-    def test_sampling_wrong_array(self):
-        with self.assertRaises(ValueError):
-            calibrate(mean_spectrum_avro_file, sampling=np.linspace(800, 350, 600), save_file=False)
-
-    def test_sampling_low(self):
-        with self.assertRaises(ValueError):
-            calibrate(mean_spectrum_avro_file, sampling=np.linspace(300, 850, 300), save_file=False)
-
-    def test_sampling_high(self):
-        with self.assertRaises(ValueError):
-            calibrate(mean_spectrum_avro_file, sampling=np.linspace(350, 1500, 200), save_file=False)
-
-    def test_sampling_both_wrong(self):
-        with self.assertRaises(ValueError):
-            calibrate(mean_spectrum_avro_file, sampling=np.linspace(200, 2000, 100), save_file=False)
+@pytest.mark.parametrize('input_file', cal_input_files)
+def test_create_spectrum(input_file):
+    def generate_single_spectrum(mean_spectrum_path):
+        # Read mean Spectrum
+        parser = InternalContinuousParser(MANDATORY_INPUT_COLS['calibrate'] + CORR_INPUT_COLUMNS)
+        parsed_spectrum_file, extension = parser.parse_file(mean_spectrum_path)
+        # Create sampled basis functions
+        sampled_basis_func = {band: SampledBasisFunctions.from_design_matrix(xp_sampling_grid, xp_design_matrices[band])
+                              for band in BANDS}
+        return _create_spectrum(parsed_spectrum_file.iloc[0], truncation=False, design_matrix=sampled_basis_func,
+                                merge=xp_merge)
+
+    instance = AbsoluteSampledSpectrum
+    spectrum = generate_single_spectrum(input_file)
+    assert isinstance(spectrum, instance), is_instance_err_message(input_file, instance)
+
+
+@pytest.mark.parametrize('input_file', cal_input_files)
+def test_calibrate_both_bands_default_calibration_model(input_file, request):
+    # Default sampling and default calibration sampling
+    spectra_df_csv, positions = calibrate(input_file, save_file=False)
+    npt.assert_array_equal(positions, sol_default_sampling_array, err_msg=npt_array_err_message(input_file))
+    try:
+        pdt.assert_frame_equal(spectra_df_csv, solution_default_df, atol=_atol, rtol=_rtol)
+    except AssertionError as e:
+        print(f'{request.node.name} failed for file: {input_file}.')
+        raise e
+
+
+@pytest.mark.parametrize('input_file', cal_input_files)
+def test_custom_sampling_default_calibration_model(input_file):
+    spectra_df_custom_sampling, positions = calibrate(input_file, sampling=np.arange(350, 1050, 200), save_file=False)
+    npt.assert_array_equal(positions, sol_custom_sampling_array, err_msg=npt_array_err_message(input_file))
+    pdt.assert_frame_equal(spectra_df_custom_sampling, solution_custom_df, atol=_atol, rtol=_rtol)
+
+
+@pytest.mark.parametrize('input_file', cal_input_files)
+def test_calibrate_both_bands_v211w_model(input_file):
+    spectra_df_csv, positions = _calibrate(input_file, save_file=False, bp_model=bp_model)
+    npt.assert_array_equal(positions, sol_v211w_default_sampling_array, err_msg=npt_array_err_message(input_file))
+    pdt.assert_frame_equal(spectra_df_csv, solution_v211w_default_df, atol=_atol, rtol=_rtol)
+
+
+@pytest.mark.parametrize('input_file', cal_input_files)
+def test_custom_sampling_v211w_model(input_file):
+    spectra_df_custom_sampling, positions = _calibrate(input_file, sampling=np.arange(350, 1050, 200), save_file=False,
+                                                       bp_model=bp_model)
+    npt.assert_array_equal(positions, sol_custom_sampling_array, err_msg=npt_array_err_message(input_file))
+    pdt.assert_frame_equal(spectra_df_custom_sampling, solution_v211w_custom_df, atol=_atol, rtol=_rtol)
+
+
+def test_sampling_wrong_array():
+    with pytest.raises(ValueError):
+        calibrate(mean_spectrum_avro_file, sampling=np.linspace(800, 350, 600), save_file=False)
+
+
+def test_sampling_low():
+    with pytest.raises(ValueError):
+        calibrate(mean_spectrum_avro_file, sampling=np.linspace(300, 850, 300), save_file=False)
+
+
+def test_sampling_high():
+    with pytest.raises(ValueError):
+        calibrate(mean_spectrum_avro_file, sampling=np.linspace(350, 1500, 200), save_file=False)
+
+
+def test_sampling_both_wrong():
+    with pytest.raises(ValueError):
+        calibrate(mean_spectrum_avro_file, sampling=np.linspace(200, 2000, 100), save_file=False)

tests/test_calibrator/test_calibrator_archive.py

@@ -1,58 +1,59 @@
-import unittest
-
 import numpy.testing as npt
 import pandas.testing as pdt
+import pytest
 
 from gaiaxpy import calibrate
-from tests.test_calibrator.calibrator_solutions import solution_default_df, sol_with_missing_sampling_array, \
-    sol_default_sampling_array, with_missing_solution_df, missing_solution_df
+from tests.test_calibrator.calibrator_solutions import (solution_default_df, sol_with_missing_sampling_array,
+                                                        sol_default_sampling_array, with_missing_solution_df,
+                                                        missing_solution_df)
 from tests.utils.utils import missing_bp_source_id
 
 _rtol = 1e-10
 _atol = 1e-10
 
 
-class TestCalibratorSingleElement(unittest.TestCase):
-
-    def test_single_element_query(self):
-        query = "SELECT * FROM gaiadr3.gaia_source WHERE source_id='5853498713190525696'"
-        output_df, sampling = calibrate(query, save_file=False)
-        source_data_output = output_df[output_df['source_id'] == 5853498713190525696]
-        source_data_solution = solution_default_df[solution_default_df['source_id'] == 5853498713190525696]
-        pdt.assert_frame_equal(source_data_output, source_data_solution, atol=_atol, rtol=_rtol)
-        npt.assert_array_equal(sampling, sol_default_sampling_array)
-
-
-class TestCalibratorMissingBPQueryInput(unittest.TestCase):
-
-    def test_missing_bp_query(self):
-        query = f"SELECT * FROM gaiadr3.gaia_source WHERE source_id IN ('5853498713190525696', " \
-                f"{missing_bp_source_id}, '5762406957886626816')"
-        output_df, sampling = calibrate(query, save_file=False)
-        sorted_output_df = output_df.sort_values('source_id', ignore_index=True)
-        sorted_solution_df = with_missing_solution_df.sort_values('source_id', ignore_index=True)
-        pdt.assert_frame_equal(sorted_output_df, sorted_solution_df, atol=_atol, rtol=_rtol)
-        npt.assert_array_equal(sampling, sol_with_missing_sampling_array)
-
-    def test_missing_bp_query_isolated(self):
-        query = f"SELECT * FROM gaiadr3.gaia_source WHERE source_id IN ({missing_bp_source_id})"
-        output_df, sampling = calibrate(query, save_file=False)
-        pdt.assert_frame_equal(output_df, missing_solution_df, atol=_atol, rtol=_rtol)
-        npt.assert_array_equal(sampling, sol_with_missing_sampling_array)
-
-
-class TestCalibratorMissingBPListInput(unittest.TestCase):
-
-    def test_missing_bp_list(self):
-        src_list = ['5853498713190525696', str(missing_bp_source_id), '5762406957886626816']
-        output_df, sampling = calibrate(src_list, save_file=False)
-        sorted_output_df = output_df.sort_values('source_id', ignore_index=True)
-        sorted_solution_df = with_missing_solution_df.sort_values('source_id', ignore_index=True)
-        pdt.assert_frame_equal(sorted_output_df, sorted_solution_df, check_dtype=False, atol=_atol, rtol=_rtol)
-        npt.assert_array_equal(sampling, sol_with_missing_sampling_array)
-
-    def test_missing_bp_isolated(self):
-        src_list = [missing_bp_source_id]
-        output_df, sampling = calibrate(src_list, save_file=False)
-        pdt.assert_frame_equal(output_df, missing_solution_df, check_dtype=False, atol=_atol, rtol=_rtol)
-        npt.assert_array_equal(sampling, sol_with_missing_sampling_array)
+@pytest.mark.archive
+def test_single_element_query():
+    query = "SELECT * FROM gaiadr3.gaia_source WHERE source_id='5853498713190525696'"
+    output_df, sampling = calibrate(query, save_file=False)
+    source_data_output = output_df[output_df['source_id'] == 5853498713190525696]
+    source_data_solution = solution_default_df[solution_default_df['source_id'] == 5853498713190525696]
+    pdt.assert_frame_equal(source_data_output, source_data_solution, atol=_atol, rtol=_rtol)
+    npt.assert_array_equal(sampling, sol_default_sampling_array)
+
+
+@pytest.mark.archive
+def test_missing_bp_query():
+    query = f"SELECT * FROM gaiadr3.gaia_source WHERE source_id IN ('5853498713190525696', " \
+            f"{missing_bp_source_id}, '5762406957886626816')"
+    output_df, sampling = calibrate(query, save_file=False)
+    sorted_output_df = output_df.sort_values('source_id', ignore_index=True)
+    sorted_solution_df = with_missing_solution_df.sort_values('source_id', ignore_index=True)
+    pdt.assert_frame_equal(sorted_output_df, sorted_solution_df, atol=_atol, rtol=_rtol)
+    npt.assert_array_equal(sampling, sol_with_missing_sampling_array)
+
+
+@pytest.mark.archive
+def test_missing_bp_query_isolated():
+    query = f"SELECT * FROM gaiadr3.gaia_source WHERE source_id IN ({missing_bp_source_id})"
+    output_df, sampling = calibrate(query, save_file=False)
+    pdt.assert_frame_equal(output_df, missing_solution_df, atol=_atol, rtol=_rtol)
+    npt.assert_array_equal(sampling, sol_with_missing_sampling_array)
+
+
+@pytest.mark.archive
+def test_missing_bp_list():
+    src_list = ['5853498713190525696', str(missing_bp_source_id), '5762406957886626816']
+    output_df, sampling = calibrate(src_list, save_file=False)
+    sorted_output_df = output_df.sort_values('source_id', ignore_index=True)
+    sorted_solution_df = with_missing_solution_df.sort_values('source_id', ignore_index=True)
+    pdt.assert_frame_equal(sorted_output_df, sorted_solution_df, check_dtype=False, atol=_atol, rtol=_rtol)
+    npt.assert_array_equal(sampling, sol_with_missing_sampling_array)
+
+
+@pytest.mark.archive
+def test_missing_bp_isolated():
+    src_list = [missing_bp_source_id]
+    output_df, sampling = calibrate(src_list, save_file=False)
+    pdt.assert_frame_equal(output_df, missing_solution_df, check_dtype=False, atol=_atol, rtol=_rtol)
+    npt.assert_array_equal(sampling, sol_with_missing_sampling_array)