From 9198f028f9d9cf84edc37f09ccb971fa8d7817b6 Mon Sep 17 00:00:00 2001 From: Paulina Grochal Date: Wed, 11 Mar 2026 18:52:00 +0100 Subject: [PATCH 1/3] DT-294 generate plot_info file for fargate run --- poetry.lock | 88 +++++++++++- pyproject.toml | 1 + pytrajplot/main.py | 106 +++++++++++++++ test/integration/test_main.py | 2 +- test/unit/test_plot_info.py | 245 ++++++++++++++++++++++++++++++++++ 5 files changed, 440 insertions(+), 2 deletions(-) create mode 100644 test/unit/test_plot_info.py diff --git a/poetry.lock b/poetry.lock index 4d821b1..174f7f2 100644 --- a/poetry.lock +++ b/poetry.lock @@ -160,6 +160,54 @@ type = "legacy" url = "https://service.meteoswiss.ch/nexus/repository/python-all/simple" reference = "meteoswiss" +[[package]] +name = "boto3" +version = "1.42.65" +description = "The AWS SDK for Python" +optional = false +python-versions = ">= 3.9" +files = [ + {file = "boto3-1.42.65-py3-none-any.whl", hash = "sha256:cc7f2e0aec6c68ee5b10232cf3e01326acf6100bc785a770385b61a0474b31f4"}, + {file = "boto3-1.42.65.tar.gz", hash = "sha256:c740af6bdaebcc1a00f3827a5729050bf6fc820ee148bf7d06f28db11c80e2a1"}, +] + +[package.dependencies] +botocore = ">=1.42.65,<1.43.0" +jmespath = ">=0.7.1,<2.0.0" +s3transfer = ">=0.16.0,<0.17.0" + +[package.extras] +crt = ["botocore[crt] (>=1.21.0,<2.0a0)"] + +[package.source] +type = "legacy" +url = "https://service.meteoswiss.ch/nexus/repository/python-all/simple" +reference = "meteoswiss" + +[[package]] +name = "botocore" +version = "1.42.65" +description = "Low-level, data-driven core of boto 3." +optional = false +python-versions = ">= 3.9" +files = [ + {file = "botocore-1.42.65-py3-none-any.whl", hash = "sha256:0283c332ce00cbd1b894e86b7bed89dd624a5ca3a4ee62ec4db3898d16652e98"}, + {file = "botocore-1.42.65.tar.gz", hash = "sha256:7d52c148df07f70c375eeda58f99b439c7c7836c25df74cccfba3bb6e12444d2"}, +] + +[package.dependencies] +jmespath = ">=0.7.1,<2.0.0" +python-dateutil = ">=2.1,<3.0.0" +urllib3 = {version = ">=1.25.4,<2.2.0 || >2.2.0,<3", markers = "python_version >= \"3.10\""} + +[package.extras] +crt = ["awscrt (==0.31.2)"] + +[package.source] +type = "legacy" +url = "https://service.meteoswiss.ch/nexus/repository/python-all/simple" +reference = "meteoswiss" + [[package]] name = "cartopy" version = "0.25.0" @@ -906,6 +954,22 @@ type = "legacy" url = "https://service.meteoswiss.ch/nexus/repository/python-all/simple" reference = "meteoswiss" +[[package]] +name = "jmespath" +version = "1.1.0" +description = "JSON Matching Expressions" +optional = false +python-versions = ">=3.9" +files = [ + {file = "jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64"}, + {file = "jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d"}, +] + +[package.source] +type = "legacy" +url = "https://service.meteoswiss.ch/nexus/repository/python-all/simple" +reference = "meteoswiss" + [[package]] name = "kiwisolver" version = "1.4.9" @@ -2480,6 +2544,28 @@ type = "legacy" url = "https://service.meteoswiss.ch/nexus/repository/python-all/simple" reference = "meteoswiss" +[[package]] +name = "s3transfer" +version = "0.16.0" +description = "An Amazon S3 Transfer Manager" +optional = false +python-versions = ">= 3.9" +files = [ + {file = "s3transfer-0.16.0-py3-none-any.whl", hash = "sha256:18e25d66fed509e3868dc1572b3f427ff947dd2c56f844a5bf09481ad3f3b2fe"}, + {file = "s3transfer-0.16.0.tar.gz", hash = "sha256:8e990f13268025792229cd52fa10cb7163744bf56e719e0b9cb925ab79abf920"}, +] + +[package.dependencies] +botocore = ">=1.37.4,<2.0a.0" + +[package.extras] +crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"] + +[package.source] +type = "legacy" +url = "https://service.meteoswiss.ch/nexus/repository/python-all/simple" +reference = "meteoswiss" + [[package]] name = "shapely" version = "2.1.2" @@ -2922,4 +3008,4 @@ reference = "meteoswiss" [metadata] lock-version = "2.0" python-versions = "~3.13" -content-hash = "0fb46ea8bd8a7d18fd6679ed8b364b11e8a31ea3a0dc052523ce8627040789bb" +content-hash = "7f5048577f4991dc4b0ba10676e4ddc482c88da76501ca85f827a1ae2e8a254f" diff --git a/pyproject.toml b/pyproject.toml index 6f4d750..9667022 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ cartopy = "^0.25.0" matplotlib = "^3.10.7" numpy = "^2.3.5" pandas = "^2.3.3" +boto3 = "^1.42.65" [tool.poetry.group.dev.dependencies] mypy = "^1.18.1" diff --git a/pytrajplot/main.py b/pytrajplot/main.py index 69a05b8..36f7104 100644 --- a/pytrajplot/main.py +++ b/pytrajplot/main.py @@ -1,8 +1,12 @@ """Command line interface of pytrajplot.""" from typing import Tuple, Dict +import logging +import os +from pathlib import Path # Third-party import click +import boto3 # First-party from pytrajplot import __version__ @@ -10,6 +14,10 @@ from pytrajplot.parse_data import check_input_dir from pytrajplot.utils import count_to_log_level +# Setup logging +log_level = os.getenv("LOG_LEVEL", "INFO").upper() +logging.basicConfig(level=log_level) +logger = logging.getLogger(__name__) def print_version(ctx: click.Context, _param: click.Parameter, value: bool) -> None: """Print the version number and exit.""" @@ -17,6 +25,79 @@ def print_version(ctx: click.Context, _param: click.Parameter, value: bool) -> N click.echo(__version__) ctx.exit(0) +def replace_variables(template_content: str) -> str: + """ + Replace $VAR with actual environment variable values. + Args: + template_content: Template string with $VARIABLE placeholders + Returns: + String with variables replaced by environment values + """ + result = template_content + # Get all environment variables as dict + env_vars = dict(os.environ) + + # Replace variables found in the template + for env_key, env_value in env_vars.items(): + placeholder = f'${env_key}' + if placeholder in result: + result = result.replace(placeholder, env_value) + logger.info(f"Replaced {placeholder} with {env_value}") + return result + + +def check_plot_info_file(input_dir: str, info_name: str, ssm_parameter_path: str = None) -> bool: + """ + Check if plot_info file exists in input directory. + If not found, fetch from SSM parameter and create it replacing variables. + Args: + input_dir: Input directory path + info_name: Name of the plot info file + ssm_parameter_path: SSM parameter path (optional, uses env var if not provided) + Returns: + bool: True if file exists or was created successfully, False otherwise + """ + input_path = Path(input_dir) + plot_info_file = input_path / info_name + + # Check if plot_info file already exists + if plot_info_file.exists(): + logger.info(f"Plot info file already exists: {plot_info_file}") + return True + + # File doesn't exist, try to create it from SSM parameter + logger.info(f"Plot info file not found: {plot_info_file}") + + try: + # Get SSM parameter path from argument or environment + ssm_param_path = ssm_parameter_path or os.environ.get('SSM_PARAMETER_PATH', '/pytrajplot/icon/plot_info') + logger.info(f"Fetching SSM parameter: {ssm_param_path}") + + # Fetch template from SSM Parameter + ssm_client = boto3.client('ssm') + response = ssm_client.get_parameter( + Name=ssm_param_path, + WithDecryption=True + ) + + # Get the template content + template_content = response['Parameter']['Value'] + logger.info(f"Template content length: {len(template_content)} chars") + + # Replace variables with environment variable values + substituted_content = replace_variables(template_content) + + # Create the plot_info file + with open(plot_info_file, 'w') as f: + f.write(substituted_content) + + logger.info(f"Successfully created plot info file: {plot_info_file}") + return True + + except Exception as e: + logger.error(f"Failed to create plot info file from SSM parameter: {str(e)}") + logger.error(f"SSM parameter path: {ssm_parameter_path or os.environ.get('SSM_PARAMETER_PATH', 'not_set')}") + return False def interpret_options(start_prefix: str, traj_prefix: str, info_name: str, language: str) -> Tuple[Dict[str, str], str]: """Reformat command line inputs. @@ -124,6 +205,17 @@ def interpret_options(start_prefix: str, traj_prefix: str, info_name: str, langu default=["pdf"], help="Choose data type(s) of final result. Default: pdf", ) +@click.option( + "--ssm-parameter-path", + type=str, + help="SSM parameter path for plot_info template. Uses SSM_PARAMETER_PATH env var if not specified.", +) +@click.option( + "--skip-ssm-fallback", + is_flag=True, + default=False, + help="Skip SSM parameter fallback if plot_info file is missing.", +) @click.option( "--version", "-V", @@ -143,7 +235,21 @@ def cli( language: str, domain: str, datatype: str, + ssm_parameter_path: str = None, + skip_ssm_fallback: bool = False, ) -> None: + # Check if plot_info file exists (create from SSM if needed) + if not skip_ssm_fallback: + plot_info_created = check_plot_info_file( + input_dir=input_dir, + info_name=info_name, + ssm_parameter_path=ssm_parameter_path + ) + + if not plot_info_created: + logger.error("Failed to check if plot_info file exists. Use --skip-ssm-fallback to continue anyway.") + raise click.ClickException("Missing plot_info file and failed to create from SSM parameter.") + prefix_dict, language = interpret_options( start_prefix=start_prefix, traj_prefix=traj_prefix, diff --git a/test/integration/test_main.py b/test/integration/test_main.py index 8649586..2eaa276 100644 --- a/test/integration/test_main.py +++ b/test/integration/test_main.py @@ -171,6 +171,7 @@ def create_args(input_dir: str, output_dir: str, opts: dict) -> list: # Positional arguments args.append(input_dir) args.append(output_dir) + args.append("--skip-ssm-fallback") # Keyword arguments for key, value in opts.items(): @@ -208,4 +209,3 @@ def test_pytrajplot(input_args, input_dir, output_dir): for rel in expected: expected_file = Path(output_path) / Path(rel).name assert expected_file.exists(), f"Expected output not found: {expected_file}" - diff --git a/test/unit/test_plot_info.py b/test/unit/test_plot_info.py new file mode 100644 index 0000000..59dacae --- /dev/null +++ b/test/unit/test_plot_info.py @@ -0,0 +1,245 @@ +"""Unit tests for pytrajplot main module.""" +import json +import os +import pytest +from pathlib import Path +from unittest.mock import Mock, patch, mock_open, MagicMock +from click.testing import CliRunner +from botocore.exceptions import ClientError + +from pytrajplot.main import replace_variables, check_plot_info_file, cli + + +class TestReplaceVariables: + """Test the replace_variables function.""" + + def test_replace_single_variable(self, monkeypatch): + monkeypatch.setenv("TEST_VAR", "test_value") + + template = "This is a $TEST_VAR template" + result = replace_variables(template) + + assert result == "This is a test_value template" + + def test_replace_multiple_variables(self, monkeypatch): + """Test replacing multiple environment variables.""" + monkeypatch.setenv("VAR1", "value1") + monkeypatch.setenv("VAR2", "value2") + monkeypatch.setenv("VAR3", "value3") + + template = "Start $VAR1 middle $VAR2 end $VAR3" + result = replace_variables(template) + + assert result == "Start value1 middle value2 end value3" + + +class TestCheckPlotInfoFile: + """Test the check_plot_info_file function.""" + + def test_plot_info_file_exists(self, tmp_path): + """Test when plot_info file already exists.""" + # Create a temporary plot_info file + plot_info_file = tmp_path / "plot_info" + plot_info_file.write_text("existing content") + + result = check_plot_info_file( + input_dir=str(tmp_path), + info_name="plot_info" + ) + + assert result is True + assert plot_info_file.read_text() == "existing content" + + @patch('boto3.client') + def test_create_plot_info_from_ssm_success(self, mock_boto3_client, tmp_path, monkeypatch): + """Test successful creation of plot_info from SSM parameter.""" + monkeypatch.setenv("FORECAST_DATE", "20240101") + monkeypatch.setenv("MODEL", "ICON") + + # Mock SSM client + mock_ssm_client = MagicMock() + mock_boto3_client.return_value = mock_ssm_client + + # Mock SSM response with template content + template_content = "Forecast: $FORECAST_DATE, Model: $MODEL" + mock_ssm_client.get_parameter.return_value = { + 'Parameter': { + 'Value': template_content + } + } + + result = check_plot_info_file( + input_dir=str(tmp_path), + info_name="plot_info", + ssm_parameter_path="/test/parameter" + ) + + assert result is True + + mock_boto3_client.assert_called_once_with('ssm') + plot_info_file = tmp_path / "plot_info" + assert plot_info_file.exists() + content = plot_info_file.read_text() + assert content == "Forecast: 20240101, Model: ICON" + + @patch('boto3.client') + def test_create_plot_info_template(self, mock_boto3_client, tmp_path, monkeypatch): + """Test plot_info template and vars substitutions.""" + monkeypatch.setenv("LAGRANTO_MODEL_BASE_TIME", "20240115_00") + monkeypatch.setenv("LM_NL_C_TTAG", "ICON-CH2-EPS") + monkeypatch.setenv("LM_NL_POLLONLM_C", "-170.0") + monkeypatch.setenv("LM_NL_POLLATLM_C", "43.0") + monkeypatch.setenv("LM_NL_STARTLON_TOT_C", "5.5") + monkeypatch.setenv("LM_NL_STARTLAT_TOT_C", "45.0") + monkeypatch.setenv("LM_NL_DLONLM_C", "0.02") + monkeypatch.setenv("LM_NL_DLATLM_C", "0.02") + monkeypatch.setenv("LM_NL_IELM_C", "250") + monkeypatch.setenv("LM_NL_JELM_C", "200") + + # Lagranto configuration template + template_content = ''' + Model base time: $LAGRANTO_MODEL_BASE_TIME + Model name: $LM_NL_C_TTAG + North pole longitude: $LM_NL_POLLONLM_C + North pole latitude: $LM_NL_POLLATLM_C + Start longitude: $LM_NL_STARTLON_TOT_C + Start latitude: $LM_NL_STARTLAT_TOT_C + Increment in longitudinal direction: $LM_NL_DLONLM_C + Increment in latitudinal direction: $LM_NL_DLATLM_C + Number of points in longitudinal direction: $LM_NL_IELM_C + Number of points in latitudinal direction: $LM_NL_JELM_C + ''' + + expected_content = ''' + Model base time: 20240115_00 + Model name: ICON-CH2-EPS + North pole longitude: -170.0 + North pole latitude: 43.0 + Start longitude: 5.5 + Start latitude: 45.0 + Increment in longitudinal direction: 0.02 + Increment in latitudinal direction: 0.02 + Number of points in longitudinal direction: 250 + Number of points in latitudinal direction: 200 + ''' + + # Mock SSM client + mock_ssm_client = MagicMock() + mock_boto3_client.return_value = mock_ssm_client + mock_ssm_client.get_parameter.return_value = { + 'Parameter': {'Value': template_content} + } + + result = check_plot_info_file( + input_dir=str(tmp_path), + info_name="plot_info" + ) + + assert result is True + + # Check content + plot_info_file = tmp_path / "plot_info" + content = plot_info_file.read_text() + assert content == expected_content + + @patch('boto3.client') + def test_create_plot_info_ssm_parameter_not_found(self, mock_boto3_client, tmp_path): + """Test SSM parameter not found error.""" + # Mock SSM client to raise exception + mock_ssm_client = MagicMock() + mock_boto3_client.return_value = mock_ssm_client + mock_ssm_client.get_parameter.side_effect = ClientError( + {'Error': {'Code': 'ParameterNotFound', 'Message': 'Parameter not found'}}, + 'GetParameter' + ) + + result = check_plot_info_file( + input_dir=str(tmp_path), + info_name="plot_info", + ssm_parameter_path="/nonexistent/parameter" + ) + + assert result is False + + plot_info_file = tmp_path / "plot_info" + assert not plot_info_file.exists() + +class TestCliIntegration: + """Test CLI integration with new functionality.""" + + def test_cli_with_existing_plot_info(self, tmp_path): + """Test CLI when plot_info file already exists.""" + # Create input directory with plot_info + input_dir = tmp_path / "input" + input_dir.mkdir() + plot_info = input_dir / "plot_info" + plot_info.write_text("existing plot_info") + + # Create some dummy trajectory files + (input_dir / "startf_test").write_text("start data") + (input_dir / "tra_geom_test").write_text("trajectory data") + + output_dir = tmp_path / "output" + + runner = CliRunner() + + # Mock the generate_pdf function since we're not testing that + with patch('pytrajplot.main.check_input_dir') as mock_check, \ + patch('pytrajplot.main.generate_pdf') as mock_generate: + + mock_check.return_value = ({}, {}) + + result = runner.invoke(cli, [ + str(input_dir), + str(output_dir), + '--skip-ssm-fallback' # Skip SSM to avoid AWS calls + ]) + + assert result.exit_code == 0 + assert "already exists" in result.output or result.exit_code == 0 + + def test_cli_skip_ssm_fallback(self, tmp_path): + """Test CLI with skip-ssm-fallback flag.""" + input_dir = tmp_path / "input" + input_dir.mkdir() + output_dir = tmp_path / "output" + + runner = CliRunner() + + with patch('pytrajplot.main.check_input_dir') as mock_check, \ + patch('pytrajplot.main.generate_pdf') as mock_generate: + + mock_check.return_value = ({}, {}) + + result = runner.invoke(cli, [ + str(input_dir), + str(output_dir), + '--skip-ssm-fallback' + ]) + + assert result.exit_code == 0 + + @patch('boto3.client') + def test_cli_ssm_failure_causes_exit(self, mock_boto3_client, tmp_path): + """Test CLI exits when SSM parameter fetch fails.""" + # Create input directory without plot_info + input_dir = tmp_path / "input" + input_dir.mkdir() + output_dir = tmp_path / "output" + + # Mock SSM client to fail + mock_ssm_client = MagicMock() + mock_boto3_client.return_value = mock_ssm_client + mock_ssm_client.get_parameter.side_effect = ClientError( + {'Error': {'Code': 'ParameterNotFound'}}, + 'GetParameter' + ) + + runner = CliRunner() + result = runner.invoke(cli, [ + str(input_dir), + str(output_dir) + ]) + + assert result.exit_code != 0 + assert "Missing plot_info file" in str(result.output) or result.exit_code != 0 From 8319ef03a27632661e857a1b14f3a0e9870475fc Mon Sep 17 00:00:00 2001 From: Paulina Grochal Date: Thu, 12 Mar 2026 13:03:57 +0100 Subject: [PATCH 2/3] DT-294 tests cleanup --- test/unit/test_plot_info.py | 34 +++------------------------------- 1 file changed, 3 insertions(+), 31 deletions(-) diff --git a/test/unit/test_plot_info.py b/test/unit/test_plot_info.py index 59dacae..8ae95ba 100644 --- a/test/unit/test_plot_info.py +++ b/test/unit/test_plot_info.py @@ -37,8 +37,7 @@ class TestCheckPlotInfoFile: """Test the check_plot_info_file function.""" def test_plot_info_file_exists(self, tmp_path): - """Test when plot_info file already exists.""" - # Create a temporary plot_info file + """Test with plot_info file in place.""" plot_info_file = tmp_path / "plot_info" plot_info_file.write_text("existing content") @@ -168,22 +167,16 @@ class TestCliIntegration: """Test CLI integration with new functionality.""" def test_cli_with_existing_plot_info(self, tmp_path): - """Test CLI when plot_info file already exists.""" - # Create input directory with plot_info + """Test CLI when plot_info file in place.""" input_dir = tmp_path / "input" input_dir.mkdir() plot_info = input_dir / "plot_info" plot_info.write_text("existing plot_info") - - # Create some dummy trajectory files - (input_dir / "startf_test").write_text("start data") - (input_dir / "tra_geom_test").write_text("trajectory data") - output_dir = tmp_path / "output" runner = CliRunner() - # Mock the generate_pdf function since we're not testing that + # Mock the generate_pdf function with patch('pytrajplot.main.check_input_dir') as mock_check, \ patch('pytrajplot.main.generate_pdf') as mock_generate: @@ -198,27 +191,6 @@ def test_cli_with_existing_plot_info(self, tmp_path): assert result.exit_code == 0 assert "already exists" in result.output or result.exit_code == 0 - def test_cli_skip_ssm_fallback(self, tmp_path): - """Test CLI with skip-ssm-fallback flag.""" - input_dir = tmp_path / "input" - input_dir.mkdir() - output_dir = tmp_path / "output" - - runner = CliRunner() - - with patch('pytrajplot.main.check_input_dir') as mock_check, \ - patch('pytrajplot.main.generate_pdf') as mock_generate: - - mock_check.return_value = ({}, {}) - - result = runner.invoke(cli, [ - str(input_dir), - str(output_dir), - '--skip-ssm-fallback' - ]) - - assert result.exit_code == 0 - @patch('boto3.client') def test_cli_ssm_failure_causes_exit(self, mock_boto3_client, tmp_path): """Test CLI exits when SSM parameter fetch fails.""" From c41cc88102b9fa68314d9ce57cb50180aa4d412e Mon Sep 17 00:00:00 2001 From: Paulina Grochal Date: Thu, 12 Mar 2026 13:33:06 +0100 Subject: [PATCH 3/3] DT-294 fix mypy typing issue --- pytrajplot/main.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pytrajplot/main.py b/pytrajplot/main.py index 36f7104..45ef792 100644 --- a/pytrajplot/main.py +++ b/pytrajplot/main.py @@ -1,5 +1,5 @@ """Command line interface of pytrajplot.""" -from typing import Tuple, Dict +from typing import Tuple, Dict, Optional import logging import os from pathlib import Path @@ -46,7 +46,7 @@ def replace_variables(template_content: str) -> str: return result -def check_plot_info_file(input_dir: str, info_name: str, ssm_parameter_path: str = None) -> bool: +def check_plot_info_file(input_dir: str, info_name: str, ssm_parameter_path: str | None = None) -> bool: """ Check if plot_info file exists in input directory. If not found, fetch from SSM parameter and create it replacing variables. @@ -235,7 +235,7 @@ def cli( language: str, domain: str, datatype: str, - ssm_parameter_path: str = None, + ssm_parameter_path: str | None = None, skip_ssm_fallback: bool = False, ) -> None: # Check if plot_info file exists (create from SSM if needed)