Added tested function 'hres_ic.set_replace_function()' to define which replacement function to use, based on the value of the command line argument '--type'.

Author: atteggiani

src/hres_ic.py

@@ -78,8 +78,30 @@ def parse_arguments():
     parser.add_argument('--type', default="era5land")
     return parser.parse_args()
 
-def main():
+def set_replace_function(data_type):
+    """
+    Set what replace function to use, based on the command line argument '--type'.
 
+    Parameters
+    ----------
+    data_type: str
+        Type of data source for the replacement. The value of the command line argument '--type'.
+
+    Returns
+    -------
+    callable or None
+        The function to be called for the replacement, or None if no replacement needs to take place.
+    """ 
+    if data_type == "era5land":
+        replace_function = replace_landsurface_with_ERA5land_IC.swap_land_era5land
+    elif data_type == "barra":
+        replace_function = replace_landsurface_with_BARRA2R_IC.swap_land_barra
+    else:
+        replace_function = None
+    return replace_function
+
+    
+def main():
     """
     The main function that creates a worker pool and generates single GRIB files 
     for requested date/times in parallel.
@@ -102,12 +124,10 @@ def main():
     print(f"replacement_file = {args.file}")
     print(f"start_time = {t}")
 
-    # If necessary replace ERA5 land/surface fields with higher-resolution options
-    if "era5land" in args.type:
-        replace_landsurface_with_ERA5land_IC.swap_land_era5land(args.mask, args.file, t)
-        replace_input_file_with_tmp_input_file(args.file)
-    elif "barra" in args.type:
-        replace_landsurface_with_BARRA2R_IC.swap_land_barra(args.mask, args.file, t)
+    # If necessary replace land/surface fields with higher-resolution options
+    replace_function = set_replace_function(args.type)
+    if replace_function is not None:
+        replace_function(args.mask, args.file, t)
         replace_input_file_with_tmp_input_file(args.file)
     else:
         print("No need to swap out IC")

tests/test_hres.py

@@ -12,7 +12,7 @@
 srcpath = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),'src') #To delete when src is a package
 sys.path.insert(0,srcpath) #To delete when src is a package
 
-from hres_ic import get_start_time, replace_input_file_with_tmp_input_file, parse_arguments
+from hres_ic import get_start_time, replace_input_file_with_tmp_input_file, parse_arguments, set_replace_function
 
 del sys.path[0] #To delete when src is a package
 
@@ -48,7 +48,7 @@ def test_parse_arguments_with_type():
     args = parse_arguments()
     assert args.type == 'newtype'
 
-@patch('sys.argv', ['program_name', '--file', 'file_path', '--start', '2024-08-12'])
+@patch('sys.argv', ['program_name', '--file', 'file_path', '--start', '202408121230'])
 def test_parse_arguments_missing_mask():
     with pytest.raises(SystemExit):
         parse_arguments()
@@ -61,4 +61,18 @@ def test_parse_arguments_missing_file():
 @patch('sys.argv', ['program_name', '--mask', 'mask_path', '--file', 'file_path'])
 def test_parse_arguments_missing_start():
     with pytest.raises(SystemExit):
-        parse_arguments()
+        parse_arguments()
+
+@patch('replace_landsurface_with_ERA5land_IC.swap_land_era5land')
+def test_set_replace_function_era5land(mock_era5land):
+    result = set_replace_function("era5land")
+    assert result == mock_era5land
+
+@patch('replace_landsurface_with_BARRA2R_IC.swap_land_barra')
+def test_set_replace_function_barra(mock_barra):
+    result = set_replace_function("barra")
+    assert result == mock_barra
+
+def test_set_replace_function_unknown():
+    result = set_replace_function("unknown")
+    assert result is None