From 7832cdc3979de68dba531f4528cc9b94b8ba9926 Mon Sep 17 00:00:00 2001 From: Eric Novotny Date: Tue, 28 Apr 2026 07:58:26 -0700 Subject: [PATCH 1/2] update ini import for other tags in file --- cwmscli/usgs/__init__.py | 12 +- cwmscli/usgs/rating_ini_file_import.py | 70 ++- tests/usgs/__init__.py | 0 tests/usgs/test_rating_ini_file_import.py | 605 ++++++++++++++++++++++ 4 files changed, 669 insertions(+), 18 deletions(-) create mode 100644 tests/usgs/__init__.py create mode 100644 tests/usgs/test_rating_ini_file_import.py diff --git a/cwmscli/usgs/__init__.py b/cwmscli/usgs/__init__.py index 316be5f..96a1502 100644 --- a/cwmscli/usgs/__init__.py +++ b/cwmscli/usgs/__init__.py @@ -105,15 +105,23 @@ def getusgs_ratings(office, days_back, api_root, api_key, api_key_loc, rating_su type=str, help="filename of ratings ini file to be processed", ) +@click.option( + "--dry-run", + is_flag=True, + default=False, + help="Preview changes without updating the database", +) @api_root_option @api_key_option @api_key_loc_option @requires(reqs.cwms, reqs.requests) -def ratingsinifileimport(filename, api_root, api_key, api_key_loc): +def ratingsinifileimport(filename, dry_run, api_root, api_key, api_key_loc): from cwmscli.usgs.rating_ini_file_import import rating_ini_file_import api_key = get_api_key(api_key, api_key_loc) - rating_ini_file_import(api_root=api_root, api_key=api_key, ini_filename=filename) + rating_ini_file_import( + api_root=api_root, api_key=api_key, ini_filename=filename, dry_run=dry_run + ) @usgs_group.command("measurements", help="Store USGS measurements into CWMS database") diff --git a/cwmscli/usgs/rating_ini_file_import.py b/cwmscli/usgs/rating_ini_file_import.py index 970efc5..9a8012e 100644 --- a/cwmscli/usgs/rating_ini_file_import.py +++ b/cwmscli/usgs/rating_ini_file_import.py @@ -11,8 +11,12 @@ } -def rating_ini_file_import(api_root, api_key, ini_filename): - init_cwms_session(cwms, api_root=api_root, api_key="apikey " + api_key) +def rating_ini_file_import(api_root, api_key, ini_filename, dry_run=False): + if dry_run: + logging.info("DRY RUN MODE - no changes will be made") + init_cwms_session(cwms, api_root=api_root) + else: + init_cwms_session(cwms, api_root=api_root, api_key="apikey " + api_key) logging.info(f"CDA connection: {api_root}") logging.info(f"Opening ini file: {ini_filename}") @@ -21,7 +25,18 @@ def rating_ini_file_import(api_root, api_key, ini_filename): ini_file.close() params = {} - keywords = ["cwms_office", "db_base", "db_exsa", "db_corr", "localid"] + keywords = [ + "cwms_office", + "cwms_database", + "db_base", + "db_exsa", + "db_corr", + "db_tail", + "db_river", + "localid", + "cwmsid", + "textfile", + ] rating_errors = [] for i in range(len(lines)): line = lines[i][:-1].strip() @@ -33,26 +48,45 @@ def rating_ini_file_import(api_root, api_key, ini_filename): continue if "=" in line: fields = line.split("=") - if fields[0] in keywords: - if fields[0] == "cwms_office": - fields[1] = fields[1].upper() - params[fields[0]] = fields[1] + key = fields[0].strip().lower() + if key in keywords: + if key == "cwms_office": + fields[1] = fields[1].strip().upper() + else: + fields[1] = fields[1].strip() + params[key] = fields[1] else: fields = parse_ini_line(line) if fields[0] in rating_types.keys(): - rating_db_type = rating_types[fields[0]]["db_type"] - if f"$(${rating_db_type})" in fields: - rating_spec = params[rating_db_type].replace( - "\$localid", params["localid"] - ) + # Find the database reference in the fields (e.g., $($db_tail), $($db_exsa), etc.) + db_key = None + for field in fields[1:]: + if field.startswith("$(") and field.endswith(")"): + # Extract the key name from $(...), e.g., "db_exsa" from "$($db_exsa)" + potential_key = field[2:-1].lstrip("$") + if potential_key in params: + db_key = potential_key + break + + if db_key: + rating_spec = params[db_key] + # Handle both localid and cwmsid substitution + if "localid" in params: + rating_spec = rating_spec.replace( + "\$localid", params["localid"] + ) + if "cwmsid" in params: + rating_spec = rating_spec.replace("\$cwmsid", params["cwmsid"]) logging.info(f"Updating rating specification: {rating_spec}") try: update_rating_spec( rating_spec, - params["cwms_office"], + params.get("cwms_office"), rating_types[fields[0]]["db_disc"], + dry_run=dry_run, ) - logging.info("SUCCESS: rating specification changes stored") + if not dry_run: + logging.info("SUCCESS: rating specification changes stored") except: logging.error( "ERROR: rating specificataion could not be update" @@ -109,7 +143,7 @@ def parse_ini_line(line): return fields -def update_rating_spec(rating_id, office_id, db_disc): +def update_rating_spec(rating_id, office_id, db_disc, dry_run=False): rating_spec = cwms.get_rating_spec(rating_id=rating_id, office_id=office_id) data = rating_spec.df data = data.drop("effective-dates", axis=1) @@ -127,4 +161,8 @@ def update_rating_spec(rating_id, office_id, db_disc): disc = data.loc[0, "description"] logging.info(f"Saving specification discription as: {disc}") data_xml = cwms.rating_spec_df_to_xml(data) - cwms.store_rating_spec(data=data_xml, fail_if_exists=False) + if dry_run: + logging.info("DRY RUN: Would store rating specification with XML:") + logging.info(data_xml) + else: + cwms.store_rating_spec(data=data_xml, fail_if_exists=False) diff --git a/tests/usgs/__init__.py b/tests/usgs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/usgs/test_rating_ini_file_import.py b/tests/usgs/test_rating_ini_file_import.py new file mode 100644 index 0000000..52800f8 --- /dev/null +++ b/tests/usgs/test_rating_ini_file_import.py @@ -0,0 +1,605 @@ +import os +import tempfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pandas as pd + +from cwmscli.usgs.rating_ini_file_import import ( + parse_ini_line, + rating_ini_file_import, + rating_types, + update_rating_spec, +) + + +class TestParseIniLine: + """Test the parse_ini_line function with various input formats.""" + + def test_simple_space_separated_fields(self): + """Test parsing simple space-separated fields.""" + line = "store_corr $($db_corr)" + fields = parse_ini_line(line) + assert fields == ["store_corr", "$($db_corr)"] + + def test_quoted_fields_with_spaces(self): + """Test parsing fields quoted with double quotes.""" + line = 'store_corr "CEDI4.Stage;Flow.USGS-BASE.USGS-NWIS"' + fields = parse_ini_line(line) + assert len(fields) == 2 + assert fields[0] == "store_corr" + assert "CEDI4.Stage;Flow.USGS-BASE.USGS-NWIS" in fields[1] + + def test_single_quoted_fields(self): + """Test parsing fields quoted with single quotes.""" + line = "store_corr 'CEDI4.Stage;Flow.USGS-BASE.USGS-NWIS'" + fields = parse_ini_line(line) + assert len(fields) == 2 + assert fields[0] == "store_corr" + + def test_tab_separated_fields(self): + """Test parsing tab-separated fields.""" + line = "store_corr\t$($db_corr)" + fields = parse_ini_line(line) + assert fields == ["store_corr", "$($db_corr)"] + + def test_escaped_quotes(self): + """Test parsing with escaped quotes.""" + line = r'field1 "quoted \"value\" here"' + fields = parse_ini_line(line) + assert len(fields) >= 1 + assert fields[0] == "field1" + + def test_empty_line(self): + """Test parsing empty line.""" + line = "" + fields = parse_ini_line(line) + assert fields == [] + + def test_single_field(self): + """Test parsing single field.""" + line = "field1" + fields = parse_ini_line(line) + assert fields == ["field1"] + + def test_multiple_spaces_between_fields(self): + """Test parsing with multiple spaces between fields.""" + line = "field1 field2 field3" + fields = parse_ini_line(line) + assert fields == ["field1", "field2", "field3"] + + def test_quoted_field_preserves_internal_spaces(self): + """Test that spaces inside quotes are preserved.""" + line = 'field1 "field with spaces" field3' + fields = parse_ini_line(line) + assert len(fields) == 3 + assert "field with spaces" in fields[1] + + +class TestRatingTypesConfig: + """Test the rating_types configuration.""" + + def test_rating_types_structure(self): + """Verify rating_types has expected structure.""" + assert "store_corr" in rating_types + assert "store_base" in rating_types + assert "store_exsa" in rating_types + + def test_store_corr_config(self): + """Test store_corr configuration.""" + config = rating_types["store_corr"] + assert config["db_type"] == "db_corr" + assert config["db_disc"] == "USGS-CORR" + + def test_store_base_config(self): + """Test store_base configuration.""" + config = rating_types["store_base"] + assert config["db_type"] == "db_base" + assert config["db_disc"] == "USGS-BASE" + + def test_store_exsa_config(self): + """Test store_exsa configuration.""" + config = rating_types["store_exsa"] + assert config["db_type"] == "db_exsa" + assert config["db_disc"] == "USGS-EXSA" + + +class TestUpdateRatingSpec: + """Test the update_rating_spec function.""" + + @patch("cwmscli.usgs.rating_ini_file_import.cwms") + def test_update_rating_spec_basic(self, mock_cwms): + """Test basic rating spec update.""" + mock_df = pd.DataFrame( + { + "active": [False], + "auto-update": [False], + "auto-activate": [False], + "source-agency": ["OTHER"], + "description": ["Old description"], + "effective-dates": ["2020-01-01"], + } + ) + mock_rating_spec = MagicMock() + mock_rating_spec.df = mock_df.copy() + + mock_cwms.get_rating_spec.return_value = mock_rating_spec + mock_cwms.rating_spec_df_to_xml.return_value = "" + + update_rating_spec("CEDI4.Stage;Flow", "MVP", "USGS-CORR") + + mock_cwms.get_rating_spec.assert_called_once_with( + rating_id="CEDI4.Stage;Flow", office_id="MVP" + ) + mock_cwms.store_rating_spec.assert_called_once() + + @patch("cwmscli.usgs.rating_ini_file_import.cwms") + def test_update_rating_spec_sets_flags(self, mock_cwms): + """Test that update_rating_spec sets all required flags.""" + mock_df = pd.DataFrame( + { + "active": [False], + "auto-update": [False], + "auto-activate": [False], + "source-agency": ["OTHER"], + "description": ["Old"], + "effective-dates": ["2020-01-01"], + } + ) + mock_rating_spec = MagicMock() + mock_rating_spec.df = mock_df.copy() + + mock_cwms.get_rating_spec.return_value = mock_rating_spec + mock_cwms.rating_spec_df_to_xml.return_value = "" + + update_rating_spec("TEST_ID", "MVP", "USGS-CORR") + + # Get the dataframe that was modified (called with positional arg) + modified_df = mock_cwms.rating_spec_df_to_xml.call_args[0][0] + assert modified_df["active"].iloc[0] + assert modified_df["auto-update"].iloc[0] + assert modified_df["auto-activate"].iloc[0] + assert modified_df["source-agency"].iloc[0] == "USGS" + + @patch("cwmscli.usgs.rating_ini_file_import.cwms") + def test_update_rating_spec_adds_description(self, mock_cwms): + """Test that update_rating_spec adds discriminator to description.""" + mock_df = pd.DataFrame( + { + "active": [False], + "auto-update": [False], + "auto-activate": [False], + "source-agency": ["OTHER"], + "description": ["Existing"], + "effective-dates": ["2020-01-01"], + } + ) + mock_rating_spec = MagicMock() + mock_rating_spec.df = mock_df.copy() + + mock_cwms.get_rating_spec.return_value = mock_rating_spec + mock_cwms.rating_spec_df_to_xml.return_value = "" + + update_rating_spec("TEST_ID", "MVP", "USGS-CORR") + + modified_df = mock_cwms.rating_spec_df_to_xml.call_args[0][0] + assert "USGS-CORR" in modified_df["description"].iloc[0] + + @patch("cwmscli.usgs.rating_ini_file_import.cwms") + def test_update_rating_spec_with_missing_description(self, mock_cwms): + """Test update_rating_spec when description column doesn't exist.""" + mock_df = pd.DataFrame( + { + "active": [False], + "auto-update": [False], + "auto-activate": [False], + "source-agency": ["OTHER"], + "effective-dates": ["2020-01-01"], + } + ) + mock_rating_spec = MagicMock() + mock_rating_spec.df = mock_df.copy() + + mock_cwms.get_rating_spec.return_value = mock_rating_spec + mock_cwms.rating_spec_df_to_xml.return_value = "" + + update_rating_spec("TEST_ID", "MVP", "USGS-CORR") + + modified_df = mock_cwms.rating_spec_df_to_xml.call_args[0][0] + assert "description" in modified_df.columns + assert modified_df["description"].iloc[0] == "USGS-CORR" + + @patch("cwmscli.usgs.rating_ini_file_import.cwms") + def test_update_rating_spec_dry_run_doesnt_store(self, mock_cwms): + """Test that update_rating_spec with dry_run=True doesn't call store_rating_spec.""" + mock_df = pd.DataFrame( + { + "active": [False], + "auto-update": [False], + "auto-activate": [False], + "source-agency": ["OTHER"], + "description": ["Old"], + "effective-dates": ["2020-01-01"], + } + ) + mock_rating_spec = MagicMock() + mock_rating_spec.df = mock_df.copy() + + mock_cwms.get_rating_spec.return_value = mock_rating_spec + mock_cwms.rating_spec_df_to_xml.return_value = "" + + update_rating_spec("TEST_ID", "MVP", "USGS-CORR", dry_run=True) + + # rating_spec_df_to_xml should still be called + mock_cwms.rating_spec_df_to_xml.assert_called_once() + # But store_rating_spec should NOT be called + mock_cwms.store_rating_spec.assert_not_called() + + +class TestRatingIniFileImport: + """Test the main rating_ini_file_import function.""" + + @patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session") + @patch("cwmscli.usgs.rating_ini_file_import.update_rating_spec") + def test_import_with_real_file(self, mock_update, mock_init_cwms): + """Test import with the actual mvp_ratings_ini.ini file.""" + desktop_file = Path.home() / "Desktop" / "mvp_ratings_ini.ini" + + if desktop_file.exists(): + # Run the import + rating_ini_file_import( + "http://localhost:8080", "test_key", str(desktop_file) + ) + + # Verify init_cwms_session was called + mock_init_cwms.assert_called_once() + # Verify update_rating_spec was called for each non-commented store_* line + assert mock_update.call_count > 0 + + def test_import_with_simple_ini_file(self): + """Test import with a simple temporary INI file.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + """ +# Test config +cwms_office=MVP +db_corr=$localid.Stage;Flow.USGS-CORR.USGS-NWIS +localid=TESTLOC +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + + mock_update.assert_called_once() + finally: + os.unlink(temp_file) + + def test_import_parameter_parsing(self): + """Test that import correctly parses configuration parameters.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + r""" +cwms_office=MVP +db_base=BASE_\$localid.SPEC +db_exsa=EXSA_\$localid.SPEC +db_corr=CORR_\$localid.SPEC +localid=TESTLOC +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + + # Verify that update_rating_spec was called with correct parameters + mock_update.assert_called_once() + args = mock_update.call_args[0] + # The rating_spec should be the db_corr value with localid substituted + assert "TESTLOC" in args[0] + assert args[1] == "MVP" # office_id + finally: + os.unlink(temp_file) + + def test_import_skips_comments(self): + """Test that import correctly skips commented lines.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + """ +cwms_office=MVP +db_corr=$localid.Stage;Flow.USGS-CORR.USGS-NWIS +db_exsa=$localid.Stage;Flow.USGS-EXSA.USGS-NWIS +localid=LOC1 +#store_corr $($db_corr) +store_exsa $($db_exsa) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + # Should only call for store_exsa, not commented store_corr + mock_update.assert_called_once() + finally: + os.unlink(temp_file) + + def test_import_handles_inline_comments(self): + """Test that import correctly handles inline comments.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + """ +cwms_office=MVP # This is the office +db_corr=$localid.Stage;Flow.USGS-CORR.USGS-NWIS +localid=TESTLOC # Location identifier +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + # Verify office_id is correctly set (without comment) + mock_update.assert_called_once() + assert mock_update.call_args[0][1] == "MVP" + finally: + os.unlink(temp_file) + + def test_import_office_id_uppercase(self): + """Test that office_id is converted to uppercase.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + """ +cwms_office=mvp +db_corr=$localid.Stage;Flow.USGS-CORR.USGS-NWIS +localid=TESTLOC +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + # Verify office_id is uppercase + mock_update.assert_called_once() + assert mock_update.call_args[0][1] == "MVP" + finally: + os.unlink(temp_file) + + def test_import_handles_multiple_locations(self): + """Test that import correctly processes multiple location blocks.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + """ +cwms_office=MVP +db_corr=$localid.Stage;Flow.USGS-CORR.USGS-NWIS +localid=LOC1 +store_corr $($db_corr) +localid=LOC2 +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + # Should be called twice, once for each location + assert mock_update.call_count == 2 + finally: + os.unlink(temp_file) + + def test_import_localid_substitution(self): + """Test that $localid is correctly substituted in specifications.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + r""" +cwms_office=MVP +db_corr=\$localid.Stage;Flow.USGS-CORR.USGS-NWIS +localid=MYLOC +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + + mock_update.assert_called_once() + # First argument should have MYLOC substituted for \$localid + rating_spec_arg = mock_update.call_args[0][0] + assert "MYLOC" in rating_spec_arg + finally: + os.unlink(temp_file) + + def test_import_cwmsid_substitution_with_nae_format(self): + """Test NAE format with cwmsid and flexible db references (db_tail, db_river).""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + r""" +CWMS_OFFICE=NAE +CWMS_DATABASE=local +db_tail=\$cwmsid.Stage-TAILWATER;Flow.USGS-EXSA.USGS-NWIS +db_river=\$cwmsid.Stage;Flow.USGS-EXSA.USGS-NWIS + +cwmsid=BMD +usgsid=01155500 +replace_exsa $(textfile) +store_exsa $($db_tail) + +cwmsid=NHD +usgsid=01151500 +replace_exsa $(textfile) +store_exsa $($db_river) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + + # Should be called twice (once for each cwmsid, replace_exsa is skipped) + assert mock_update.call_count == 2 + + # First call should use db_tail with BMD substituted + call1_args = mock_update.call_args_list[0][0] + assert "BMD" in call1_args[0] + assert "TAILWATER" in call1_args[0] + assert call1_args[1] == "NAE" + + # Second call should use db_river with NHD substituted + call2_args = mock_update.call_args_list[1][0] + assert "NHD" in call2_args[0] + assert "TAILWATER" not in call2_args[0] + assert call2_args[1] == "NAE" + finally: + os.unlink(temp_file) + + def test_dry_run_mode_calls_update_with_flag(self): + """Test that dry_run mode calls update_rating_spec with dry_run=True.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + r""" +cwms_office=MVP +db_corr=\$localid.Stage;Flow.USGS-CORR.USGS-NWIS +localid=TESTLOC +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + # Run with dry_run=True + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file, dry_run=True + ) + + # update_rating_spec should be called WITH dry_run=True + mock_update.assert_called_once() + assert mock_update.call_args[1]["dry_run"] is True + finally: + os.unlink(temp_file) + + def test_dry_run_mode_with_multiple_entries(self): + """Test that dry_run processes all entries with dry_run flag.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + r""" +CWMS_OFFICE=NAE +CWMS_DATABASE=local +db_tail=\$cwmsid.Stage-TAILWATER;Flow.USGS-EXSA.USGS-NWIS +db_river=\$cwmsid.Stage;Flow.USGS-EXSA.USGS-NWIS + +cwmsid=BMD +usgsid=01155500 +store_exsa $($db_tail) + +cwmsid=NHD +usgsid=01151500 +store_exsa $($db_river) + +cwmsid=NSD +usgsid=01153000 +store_exsa $($db_tail) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + # Run with dry_run=True + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file, dry_run=True + ) + + # update_rating_spec should be called 3 times with dry_run=True + assert mock_update.call_count == 3 + for call in mock_update.call_args_list: + assert call[1]["dry_run"] is True + finally: + os.unlink(temp_file) + + def test_normal_mode_calls_updates(self): + """Test that normal mode (not dry_run) calls update_rating_spec with dry_run=False.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + r""" +cwms_office=MVP +db_corr=\$localid.Stage;Flow.USGS-CORR.USGS-NWIS +localid=TESTLOC +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + # Run with dry_run=False (default) + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file, dry_run=False + ) + + # update_rating_spec SHOULD be called with dry_run=False + mock_update.assert_called_once() + assert mock_update.call_args[1]["dry_run"] is False + finally: + os.unlink(temp_file) From 2afe42141a06211ed3a14eb1e89a0554df2ee987 Mon Sep 17 00:00:00 2001 From: Eric Novotny Date: Tue, 28 Apr 2026 08:07:09 -0700 Subject: [PATCH 2/2] update to work with any custum tags --- cwmscli/usgs/rating_ini_file_import.py | 34 ++++++----------------- tests/usgs/test_rating_ini_file_import.py | 33 ++++++++++++++++++++++ 2 files changed, 42 insertions(+), 25 deletions(-) diff --git a/cwmscli/usgs/rating_ini_file_import.py b/cwmscli/usgs/rating_ini_file_import.py index 9a8012e..a194a08 100644 --- a/cwmscli/usgs/rating_ini_file_import.py +++ b/cwmscli/usgs/rating_ini_file_import.py @@ -25,18 +25,6 @@ def rating_ini_file_import(api_root, api_key, ini_filename, dry_run=False): ini_file.close() params = {} - keywords = [ - "cwms_office", - "cwms_database", - "db_base", - "db_exsa", - "db_corr", - "db_tail", - "db_river", - "localid", - "cwmsid", - "textfile", - ] rating_errors = [] for i in range(len(lines)): line = lines[i][:-1].strip() @@ -49,12 +37,10 @@ def rating_ini_file_import(api_root, api_key, ini_filename, dry_run=False): if "=" in line: fields = line.split("=") key = fields[0].strip().lower() - if key in keywords: - if key == "cwms_office": - fields[1] = fields[1].strip().upper() - else: - fields[1] = fields[1].strip() - params[key] = fields[1] + value = fields[1].strip() + if key == "cwms_office": + value = value.upper() + params[key] = value else: fields = parse_ini_line(line) if fields[0] in rating_types.keys(): @@ -70,13 +56,11 @@ def rating_ini_file_import(api_root, api_key, ini_filename, dry_run=False): if db_key: rating_spec = params[db_key] - # Handle both localid and cwmsid substitution - if "localid" in params: - rating_spec = rating_spec.replace( - "\$localid", params["localid"] - ) - if "cwmsid" in params: - rating_spec = rating_spec.replace("\$cwmsid", params["cwmsid"]) + # Substitute any custom parameters found in rating_spec + for param_key, param_value in params.items(): + placeholder = f"\\${param_key}" + if placeholder in rating_spec: + rating_spec = rating_spec.replace(placeholder, param_value) logging.info(f"Updating rating specification: {rating_spec}") try: update_rating_spec( diff --git a/tests/usgs/test_rating_ini_file_import.py b/tests/usgs/test_rating_ini_file_import.py index 52800f8..15bd7c0 100644 --- a/tests/usgs/test_rating_ini_file_import.py +++ b/tests/usgs/test_rating_ini_file_import.py @@ -575,6 +575,39 @@ def test_dry_run_mode_with_multiple_entries(self): finally: os.unlink(temp_file) + def test_import_custom_tag_substitution(self): + """Test that custom tags (not just localid/cwmsid) are substituted.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: + f.write( + r""" +cwms_office=MVP +db_corr=\$location.\$parameter.USGS-CORR.USGS-NWIS +location=TESTLOC +parameter=Stage;Flow +store_corr $($db_corr) +""" + ) + temp_file = f.name + + try: + with patch("cwmscli.usgs.rating_ini_file_import.init_cwms_session"): + with patch( + "cwmscli.usgs.rating_ini_file_import.update_rating_spec" + ) as mock_update: + rating_ini_file_import( + "http://localhost:8080", "test_key", temp_file + ) + + mock_update.assert_called_once() + rating_spec_arg = mock_update.call_args[0][0] + # Both custom tags should be substituted + assert "TESTLOC" in rating_spec_arg + assert "Stage;Flow" in rating_spec_arg + assert "\$location" not in rating_spec_arg + assert "\$parameter" not in rating_spec_arg + finally: + os.unlink(temp_file) + def test_normal_mode_calls_updates(self): """Test that normal mode (not dry_run) calls update_rating_spec with dry_run=False.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".ini", delete=False) as f: