diff --git a/.github/actions/gitleaks/action.yml b/.github/actions/gitleaks/action.yml new file mode 100644 index 0000000..aaef2c8 --- /dev/null +++ b/.github/actions/gitleaks/action.yml @@ -0,0 +1,39 @@ +name: Gitleaks scan +description: Downloads a checksum-verified gitleaks release and scans the commits introduced by the current push or pull request for secrets. Extracted so this repo's own CI and monorepo's synced copy of this repo can call the identical steps locally -- see the parent engagement README's "Sharing CI with the monorepo" pattern, established for infrastructure and applied here. + +inputs: + working-directory: + description: Directory whose git history and .gitleaks.toml config to scan + required: false + default: "." + +runs: + using: composite + steps: + - name: Install gitleaks + shell: bash + working-directory: ${{ inputs.working-directory }} + run: | + set -euo pipefail + VER=8.21.2 + curl -sSfL --max-time 60 -o gitleaks.tar.gz "https://github.com/gitleaks/gitleaks/releases/download/v${VER}/gitleaks_${VER}_linux_x64.tar.gz" + echo "5bc41815076e6ed6ef8fbecc9d9b75bcae31f39029ceb55da08086315316e3ba gitleaks.tar.gz" | sha256sum -c - + tar -xzf gitleaks.tar.gz gitleaks + rm gitleaks.tar.gz + ./gitleaks version + + - name: Scan new commits for secrets + shell: bash + working-directory: ${{ inputs.working-directory }} + run: | + set -euo pipefail + if [ "${{ github.event_name }}" = "pull_request" ]; then + RANGE="${{ github.event.pull_request.base.sha }}..${{ github.sha }}" + elif [ "${{ github.event.before }}" = "0000000000000000000000000000000000000000" ]; then + RANGE="-1 ${{ github.sha }}" # first push of a branch: scan tip commit + else + RANGE="${{ github.event.before }}..${{ github.sha }}" + fi + echo "Scanning commit range: $RANGE" + ./gitleaks detect --source . --log-opts="$RANGE" \ + --config .gitleaks.toml --redact --verbose diff --git a/.github/actions/test-minimal-install/action.yml b/.github/actions/test-minimal-install/action.yml new file mode 100644 index 0000000..74219af --- /dev/null +++ b/.github/actions/test-minimal-install/action.yml @@ -0,0 +1,29 @@ +name: Test minimal install +description: Regression guard for the model_runner lazy-import bug -- a plain `pip install .`, with neither pandas nor pywr present, must still let create-synthetic-model run. Extracted so this repo's own CI and monorepo's synced copy of this repo can call the identical steps locally -- see the parent engagement README's "Sharing CI with the monorepo" pattern, established for infrastructure and applied here. + +inputs: + working-directory: + description: Directory containing the pywr_utils package to test + required: false + default: "." + +runs: + using: composite + steps: + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - shell: bash + working-directory: ${{ inputs.working-directory }} + run: pip install . + - shell: bash + working-directory: ${{ inputs.working-directory }} + run: | + pywr-utils create-synthetic-model --zones 2 --transfers 1 --output /tmp/synthetic.json --show-summary + test -f /tmp/synthetic.json + - name: confirm pandas and pywr are genuinely absent + shell: bash + working-directory: ${{ inputs.working-directory }} + run: | + ! python -c "import pandas" 2>/dev/null + ! python -c "import pywr" 2>/dev/null diff --git a/.github/actions/test/action.yml b/.github/actions/test/action.yml new file mode 100644 index 0000000..4f32981 --- /dev/null +++ b/.github/actions/test/action.yml @@ -0,0 +1,41 @@ +name: Test +description: Installs the package with its run/dev extras, runs the pytest suite, and uploads a coverage report. Extracted so this repo's own CI and monorepo's synced copy of this repo can call the identical steps locally -- see the parent engagement README's "Sharing CI with the monorepo" pattern, established for infrastructure and applied here. + +inputs: + working-directory: + description: Directory containing the pywr_utils package to test + required: false + default: "." + +runs: + using: composite + steps: + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - shell: bash + working-directory: ${{ inputs.working-directory }} + # pywr's lpsolve extension needs headers this runner doesn't have; glpk (also needed) is prebaked into the ci-runners image. + env: + PYWR_BUILD_LPSOLVE: "false" + run: pip install -e ".[run,dev]" + - shell: bash + working-directory: ${{ inputs.working-directory }} + run: pytest -v + - name: write coverage summary + # Fully native, no admin-gated feature or third-party action needed: coverage.py's own markdown formatter piped into GitHub Actions' own step-summary mechanism. Shows up on every run's summary page (and one click away from the PR checks list) regardless of whether Code Quality is ever enabled. + if: always() + shell: bash + working-directory: ${{ inputs.working-directory }} + run: | + echo '## Coverage' >> "$GITHUB_STEP_SUMMARY" + python -m coverage report --format=markdown >> "$GITHUB_STEP_SUMMARY" + - name: upload coverage report + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository + uses: actions/upload-code-coverage@v1 + with: + file: ${{ inputs.working-directory }}/coverage.xml + language: Python + label: code-coverage/pytest + # TEMPORARY: Code Quality isn't enabled for this repo yet (Settings > Security > Code quality) — that's an org/repo admin action, and neither CI nor a push-only collaborator can do it. Without this, the upload 403s and fails the whole job. Remove fail-on-error once an admin enables Code Quality and this step succeeds for real. + fail-on-error: false diff --git a/.github/actions/typecheck/action.yml b/.github/actions/typecheck/action.yml new file mode 100644 index 0000000..86f4dd9 --- /dev/null +++ b/.github/actions/typecheck/action.yml @@ -0,0 +1,40 @@ +name: Typecheck and lint +description: Runs mypy, ruff, vulture, bandit, and pip-audit against the Python package. Extracted so this repo's own CI and monorepo's synced copy of this repo can call the identical steps locally -- see the parent engagement README's "Sharing CI with the monorepo" pattern, established for infrastructure and applied here. + +inputs: + working-directory: + description: Directory containing the pywr_utils package to check + required: false + default: "." + +runs: + using: composite + steps: + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - shell: bash + working-directory: ${{ inputs.working-directory }} + run: pip install -e ".[dev]" + - shell: bash + working-directory: ${{ inputs.working-directory }} + run: mypy pywr_utils/ + - name: lint (ruff) + shell: bash + working-directory: ${{ inputs.working-directory }} + run: ruff check . + - name: dead-code scan (vulture) + shell: bash + working-directory: ${{ inputs.working-directory }} + run: vulture pywr_utils vulture_whitelist.py --min-confidence 80 + - name: security scan (bandit) + shell: bash + working-directory: ${{ inputs.working-directory }} + run: bandit -r pywr_utils + - name: dependency vulnerability scan (pip-audit) + shell: bash + working-directory: ${{ inputs.working-directory }} + # pip itself is not a declared dependency, just whatever setup-python provisions; keep it current so a stale bundled pip doesn't produce a false CVE finding here. + run: | + pip install -U pip + pip-audit --skip-editable diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..60c1fb5 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,23 @@ +version: 2 +updates: + - package-ecosystem: "pip" + directory: "/" + schedule: + interval: "weekly" + cooldown: + default-days: 7 + groups: + python-dependencies: + patterns: + - "*" + + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + cooldown: + default-days: 7 + groups: + github-actions: + patterns: + - "*" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..e276985 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,36 @@ +name: CI + +on: + push: + branches: [main] + pull_request: {} + workflow_dispatch: {} + +permissions: + contents: read + code-quality: write + +jobs: + typecheck: + runs-on: gha-ci-runners + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/typecheck + + test: + runs-on: gha-ci-runners + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.head.sha || github.sha }} + - uses: ./.github/actions/test + + test-minimal-install: + # Regression guard for the model_runner lazy-import bug: a plain `pip install .`, with neither pandas nor pywr present, must still let create-synthetic-model run — only the `run` extra needs either. + runs-on: gha-ci-runners + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/test-minimal-install diff --git a/.github/workflows/gitleaks.yml b/.github/workflows/gitleaks.yml new file mode 100644 index 0000000..09a9dca --- /dev/null +++ b/.github/workflows/gitleaks.yml @@ -0,0 +1,22 @@ +name: gitleaks + +# Scan only the NEW commits of each push / PR for secrets, matching the pattern already established on cluster-api-v2. No known historical-secret backlog exists in this repo, but scanning new commits only (rather than full history on every run) keeps this consistent with the rest of the org and cheap to run. +on: + pull_request: + push: + branches: [main] + workflow_dispatch: {} + +permissions: + contents: read + +jobs: + gitleaks: + # Self-hosted fleet: the org's billing block stops GitHub-hosted jobs, and the runner container has no sudo -- the binary runs from the workspace by path. + runs-on: gha-ci-runners + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: ./.github/actions/gitleaks diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c293e0e --- /dev/null +++ b/.gitignore @@ -0,0 +1,208 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +#uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock +#poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. +# https://pdm-project.org/en/latest/usage/project/#working-with-version-control +#pdm.lock +#pdm.toml +.pdm-python +.pdm-build/ + +# pixi +# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. +#pixi.lock +# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one +# in the .venv directory. It is recommended not to include this directory in version control. +.pixi + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Cursor +# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to +# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data +# refer to https://docs.cursor.com/context/ignore-files +.cursorignore +.cursorindexingignore + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ +models/ diff --git a/.gitleaks.toml b/.gitleaks.toml new file mode 100644 index 0000000..9417ecf --- /dev/null +++ b/.gitleaks.toml @@ -0,0 +1,3 @@ +# gitleaks config — extends the built-in rules. No repo-specific rules needed here (unlike cluster-api-v2's own .gitleaks.toml, which hardcodes patterns for its own known-leaked literal secrets). +[extend] +useDefault = true diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..9c18a23 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,25 @@ +default_install_hook_types: [pre-commit, pre-push] + +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.16.2 + hooks: + # Auto-fix lint issues on staged files and block the commit if any finding remains after the fix. The codebase reached zero findings under the real default ruleset (see the fix/remaining-ruff-findings PR chain), so this hook now enforces that state locally, matching CI's ruff check . rather than just warning about drift from it. + - id: ruff-check + args: [--fix] + + - repo: local + hooks: + - id: gitleaks-pre-push + name: gitleaks (scan commits about to be pushed) + language: system + pass_filenames: false + stages: [pre-push] + entry: >- + bash -c ' + UPSTREAM=$(git rev-parse --abbrev-ref --symbolic-full-name @{upstream} 2>/dev/null || echo ""); + if [ -n "$UPSTREAM" ]; then RANGE="${UPSTREAM}..HEAD"; else RANGE="-1 HEAD"; fi; + echo "Scanning commit range: $RANGE"; + gitleaks detect --source . --log-opts="$RANGE" --config .gitleaks.toml --redact --verbose + ' diff --git a/README.md b/README.md index c90803f..e2fd577 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ pywr-utils create-synthetic-model --zones 5 --transfers 10 1. Clone the repository 2. Install in development mode: `pip install -e .` 3. Run the CLI: `pywr-utils --help` +4. One-time, so the pre-push gitleaks hook actually runs: `pip install pre-commit && pre-commit install` (needs the `gitleaks` binary on PATH — `brew install gitleaks` or see the install step in `.github/workflows/gitleaks.yml`). CI runs the same scan on every push/PR regardless, this just catches a leak before it leaves your machine. ## License diff --git a/incremental_timing_results.csv b/incremental_timing_results.csv new file mode 100644 index 0000000..fe9d444 --- /dev/null +++ b/incremental_timing_results.csv @@ -0,0 +1,3 @@ +zones,transfers,setup_time,run_time,total_time,error +10,10,0.01154470443725586,0.050969600677490234,0.0625143051147461, +10,20,1.1486048698425293,1.2410962581634521,2.3897011280059814, diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..0e28833 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,55 @@ +[build-system] +requires = ["setuptools>=45", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "pywr-utils" +version = "0.1.0" +description = "Utilities for PYWR water resource modeling" +readme = "README.md" +requires-python = ">=3.7" +license = {text = "MIT"} +authors = [ + {name = "Your Name", email = "your.email@example.com"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", +] +dependencies = [ + "click>=8.0.0", +] + +[project.optional-dependencies] +run = ["pandas", "pywr"] +dev = ["pytest>=8.0", "pytest-cov>=5.0", "mypy>=1.10", "pandas-stubs", "vulture>=2.11", "bandit>=1.7", "pip-audit>=2.7", "ruff==0.16.2"] + +[project.scripts] +pywr-utils = "pywr_utils.cli:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["pywr_utils*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "--cov=pywr_utils --cov-report=term-missing --cov-report=xml" + +[tool.coverage.report] +# Ratchet, not an aspirational target: set at today's measured coverage (99.23%, after a full unit/integration/e2e pass on model_runner.py) minus a small buffer for environment variance, so coverage can't silently regress further without failing CI. Raise this number as coverage genuinely improves. +fail_under = 97 + +[[tool.mypy.overrides]] +# pywr ships no type stubs and no py.typed marker (not a types-pywr package situation — it's a niche scientific library that's simply untyped). +module = "pywr.*" +ignore_missing_imports = true + +[tool.ruff.lint] +# No select override here deliberately: ruff's own real default ruleset applies. This stays deterministic not by pinning the rule list but by pinning the ruff *version* itself exactly everywhere it runs (see the ruff dev dependency and .pre-commit-config.yaml), so "the default" always means the same concrete rule list until someone deliberately bumps the version pin. diff --git a/pywr_utils/__init__.py b/pywr_utils/__init__.py new file mode 100644 index 0000000..5069c99 --- /dev/null +++ b/pywr_utils/__init__.py @@ -0,0 +1,3 @@ +"""PYWR utilities package.""" + +__version__ = "0.1.0" diff --git a/pywr_utils/cli.py b/pywr_utils/cli.py new file mode 100644 index 0000000..564d5fb --- /dev/null +++ b/pywr_utils/cli.py @@ -0,0 +1,161 @@ +"""Command-line interface for pywr-utils.""" + +import os + +import click + +from .model_creation import SyntheticModelCreator + + +@click.group() +@click.version_option(version="0.1.0", prog_name="pywr-utils") +def main(): + """PYWR utilities command-line interface.""" + + +@main.command("create-synthetic-model") +@click.option( + "--zones", + "-z", + type=int, + default=10, + help="Number of zones to create in the synthetic model (default: 10)" +) +@click.option( + "--transfers", + "-t", + type=int, + default=10, + help="Number of transfers to create in the synthetic model (default: 10)" +) +@click.option( + "--output", + "-o", + type=click.Path(), + help="Output filename for the model (defaults to synthetic_model_z_t.json)" +) +@click.option( + "--show-summary", + "-s", + is_flag=True, + help="Show a summary of the created model" +) +@click.option( + "--run", + "-r", + is_flag=True, + help="Run the model after creation" +) +def create_synthetic_model(zones, transfers, output, show_summary, run): + """Create a synthetic PYWR model with specified zones and transfers.""" + click.echo(f"Creating synthetic model with {zones} zones and {transfers} transfers...") + + # Create the model + creator = SyntheticModelCreator(zones, transfers) + + # Generate default filename if not provided + if not output: + output = os.path.join("models", f"synthetic_model_{zones}z_{transfers}t.json") + + # Ensure output directory exists + output_dir = os.path.dirname(output) + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir) + + try: + # Save the model + saved_file = creator.save_model(output) + click.echo(f"✓ Model saved to: {saved_file}") + + # Show summary if requested + if show_summary: + click.echo("\n" + creator.get_model_summary()) + + # Show basic statistics + model = creator.build_model() + click.echo("\nModel Statistics:") + click.echo(f" • Total nodes: {len(model['nodes'])}") + click.echo(f" • Total edges: {len(model['edges'])}") + click.echo(f" • Parameters: {len(model['parameters'])}") + + click.echo("\n✓ Synthetic model creation completed successfully!") + + # Run the model if requested + if run: + click.echo("\nRunning model...") + from .model_runner import run_file + run_file(saved_file) + + except Exception as e: + click.echo(f"✗ Error creating model: {e!s}", err=True) + raise click.Abort() from e + + +@main.command("run-incremental-sizes") +@click.option( + "--max-zones", + type=int, + default=100, + help="Maximum number of zones to test (default: 100)" +) +@click.option( + "--zone-increment", + type=int, + default=10, + help="Increment for zones (default: 10)" +) +@click.option( + "--max-transfers", + type=int, + default=None, + help="Maximum number of transfers per zone size (defaults to max-zones)" +) +@click.option( + "--transfer-increment", + type=int, + default=10, + help="Increment for transfers (default: 10)" +) +@click.option( + "--output", + "-o", + type=click.Path(), + default="incremental_timing_results.csv", + help="Output CSV filename for results (default: incremental_timing_results.csv)" +) +@click.option( + "--show-graph", + "-g", + is_flag=True, + help="Display a terminal graph of setup time vs transfers" +) +def run_incremental_sizes_cmd(max_zones, zone_increment, max_transfers, transfer_increment, output, show_graph): + """Run incremental model sizes and measure setup times for performance analysis.""" + click.echo("Running incremental size tests...") + click.echo(f" Max zones: {max_zones}") + click.echo(f" Zone increment: {zone_increment}") + click.echo(f" Max transfers: {max_transfers if max_transfers else max_zones}") + click.echo(f" Transfer increment: {transfer_increment}") + click.echo(f" Output file: {output}") + + try: + from .model_runner import run_incremental_sizes + run_incremental_sizes( + max_zones=max_zones, + zone_increment=zone_increment, + max_transfers=max_transfers, + transfer_increment=transfer_increment, + output_csv=output, + show_graph=show_graph + ) + + click.echo("\n✓ Incremental size testing completed successfully!") + click.echo(f"Results saved to: {output}") + + except Exception as e: + click.echo(f"✗ Error running incremental sizes: {e!s}", err=True) + raise click.Abort() from e + + +if __name__ == "__main__": + main() diff --git a/pywr_utils/model_creation.py b/pywr_utils/model_creation.py new file mode 100644 index 0000000..4d1d45a --- /dev/null +++ b/pywr_utils/model_creation.py @@ -0,0 +1,238 @@ +"""Model creation utilities for PYWR synthetic models.""" + +from __future__ import annotations + +import json +from typing import Any + + +class SyntheticModelCreator: + """Creates synthetic PYWR models with specified inputs and transfers.""" + + def __init__(self, inputs: int, transfers: int): + """Initialize the model creator. + + Args: + inputs: Number of inputs to create + transfers: Number of transfers to create + """ + self.inputs = inputs + self.transfers = transfers + self.model_data = { + "metadata": { + "title": f"Synthetic Model - {inputs} inputs, {transfers} Transfers", + "description": "Auto-generated synthetic PYWR model", + "minimum_version": "1.0" + }, + "timestepper": { + "start": "2020-01-01", + "end": "2020-12-31", + "timestep": 1 + }, + "nodes": [], + "edges": [], + "parameters": [] + } + + def create_inputs(self) -> list[dict[str, Any]]: + """Create input nodes for the model. + + Returns: + List of input node definitions + """ + inputs = [] + for i in range(self.inputs): + input = { + "name": f"input_{i+1}", + "type": "input", + "max_flow": { + "type": "constant", + "value": 10.0 + i * 5 # Varying flow for each input + }, + "position": { + "geographic": [0.0 + i * 0.1, 51.0 + i * 0.1] + } + } + inputs.append(input) + return inputs + + def create_link_nodes(self) -> list[dict[str, Any]]: + """Create link nodes for each input. + + Returns: + List of link node definitions + """ + link_nodes = [] + for i in range(self.inputs): + link = { + "name": f"link_{i+1}", + "type": "link", + "position": { + "geographic": [0.05 + i * 0.1, 51.05 + i * 0.1] + } + } + link_nodes.append(link) + return link_nodes + + def create_demand_nodes(self) -> list[dict[str, Any]]: + """Create demand nodes for the model. + + Returns: + List of demand node definitions + """ + demand_nodes = [] + for i in range(self.inputs): + demand = { + "name": f"demand_{i+1}", + "type": "output", + "max_flow": 8.0 + i * 2, # Varying demand + "cost": -10.0, # Negative cost to prioritize meeting demand + "position": { + "geographic": [0.1 + i * 0.1, 51.1 + i * 0.1] + } + } + demand_nodes.append(demand) + return demand_nodes + + def create_transfer_nodes(self) -> list[dict[str, Any]]: + """Create transfer link nodes. + + Returns: + List of transfer node definitions + """ + transfer_nodes = [] + for i in range(self.transfers): + # Create transfers between adjacent inputs (cycling if needed) + from_input = i % self.inputs + to_input = (i + 1) % self.inputs + + transfer = { + "name": f"transfer_{i+1}", + "type": "link", + "max_flow": 15.0, + "cost": 1.0 + i * 0.5, # Varying transfer costs + "position": { + "geographic": [0.075 + from_input * 0.1, 51.075 + to_input * 0.1] + } + } + transfer_nodes.append(transfer) + return transfer_nodes + + def create_edges(self) -> list[list[str]]: + """Create edges connecting the nodes. + + Returns: + List of edge definitions [from_node, to_node] + """ + edges = [] + + # Connect inputs to the links + for i in range(self.inputs): + edges.append([f"input_{i+1}", f"link_{i+1}"]) + + # Connect links to demands + for i in range(self.inputs): + edges.append([f"link_{i+1}", f"demand_{i+1}"]) + + # Connect transfer links + for i in range(self.transfers): + from_input = i % self.inputs + to_input = (i + 1) % self.inputs + # From link to transfer + edges.append([f"link_{from_input+1}", f"transfer_{i+1}"]) + # From transfer to destination link + edges.append([f"transfer_{i+1}", f"link_{to_input+1}"]) + + return edges + + def create_parameters(self) -> dict[str, Any]: + """Create model parameters. + + Returns: + List of parameter definitions + """ + parameters: dict[str, Any] = {} + + return parameters + + def build_model(self) -> dict[str, Any]: + """Build the complete model structure. + + Returns: + Complete PYWR model dictionary + """ + # Create all node types + inputs = self.create_inputs() + link_nodes = self.create_link_nodes() + demand_nodes = self.create_demand_nodes() + transfer_nodes = self.create_transfer_nodes() + + # Combine all nodes + all_nodes = inputs + link_nodes + demand_nodes + transfer_nodes + self.model_data["nodes"] = all_nodes + + # Create edges + self.model_data["edges"] = self.create_edges() + + # Create parameters + self.model_data["parameters"] = self.create_parameters() + + return self.model_data + + def save_model(self, filename: str) -> str: + """Save the model to a JSON file. + + Args: + filename: Output filename + + Returns: + Path to the saved file + """ + model = self.build_model() + + with open(filename, 'w') as f: + json.dump(model, f, indent=2) + + return filename + + def get_model_summary(self) -> str: + """Get a summary of the created model. + + Returns: + String summary of the model + """ + model = self.build_model() + + summary_lines = [ + "Synthetic PYWR Model Summary:", + f" inputs: {self.inputs}", + f" Transfers: {self.transfers}", + f" Total Nodes: {len(model['nodes'])}", + f" - Catchments: {self.inputs}", + f" - Reservoirs: {self.inputs}", + f" - Demands: {self.inputs}", + f" - Transfers: {self.transfers}", + f" Total Edges: {len(model['edges'])}", + f" Parameters: {len(model['parameters'])}", + f" Time Period: {model['timestepper']['start']} to {model['timestepper']['end']}" + ] + + return "\n".join(summary_lines) + + +def create_synthetic_model(inputs: int, transfers: int, output_file: str | None = None) -> dict[str, Any]: + """Create a synthetic PYWR model. + + Args: + inputs: Number of inputs to create transfers: Number of transfers to create output_file: Optional output filename + + Returns: + Model dictionary + """ + creator = SyntheticModelCreator(inputs, transfers) + model = creator.build_model() + + if output_file: + creator.save_model(output_file) + + return model diff --git a/pywr_utils/model_runner.py b/pywr_utils/model_runner.py new file mode 100644 index 0000000..cc12f0f --- /dev/null +++ b/pywr_utils/model_runner.py @@ -0,0 +1,431 @@ +import json +import logging +import os +import time + +import pandas as pd +from pywr.model import Model +from pywr.recorders.progress import ProgressRecorder + +logging.basicConfig(level='INFO') +logger = logging.getLogger(__name__) + +def run_file(filepath, output_file=None): + """Run a PYWR model from file and return timing information.""" + pfr = PywrFileRunner() + pfr.load_pywr_model_from_file(filepath) + if not output_file: + #extract just the filename from the path + file_name = os.path.basename(filepath) + output_file = f"output_{file_name.split('.')[0]}.csv" + timing_info = pfr.run_pywr_model(output_file) + return timing_info + +class PywrFileRunner: + def __init__(self): + self.model = None + self.log = logging.getLogger(__name__) + + def load_pywr_model_from_file(self, filename, solver=None): + modelpath = os.path.abspath(filename) + + with open(modelpath, 'r') as f: + pywr_data = json.load(f) + + self.model = Model.load(pywr_data, solver=solver or 'glpk') + + def run_pywr_model(self, outfile="output_1.csv", create_csv=True): + """Run the PYWR model and return timing information. + + Args: + outfile: Output CSV filename (ignored if create_csv=False) create_csv: Whether to create CSV output file (default: True) + """ + setup_time = 0 + run_time = 0 + + try: + # Add a progress recorder to monitor the run. + ProgressRecorder(self.model) + + print("Setting up the model...") + # Time the setup phase + setup_start = time.time() + # Force a setup regardless of whether the model has been run or setup before + self.model.setup() + setup_time = time.time() - setup_start + print(f"Model setup complete in {setup_time:.3f} seconds. Running the model...") + + # Time the run phase + run_start = time.time() + run_stats = self.model.run() + run_time = time.time() - run_start + self.log.info(run_stats) + + # dataframes_to_output = [] columns = [] + columns = [] + data = [] + print("Model run complete. Processing results...") + for r in self.model.recorders: + if hasattr(r, 'values'): + try: + if sum(r.values()) > 0: + columns.append(r.name) + data.append(list(r.values())) + except NotImplementedError as e: + self.log.error(f"Error processing recorder {r.name}: {e}") + if 'total' in r.name: + print(f"{r.name}: {next(iter(r.values()))}") + + # Only create CSV if requested + if create_csv: + df = pd.DataFrame(data, index=columns) + df.to_csv(outfile) + print(f"Output written to {os.path.abspath(outfile)}") + + self.log.info("Model run complete.") + return { + 'setup_time': setup_time, + 'run_time': run_time, + 'total_time': setup_time + run_time + } + except Exception as e: + self.log.exception("Error running the PYWR model") + return { + 'setup_time': setup_time, + 'run_time': run_time, + 'total_time': setup_time + run_time, + 'error': str(e) + } + + +def run_incremental_sizes(max_zones=100, zone_increment=10, max_transfers=None, transfer_increment=10, output_csv="incremental_timing_results.csv", show_graph=False): + """ + Run incremental model sizes and measure setup times. + + Args: + max_zones: Maximum number of zones to test (default: 100) zone_increment: Increment for zones (default: 10) max_transfers: Maximum number of transfers per zone size (defaults to max_zones, but overrides zone limitation if specified) transfer_increment: Increment for transfers (default: 10) output_csv: Output CSV filename for results show_graph: Show a terminal graph of time vs transfers (default: False) + """ + from .model_creation import SyntheticModelCreator + + # Track whether max_transfers was explicitly specified + max_transfers_specified = max_transfers is not None + + # Default max_transfers to max_zones if not specified + if max_transfers is None: + max_transfers = max_zones + + results = [] + + print(f"Running incremental size tests up to {max_zones} zones and {max_transfers} transfers...") + + # Ensure output directory exists + output_dir = os.path.dirname(output_csv) if os.path.dirname(output_csv) else "." + if output_dir and output_dir != "." and not os.path.exists(output_dir): + os.makedirs(output_dir) + + # Iterate through zone sizes + for zones in range(zone_increment, max_zones + 1, zone_increment): + print(f"\nTesting with {zones} zones...") + + # Determine max transfers for this zone size + if max_transfers_specified: + # If max_transfers was explicitly specified, use it regardless of zone count + max_transfers_for_zones = max_transfers + else: + # If max_transfers was not specified, limit to zone count (original behavior) + max_transfers_for_zones = min(max_transfers, zones) + + for transfers in range(transfer_increment, max_transfers_for_zones + 1, transfer_increment): + print(f" Creating model with {zones} zones and {transfers} transfers...") + + try: + # Create the model + creator = SyntheticModelCreator(zones, transfers) + + # Ensure models directory exists + models_dir = "models" + if not os.path.exists(models_dir): + os.makedirs(models_dir) + print(f" Created models directory: {os.path.abspath(models_dir)}") + + # Create a named JSON file for the current model in the models folder + model_filename = os.path.join(models_dir, f"current_model_{zones}z_{transfers}t.json") + model_data = creator.build_model() + + print(f" Creating model file: {model_filename}") + with open(model_filename, 'w') as f: + json.dump(model_data, f, indent=2) + + print(f" Model file created: {os.path.abspath(model_filename)}") + + # Run the model and capture timing (without creating CSV output) + pfr = PywrFileRunner() + pfr.load_pywr_model_from_file(model_filename) + timing_info = pfr.run_pywr_model(create_csv=False) + + # Store results + result = { + 'zones': zones, + 'transfers': transfers, + 'setup_time': timing_info.get('setup_time', 0), + 'run_time': timing_info.get('run_time', 0), + 'total_time': timing_info.get('total_time', 0), + 'error': timing_info.get('error', '') + } + results.append(result) + + print(f" Setup time: {timing_info.get('setup_time', 0):.3f}s") + + # Show updated graph if requested (after each model run) + if show_graph and len(results) > 0: + current_df = pd.DataFrame(results) + _display_terminal_graph(current_df) + + # Clean up the model file after successful run + try: + os.unlink(model_filename) + print(f" Cleaned up: {model_filename}") + except OSError as e: + print(f" Warning: Could not delete {model_filename}: {e}") + + except Exception as e: + logger.exception("Error running model with %s zones and %s transfers", zones, transfers) + + # Clean up model file if it exists, even on error + models_dir = "models" + model_filename = os.path.join(models_dir, f"current_model_{zones}z_{transfers}t.json") + if os.path.exists(model_filename): + try: + os.unlink(model_filename) + print(f" Cleaned up after error: {model_filename}") + except OSError: + pass + + result = { + 'zones': zones, + 'transfers': transfers, + 'setup_time': 0, + 'run_time': 0, + 'total_time': 0, + 'error': str(e) + } + results.append(result) + + # Show updated graph even after errors if requested + if show_graph and len(results) > 0: + current_df = pd.DataFrame(results) + _display_terminal_graph(current_df) + + # Save results to CSV + if results: + df = pd.DataFrame(results) + output_path = os.path.join(output_dir, output_csv) + try: + df.to_csv(output_path, index=False) + print(f"\nResults saved to {os.path.abspath(output_path)}") + + # Verify file was created + if os.path.exists(output_path): + print(f"✓ Output file confirmed: {os.path.abspath(output_path)}") + file_size = os.path.getsize(output_path) + print(f" File size: {file_size} bytes") + else: + print(f"✗ Warning: Output file was not created at {output_path}") + + except Exception as e: + print(f"✗ Error saving CSV: {e!s}") + raise + + # Print summary statistics + print("\nSummary:") + print(f"Total models tested: {len(results)}") + if len(results) > 0: + print(f"Average setup time: {df['setup_time'].mean():.3f}s") + print(f"Max setup time: {df['setup_time'].max():.3f}s") + print(f"Min setup time: {df['setup_time'].min():.3f}s") + + return df + else: + print("No results to save.") + return pd.DataFrame() + + +def _display_terminal_graph(df): + """Display a 2D line chart of setup time vs transfers in the terminal.""" + print(f"\n{'='*80}") + print(f"SETUP TIME vs TRANSFERS (2D Line Chart) - {len(df)} models tested") + print(f"{'='*80}") + + if df.empty: + print("No data to plot.") + return + + # Get unique zone counts for different lines + zone_counts = sorted(df['zones'].unique()) + + # Determine overall ranges for consistent scaling + all_transfers = df['transfers'].tolist() + all_setup_times = df['setup_time'].tolist() + + if not all_transfers or not all_setup_times: + print("No valid data to plot.") + return + + min_transfers = min(all_transfers) + max_transfers = max(all_transfers) + min_time = min(all_setup_times) + max_time = max(all_setup_times) + + # Chart dimensions + chart_width = 60 + chart_height = 20 + + # Create the 2D grid + grid = [[' ' for _ in range(chart_width)] for _ in range(chart_height)] + + # Characters for different zone counts + zone_chars = ['●', '■', '▲', '♦', '★', '◆', '▼', '◄', '►', '♠'] + + # Plot data for each zone count + legend_info = [] + + for i, zones in enumerate(zone_counts): + zone_data = df[df['zones'] == zones].sort_values('transfers') + if zone_data.empty: + continue + + char = zone_chars[i % len(zone_chars)] + legend_info.append((zones, char)) + + # Plot points for this zone count + prev_x = prev_y = None + + for _, row in zone_data.iterrows(): + transfers = row['transfers'] + setup_time = row['setup_time'] + + # Scale to grid coordinates + if max_transfers > min_transfers: + x = int(((transfers - min_transfers) / (max_transfers - min_transfers)) * (chart_width - 1)) + else: + x = chart_width // 2 + + if max_time > min_time: + y = chart_height - 1 - int(((setup_time - min_time) / (max_time - min_time)) * (chart_height - 1)) + else: + y = chart_height // 2 + + # Ensure coordinates are within bounds + x = max(0, min(chart_width - 1, x)) + y = max(0, min(chart_height - 1, y)) + + # Plot the point + grid[y][x] = char + + # Draw line from previous point if exists + if prev_x is not None and prev_y is not None: + _draw_line(grid, prev_x, prev_y, x, y, char, chart_width, chart_height) + + prev_x, prev_y = x, y + + # Print the chart + print("\nSetup Time (s)") + print("^") + + # Y-axis labels and grid + for y in range(chart_height): + # Calculate the time value for this y position + time_val = max_time - ((y / (chart_height - 1)) * (max_time - min_time)) if max_time > min_time else min_time + + # Print y-axis label every few lines + if y % 4 == 0: + print(f"{time_val:6.3f} |", end="") + else: + print(" |", end="") + + # Print the row + print(''.join(grid[y])) + + # X-axis + print(f" +{'-' * chart_width}") + print(" ", end="") + + # X-axis labels + for i in range(0, chart_width, 10): + if max_transfers > min_transfers: + transfer_val = min_transfers + ((i / (chart_width - 1)) * (max_transfers - min_transfers)) + else: + transfer_val = min_transfers + print(f"{transfer_val:8.0f}", end=" ") + + print(f"\n{' ' * 35}Transfers") + + # Legend + print("\nLegend:") + for zones, char in legend_info: + print(f" {char} = {zones} zones") + + # Statistics + print("\nStatistics:") + print(f"Transfer range: {min_transfers} - {max_transfers}") + print(f"Setup time range: {min_time:.3f}s - {max_time:.3f}s") + + # Show correlation for each zone count + if len(zone_counts) > 1: + print("\nCorrelations by zone count:") + for zones in zone_counts: + zone_data = df[df['zones'] == zones] + if len(zone_data) > 1: + correlation = zone_data['transfers'].corr(zone_data['setup_time']) + if not pd.isna(correlation): + print(f" {zones} zones: {correlation:.3f}") + + # Overall correlation + if len(df) > 1: + overall_correlation = df['transfers'].corr(df['setup_time']) + if not pd.isna(overall_correlation): + print(f"\nOverall correlation: {overall_correlation:.3f}") + if overall_correlation > 0.7: + print(" → Strong positive correlation") + elif overall_correlation > 0.3: + print(" → Moderate positive correlation") + elif overall_correlation < -0.7: + print(" → Strong negative correlation") + elif overall_correlation < -0.3: + print(" → Moderate negative correlation") + else: + print(" → Weak correlation") + + print(f"{'='*80}") + + +def _draw_line(grid, x1, y1, x2, y2, char, width, height): + """Draw a line between two points on the grid using Bresenham's algorithm.""" + dx = abs(x2 - x1) + dy = abs(y2 - y1) + + x, y = x1, y1 + + x_inc = 1 if x1 < x2 else -1 + y_inc = 1 if y1 < y2 else -1 + + error = dx - dy + + while True: + # Ensure coordinates are within bounds + if 0 <= x < width and 0 <= y < height and grid[y][x] == ' ': # Don't overwrite existing points + grid[y][x] = '·' # Use a different character for line segments + + if x == x2 and y == y2: + break + + error2 = 2 * error + + if error2 > -dy: + error -= dy + x += x_inc + + if error2 < dx: + error += dx + y += y_inc + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c1a7b9a --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +click>=8.0.0 diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..ce63da7 --- /dev/null +++ b/setup.py @@ -0,0 +1,34 @@ +from setuptools import find_packages, setup + +setup( + name="pywr-utils", + version="0.1.0", + description="Utilities for PYWR water resource modeling", + author="Your Name", + author_email="your.email@example.com", + packages=find_packages(), + install_requires=[ + "click>=8.0.0", + ], + extras_require={ + "run": ["pandas", "pywr"], + "dev": ["pytest>=8.0", "pytest-cov>=5.0"], + }, + entry_points={ + "console_scripts": [ + "pywr-utils=pywr_utils.cli:main", + ], + }, + python_requires=">=3.7", + classifiers=[ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + ], +) diff --git a/tests/e2e/test_cli_commands.py b/tests/e2e/test_cli_commands.py new file mode 100644 index 0000000..b276596 --- /dev/null +++ b/tests/e2e/test_cli_commands.py @@ -0,0 +1,208 @@ +"""E2E tests for the pywr-utils command-line interface. + +These drive the CLI the way a user does — through the Click entrypoint with argv-style arguments — rather than calling the underlying functions. The final test in this file goes one step further and shells out to the installed `pywr-utils` console script, which is the only way to prove the packaging metadata in pyproject.toml actually produces a working command. +""" +import json +import runpy +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +from click.testing import CliRunner + +from pywr_utils import model_runner +from pywr_utils.cli import main + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def workdir(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + return tmp_path + + +def test_help_lists_both_commands(runner): + result = runner.invoke(main, ["--help"]) + + assert result.exit_code == 0 + assert "create-synthetic-model" in result.output + assert "run-incremental-sizes" in result.output + + +def test_version_is_reported(runner): + result = runner.invoke(main, ["--version"]) + + assert result.exit_code == 0 + assert "pywr-utils" in result.output + + +def test_create_synthetic_model_writes_the_requested_model(runner, workdir): + output = workdir / "model.json" + + result = runner.invoke(main, ["create-synthetic-model", "--zones", "3", "--transfers", "2", "--output", str(output)]) + + assert result.exit_code == 0, result.output + model = json.loads(output.read_text()) + assert f"Total nodes: {len(model['nodes'])}" in result.output + assert f"Total edges: {len(model['edges'])}" in result.output + assert f"Parameters: {len(model['parameters'])}" in result.output + + +def test_short_option_names_are_accepted(runner, workdir): + output = workdir / "short.json" + + result = runner.invoke(main, ["create-synthetic-model", "-z", "2", "-t", "1", "-o", str(output)]) + + assert result.exit_code == 0, result.output + assert output.exists() + + +def test_output_defaults_to_a_models_directory_named_after_the_size(runner, workdir): + result = runner.invoke(main, ["create-synthetic-model", "--zones", "3", "--transfers", "2"]) + + assert result.exit_code == 0, result.output + expected = workdir / "models" / "synthetic_model_3z_2t.json" + assert expected.exists(), f"expected the default path to be created, got {list(workdir.rglob('*.json'))}" + + +def test_a_missing_output_directory_is_created(runner, workdir): + output = workdir / "nested" / "deeper" / "model.json" + + result = runner.invoke(main, ["create-synthetic-model", "-z", "2", "-t", "1", "-o", str(output)]) + + assert result.exit_code == 0, result.output + assert output.exists() + + +def test_show_summary_prints_the_model_summary(runner, workdir): + output = workdir / "model.json" + args = ["create-synthetic-model", "-z", "4", "-t", "3", "-o", str(output)] + + without_summary = runner.invoke(main, args) + with_summary = runner.invoke(main, args + ["--show-summary"]) + + assert with_summary.exit_code == 0, with_summary.output + assert "Synthetic PYWR Model Summary:" in with_summary.output + assert "Synthetic PYWR Model Summary:" not in without_summary.output + assert "Total Nodes:" in with_summary.output + + +def test_the_run_flag_solves_the_model_it_just_created(runner, workdir): + output = workdir / "model.json" + + result = runner.invoke(main, ["create-synthetic-model", "-z", "3", "-t", "2", "-o", str(output), "--run"]) + + assert result.exit_code == 0, result.output + assert (workdir / f"output_{output.stem}.csv").exists(), "the --run flag should leave the solver's results behind" + + +def test_an_unwritable_output_path_aborts_with_an_error(runner, workdir): + blocker = workdir / "blocker" + blocker.write_text("this is a file, not a directory") + output = blocker / "model.json" + + result = runner.invoke(main, ["create-synthetic-model", "-z", "2", "-t", "1", "-o", str(output)]) + + assert result.exit_code != 0 + assert "Error creating model" in result.output + + +def test_run_incremental_sizes_sweeps_and_saves_results(runner, workdir): + result = runner.invoke( + main, + ["run-incremental-sizes", "--max-zones", "4", "--zone-increment", "2", "--transfer-increment", "2", "--output", "timings.csv"], + ) + + assert result.exit_code == 0, result.output + assert "Max zones: 4" in result.output + assert "completed successfully" in result.output + saved = (workdir / "timings.csv").read_text() + assert saved.splitlines()[0].split(",") == ["zones", "transfers", "setup_time", "run_time", "total_time", "error"] + + +def test_run_incremental_sizes_echoes_the_zone_count_when_max_transfers_is_omitted(runner, workdir): + result = runner.invoke( + main, + ["run-incremental-sizes", "--max-zones", "2", "--zone-increment", "2", "--transfer-increment", "2", "--output", "timings.csv"], + ) + + assert result.exit_code == 0, result.output + assert "Max transfers: 2" in result.output + + +def test_run_incremental_sizes_echoes_an_explicit_max_transfers(runner, workdir): + result = runner.invoke( + main, + [ + "run-incremental-sizes", + "--max-zones", "2", + "--zone-increment", "2", + "--max-transfers", "4", + "--transfer-increment", "2", + "--output", "timings.csv", + ], + ) + + assert result.exit_code == 0, result.output + assert "Max transfers: 4" in result.output + + +def test_show_graph_draws_the_chart(runner, workdir): + result = runner.invoke( + main, + ["run-incremental-sizes", "--max-zones", "2", "--zone-increment", "2", "--transfer-increment", "2", "--output", "timings.csv", "--show-graph"], + ) + + assert result.exit_code == 0, result.output + assert "SETUP TIME vs TRANSFERS" in result.output + + +def test_a_failing_sweep_aborts_with_an_error(runner, workdir, monkeypatch): + failure = "sweep exploded" + + def explode(**kwargs): + raise RuntimeError(failure) + + monkeypatch.setattr(model_runner, "run_incremental_sizes", explode) + + result = runner.invoke(main, ["run-incremental-sizes", "--max-zones", "2", "--zone-increment", "2", "--output", "timings.csv"]) + + assert result.exit_code != 0 + assert failure in result.output + + +def test_the_module_is_runnable_with_python_dash_m(monkeypatch): + """`python -m pywr_utils.cli` is a supported way in, so the __main__ guard has to actually dispatch to the Click group.""" + monkeypatch.setattr(sys, "argv", ["pywr-utils", "--help"]) + + with pytest.raises(SystemExit) as exit_info: + runpy.run_module("pywr_utils.cli", run_name="__main__") + + assert exit_info.value.code == 0 + + +def test_the_installed_console_script_creates_a_model(tmp_path): + """The truest end-to-end check available: run the `pywr-utils` command that pyproject.toml's [project.scripts] entry installs, in a real subprocess, and confirm it produces a model file.""" + # Resolve the script next to the interpreter running the tests rather than trusting PATH, which does not include the virtualenv's bin directory unless it has been activated. + alongside_interpreter = Path(sys.executable).parent / "pywr-utils" + executable = str(alongside_interpreter) if alongside_interpreter.exists() else shutil.which("pywr-utils") + assert executable, "pywr-utils console script was not found; the package is not installed" + output = tmp_path / "model.json" + + completed = subprocess.run( + [executable, "create-synthetic-model", "--zones", "3", "--transfers", "2", "--output", str(output), "--show-summary"], + capture_output=True, + text=True, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + assert "Synthetic PYWR Model Summary:" in completed.stdout + model = json.loads(output.read_text()) + assert {node["name"] for node in model["nodes"]} >= {"input_1", "link_1", "demand_1", "transfer_1"} diff --git a/tests/integration/test_incremental_sweep_failures.py b/tests/integration/test_incremental_sweep_failures.py new file mode 100644 index 0000000..2aa59c8 --- /dev/null +++ b/tests/integration/test_incremental_sweep_failures.py @@ -0,0 +1,125 @@ +"""INTEGRATION tests for the failure and cleanup paths of run_incremental_sizes. + +The sweep runs real models, so a failure part way through has to leave the working directory tidy and still surface a usable results frame. These tests drive the real function and inject failures at the filesystem boundary rather than stubbing the sweep itself. +""" +import os + +import pandas as pd +import pytest + +from pywr_utils import model_runner +from pywr_utils.model_runner import PywrFileRunner, run_incremental_sizes + +SWEEP_ARGS = {"max_zones": 2, "zone_increment": 2, "transfer_increment": 2, "output_csv": "timings.csv"} + + +@pytest.fixture +def workdir(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + return tmp_path + + +def test_an_undeletable_model_file_warns_without_failing_the_sweep(workdir, monkeypatch, capsys): + def refuse_to_delete(path): + raise OSError("permission denied") + + monkeypatch.setattr(os, "unlink", refuse_to_delete) + + results = run_incremental_sizes(**SWEEP_ARGS) + output = capsys.readouterr().out + + assert len(results) == 1 + assert (results["error"] == "").all(), "a cleanup failure must not be recorded as a model failure" + assert "Warning: Could not delete" in output + assert (workdir / "timings.csv").exists() + + +def test_a_model_that_fails_after_being_written_is_still_cleaned_up(workdir, capsys, monkeypatch): + failure = "load exploded" + + def explode(self, filename, solver=None): + raise RuntimeError(failure) + + monkeypatch.setattr(PywrFileRunner, "load_pywr_model_from_file", explode) + + results = run_incremental_sizes(**SWEEP_ARGS) + output = capsys.readouterr().out + + assert (results["error"] == failure).all() + assert "Cleaned up after error" in output + assert list((workdir / "models").iterdir()) == [], "the half-finished model file must not be left behind" + + +def test_a_failed_model_whose_file_also_resists_deletion_still_returns_results(workdir, monkeypatch): + """Both the run and the cleanup fail. The sweep must still hand back the error row rather than letting the cleanup's own exception escape and destroy the results of every model before it.""" + + def explode(self, filename, solver=None): + raise RuntimeError("load exploded") + + def refuse_to_delete(path): + raise OSError("permission denied") + + monkeypatch.setattr(PywrFileRunner, "load_pywr_model_from_file", explode) + monkeypatch.setattr(os, "unlink", refuse_to_delete) + + results = run_incremental_sizes(**SWEEP_ARGS) + + assert (results["error"] == "load exploded").all() + assert (workdir / "timings.csv").exists() + + +def test_the_chart_is_still_drawn_for_a_failed_model(workdir, monkeypatch): + rendered = [] + monkeypatch.setattr(model_runner, "_display_terminal_graph", lambda df: rendered.append(len(df))) + + def explode(self, filename, solver=None): + raise RuntimeError("load exploded") + + monkeypatch.setattr(PywrFileRunner, "load_pywr_model_from_file", explode) + + run_incremental_sizes(show_graph=True, **SWEEP_ARGS) + + assert rendered == [1], "an error row still belongs on the chart" + + +def test_a_failure_writing_the_summary_is_reported_and_propagated(workdir, monkeypatch, capsys): + """The sweep is expensive, so a failure to persist its results must be loud rather than swallowed into a silently missing file.""" + failure = "disk full" + + def refuse_to_write(self, *args, **kwargs): + raise OSError(failure) + + monkeypatch.setattr(pd.DataFrame, "to_csv", refuse_to_write) + + with pytest.raises(OSError, match=failure): + run_incremental_sizes(**SWEEP_ARGS) + + assert "Error saving CSV" in capsys.readouterr().out + + +def test_a_summary_that_silently_fails_to_appear_is_flagged(workdir, monkeypatch, capsys): + monkeypatch.setattr(pd.DataFrame, "to_csv", lambda self, *args, **kwargs: None) + + run_incremental_sizes(**SWEEP_ARGS) + output = capsys.readouterr().out + + assert "Output file was not created" in output + assert not (workdir / "timings.csv").exists() + + +def test_the_confirmed_summary_reports_its_size_on_disk(workdir, capsys): + run_incremental_sizes(**SWEEP_ARGS) + output = capsys.readouterr().out + + assert "Output file confirmed" in output + assert f"File size: {(workdir / 'timings.csv').stat().st_size} bytes" in output + + +@pytest.mark.xfail( + strict=True, + reason="run_incremental_sizes joins the output directory onto a path that already contains it, so an output_csv with any directory component is written to // and the whole sweep's results are lost to an OSError after every model has already been run.", +) +def test_an_output_path_with_a_directory_component_is_written_where_it_was_asked_for(workdir): + run_incremental_sizes(max_zones=2, zone_increment=2, transfer_increment=2, output_csv="results/timings.csv") + + assert (workdir / "results" / "timings.csv").exists() diff --git a/tests/integration/test_model_runner_api.py b/tests/integration/test_model_runner_api.py new file mode 100644 index 0000000..14e3f59 --- /dev/null +++ b/tests/integration/test_model_runner_api.py @@ -0,0 +1,255 @@ +"""INTEGRATION tests for pywr_utils.model_runner — PywrFileRunner, run_file and run_incremental_sizes driven as Python APIs against a real PyWr model. + +Nothing here is mocked. Each test generates a genuine model with SyntheticModelCreator, hands it to PyWr, and lets the real GLPK solver run it. That is the point of this layer: it proves the loading, solving, recorder-harvesting and CSV-writing path actually works, not merely that the functions were called. +""" +import json +import logging + +import pandas as pd +import pytest + +from pywr_utils import model_runner +from pywr_utils.model_creation import SyntheticModelCreator +from pywr_utils.model_runner import PywrFileRunner, run_file, run_incremental_sizes + +DEFAULT_SOLVER = "glpk" +ALTERNATIVE_SOLVER = "glpk-edge" + + +@pytest.fixture +def model_file(tmp_path): + """A small but genuinely solvable model — three inputs feeding three demands, with two transfers between them.""" + path = tmp_path / "model.json" + SyntheticModelCreator(3, 2).save_model(str(path)) + return path + + +@pytest.fixture +def model_file_with_recorders(tmp_path): + """The same model with two flow recorders attached, one of them named so it trips the 'total' reporting branch.""" + model = SyntheticModelCreator(3, 2).build_model() + model["recorders"] = { + "total_demand_1": {"type": "totalflownoderecorder", "node": "demand_1"}, + "demand_2_flow": {"type": "totalflownoderecorder", "node": "demand_2"}, + } + path = tmp_path / "model_with_recorders.json" + path.write_text(json.dumps(model)) + return path + + +def test_load_then_run_solves_the_model_and_reports_timings(model_file, tmp_path): + runner = PywrFileRunner() + runner.load_pywr_model_from_file(str(model_file)) + outfile = tmp_path / "results.csv" + + timings = runner.run_pywr_model(str(outfile)) + + assert "error" not in timings, timings + assert timings["setup_time"] > 0 + assert timings["run_time"] > 0 + assert timings["total_time"] == pytest.approx(timings["setup_time"] + timings["run_time"]) + assert outfile.exists() + + +def test_loading_builds_a_model_matching_the_json_on_disk(model_file): + expected = json.loads(model_file.read_text()) + runner = PywrFileRunner() + + runner.load_pywr_model_from_file(str(model_file)) + + assert {node.name for node in runner.model.nodes} == {node["name"] for node in expected["nodes"]} + + +def test_load_defaults_to_the_glpk_solver(model_file): + runner = PywrFileRunner() + + runner.load_pywr_model_from_file(str(model_file)) + + assert runner.model.solver.name == DEFAULT_SOLVER + + +def test_load_passes_an_explicit_solver_through_to_pywr(model_file): + """Regression guard: the solver argument used to be accepted and then silently discarded in favour of a hardcoded glpk, so a caller asking for a different solver got the default without any error.""" + runner = PywrFileRunner() + + runner.load_pywr_model_from_file(str(model_file), solver=ALTERNATIVE_SOLVER) + + assert ALTERNATIVE_SOLVER != DEFAULT_SOLVER, "this test is vacuous unless the two solvers really differ" + assert runner.model.solver.name == ALTERNATIVE_SOLVER + + +def test_an_unknown_solver_is_rejected_rather_than_falling_back(model_file): + runner = PywrFileRunner() + + with pytest.raises(KeyError): + runner.load_pywr_model_from_file(str(model_file), solver="not-a-real-solver") + + +def test_a_model_run_with_either_solver_produces_the_same_results(model_file, tmp_path): + """If the solver argument were being ignored, this test would still pass — but paired with the assertion above that the chosen solver is actually installed on the model, it shows the alternative solver is genuinely exercised rather than merely accepted.""" + results = {} + for solver in (DEFAULT_SOLVER, ALTERNATIVE_SOLVER): + runner = PywrFileRunner() + runner.load_pywr_model_from_file(str(model_file), solver=solver) + outfile = tmp_path / f"results_{solver}.csv" + timings = runner.run_pywr_model(str(outfile)) + assert "error" not in timings, timings + results[solver] = outfile.read_text() + + assert results[DEFAULT_SOLVER] == results[ALTERNATIVE_SOLVER] + + +def test_recorder_values_are_written_to_the_output_csv(model_file_with_recorders, tmp_path): + runner = PywrFileRunner() + runner.load_pywr_model_from_file(str(model_file_with_recorders)) + outfile = tmp_path / "recorded.csv" + + runner.run_pywr_model(str(outfile)) + + written = pd.read_csv(outfile, index_col=0) + reported = {recorder.name: list(recorder.values()) for recorder in runner.model.recorders if recorder.name in written.index} + assert reported, "expected at least one recorder to reach the CSV" + for name, values in reported.items(): + assert written.loc[name].tolist() == pytest.approx(values) + + +def test_recorders_named_total_are_echoed_to_stdout(model_file_with_recorders, tmp_path, capsys): + runner = PywrFileRunner() + runner.load_pywr_model_from_file(str(model_file_with_recorders)) + + runner.run_pywr_model(str(tmp_path / "out.csv")) + output = capsys.readouterr().out + + total_recorder = next(r for r in runner.model.recorders if "total" in r.name) + assert f"{total_recorder.name}: {next(iter(total_recorder.values()))}" in output + + +def test_a_recorder_that_cannot_report_values_is_logged_and_skipped(model_file, tmp_path, caplog): + """The progress recorder the runner attaches has no meaningful values(); it must be logged and stepped over rather than aborting the whole results harvest.""" + runner = PywrFileRunner() + runner.load_pywr_model_from_file(str(model_file)) + outfile = tmp_path / "out.csv" + + with caplog.at_level(logging.ERROR, logger="pywr_utils.model_runner"): + timings = runner.run_pywr_model(str(outfile)) + + assert "error" not in timings, timings + assert any("progressrecorder" in record.message for record in caplog.records) + assert "progressrecorder" not in outfile.read_text() + + +def test_create_csv_false_skips_writing_the_output_file(model_file, tmp_path): + runner = PywrFileRunner() + runner.load_pywr_model_from_file(str(model_file)) + outfile = tmp_path / "should_not_exist.csv" + + timings = runner.run_pywr_model(str(outfile), create_csv=False) + + assert "error" not in timings, timings + assert not outfile.exists() + + +def test_running_without_a_loaded_model_reports_the_error_instead_of_raising(tmp_path, caplog): + runner = PywrFileRunner() + + with caplog.at_level(logging.ERROR, logger="pywr_utils.model_runner"): + timings = runner.run_pywr_model(str(tmp_path / "out.csv")) + + assert timings["error"] + assert timings["total_time"] == pytest.approx(timings["setup_time"] + timings["run_time"]) + error_records = [record for record in caplog.records if record.levelno == logging.ERROR] + assert error_records, "the failure should be logged through the module's own logger, not only returned" + assert error_records[-1].name == "pywr_utils.model_runner", "the failure should be logged through the module's own named logger, not the root logger" + assert error_records[-1].exc_info is not None, "the traceback should be attached via exc_info rather than repeated in the log message" + + +def test_run_file_derives_the_output_name_from_the_input_filename(model_file, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + timings = run_file(str(model_file)) + + assert "error" not in timings, timings + assert (tmp_path / f"output_{model_file.stem}.csv").exists() + + +def test_run_file_honours_an_explicit_output_file(model_file, tmp_path): + outfile = tmp_path / "explicit.csv" + + timings = run_file(str(model_file), output_file=str(outfile)) + + assert "error" not in timings, timings + assert outfile.exists() + + +def test_incremental_sizes_runs_each_combination_and_saves_a_summary(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + results = run_incremental_sizes(max_zones=4, zone_increment=2, transfer_increment=2, output_csv="timings.csv") + + # Transfers are capped at the zone count when max_transfers is not given: 2 zones allow 2 transfers, 4 zones allow 2 and 4. + assert list(zip(results["zones"], results["transfers"])) == [(2, 2), (4, 2), (4, 4)] + assert (results["setup_time"] > 0).all() + assert (results["error"] == "").all(), results["error"].tolist() + assert pd.read_csv(tmp_path / "timings.csv").shape[0] == len(results) + + +def test_incremental_sizes_cleans_up_the_generated_model_files(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + run_incremental_sizes(max_zones=2, zone_increment=2, transfer_increment=2, output_csv="timings.csv") + + assert list((tmp_path / "models").iterdir()) == [] + + +def test_an_explicit_max_transfers_is_not_capped_by_the_zone_count(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + results = run_incremental_sizes(max_zones=2, zone_increment=2, max_transfers=4, transfer_increment=2, output_csv="timings.csv") + + assert list(results["transfers"]) == [2, 4], "an explicit max_transfers should override the zone-count cap" + + +def test_a_failing_model_is_recorded_as_an_error_row_rather_than_aborting_the_sweep(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + failure = "synthetic build failure" + + def explode(self): + raise RuntimeError(failure) + + monkeypatch.setattr(SyntheticModelCreator, "build_model", explode) + + results = run_incremental_sizes(max_zones=4, zone_increment=2, transfer_increment=2, output_csv="timings.csv") + + assert len(results) == 3 + assert (results["error"] == failure).all() + assert (results["total_time"] == 0).all() + + +def test_no_size_combinations_produces_an_empty_frame_and_no_csv(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + results = run_incremental_sizes(max_zones=1, zone_increment=2, transfer_increment=2, output_csv="timings.csv") + + assert results.empty + assert not (tmp_path / "timings.csv").exists() + + +def test_show_graph_renders_the_terminal_chart_after_each_run(tmp_path, monkeypatch, capsys): + monkeypatch.chdir(tmp_path) + rendered = [] + monkeypatch.setattr(model_runner, "_display_terminal_graph", lambda df: rendered.append(len(df))) + + run_incremental_sizes(max_zones=4, zone_increment=2, transfer_increment=2, output_csv="timings.csv", show_graph=True) + + # The chart is redrawn after every model, each time with one more row than the last. + assert rendered == [1, 2, 3] + + +def test_show_graph_off_renders_nothing(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + rendered = [] + monkeypatch.setattr(model_runner, "_display_terminal_graph", lambda df: rendered.append(len(df))) + + run_incremental_sizes(max_zones=2, zone_increment=2, transfer_increment=2, output_csv="timings.csv") + + assert rendered == [] diff --git a/tests/test_cli_lazy_import.py b/tests/test_cli_lazy_import.py new file mode 100644 index 0000000..8090335 --- /dev/null +++ b/tests/test_cli_lazy_import.py @@ -0,0 +1,60 @@ +"""Regression tests for the model_runner lazy-import fix. + +create-synthetic-model must not require pandas or pywr to be installed, since it only depends on model_creation (stdlib json/typing). Before the fix, cli.py imported model_runner (which needs both) at module load time, so every CLI invocation failed on a fresh install before pywr's own native build was ever reached. +""" +import importlib +import sys + +import pytest +from click.testing import CliRunner + +BLOCKED_MODULES = ("pandas", "pywr") + + +@pytest.fixture +def block_pandas_and_pywr(monkeypatch): + """Simulate a fresh install with neither pandas nor pywr present. + + Setting `sys.modules[name] = None` is the standard idiom for this: Python's import system raises ImportError for any name mapped to None in sys.modules, without needing to intercept `__import__` or actually uninstall the package. + """ + for name in BLOCKED_MODULES: + monkeypatch.setitem(sys.modules, name, None) + for name in ("pywr_utils.cli", "pywr_utils.model_runner"): + monkeypatch.delitem(sys.modules, name, raising=False) + yield + for name in ("pywr_utils.cli", "pywr_utils.model_runner"): + monkeypatch.delitem(sys.modules, name, raising=False) + + +def test_cli_module_imports_without_pandas_or_pywr(block_pandas_and_pywr): + module = importlib.import_module("pywr_utils.cli") + assert hasattr(module, "create_synthetic_model") + + +def test_create_synthetic_model_command_without_pandas_or_pywr(tmp_path, block_pandas_and_pywr): + cli = importlib.import_module("pywr_utils.cli") + output_path = tmp_path / "synthetic.json" + + runner = CliRunner() + result = runner.invoke( + cli.main, + ["create-synthetic-model", "--zones", "2", "--transfers", "1", "--output", str(output_path)], + ) + + assert result.exit_code == 0, result.output + assert output_path.exists() + + +def test_create_synthetic_model_command_run_flag_requires_pywr(tmp_path, block_pandas_and_pywr): + """--run does need pywr, and must fail with the real ImportError rather than something masking it — the lazy import only defers the cost, it doesn't remove the dependency for the code path that actually uses it.""" + cli = importlib.import_module("pywr_utils.cli") + output_path = tmp_path / "synthetic.json" + + runner = CliRunner() + result = runner.invoke( + cli.main, + ["create-synthetic-model", "--zones", "2", "--transfers", "1", "--output", str(output_path), "--run"], + ) + + assert result.exit_code != 0 + assert output_path.exists(), "model should still be saved before the run step is attempted" diff --git a/tests/test_vulture_whitelist.py b/tests/test_vulture_whitelist.py new file mode 100644 index 0000000..a3a9495 --- /dev/null +++ b/tests/test_vulture_whitelist.py @@ -0,0 +1,34 @@ +"""Regression tests for `vulture_whitelist.py`. + +`vulture_whitelist.py` is parsed as plain Python by both vulture (an AST walk, to register names as "used" so its dead-code scan doesn't flag them) and ruff (full name resolution). Every name the file references therefore has to actually be defined there via a real import or assignment -- a bare, unimported name is not valid Python and raises `NameError` the moment the module is executed, which is exactly what ruff's F821 (undefined-name) rule flags statically without needing to run it. +""" + +import runpy +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +WHITELIST_PATH = REPO_ROOT / "vulture_whitelist.py" + + +def test_vulture_whitelist_has_no_undefined_names(): + """Executing vulture_whitelist.py must not raise NameError. + + A stale entry -- e.g. a bare name left behind after the unused import it was whitelisting was itself removed elsewhere in the codebase -- is invalid Python: running the module raises NameError for that name. Vulture only walks the AST and never surfaces this, but the file must still be valid, executable Python. + """ + runpy.run_path(str(WHITELIST_PATH), run_name="__main__") + + +def test_no_f821_findings_anywhere_in_the_repo(): + """ruff's F821 (undefined-name) check must be clean across the whole repository. + + A direct, automated version of the same check this fix was verified against manually, so any future undefined-name regression -- in vulture_whitelist.py or anywhere else -- fails the suite instead of only being caught by a separate lint step. + """ + result = subprocess.run( + [sys.executable, "-m", "ruff", "check", "--select", "F821", str(REPO_ROOT)], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/unit/test_model_creation.py b/tests/unit/test_model_creation.py new file mode 100644 index 0000000..559330c --- /dev/null +++ b/tests/unit/test_model_creation.py @@ -0,0 +1,96 @@ +"""UNIT tests for pywr_utils.model_creation — SyntheticModelCreator and the create_synthetic_model helper, exercised in isolation with no pywr, pandas or filesystem-heavy dependencies. + +The class is pure data construction: given an input count and a transfer count it emits a PyWr model dictionary. Every assertion below derives its expected value from those two counts rather than hardcoding a total, so the tests stay honest if the node layout changes shape. +""" +import json + +import pytest + +from pywr_utils.model_creation import SyntheticModelCreator, create_synthetic_model + +# One input produces three nodes: an input, a link and a demand. Each transfer adds one further link node. +NODES_PER_INPUT = 3 + +# Each input contributes an input->link edge and a link->demand edge; each transfer contributes a link->transfer edge and a transfer->link edge. +EDGES_PER_INPUT = 2 +EDGES_PER_TRANSFER = 2 + + +@pytest.mark.parametrize("inputs,transfers", [(1, 0), (3, 2), (5, 12)]) +def test_build_model_node_and_edge_counts_follow_from_inputs_and_transfers(inputs, transfers): + model = SyntheticModelCreator(inputs, transfers).build_model() + + assert len(model["nodes"]) == inputs * NODES_PER_INPUT + transfers + assert len(model["edges"]) == inputs * EDGES_PER_INPUT + transfers * EDGES_PER_TRANSFER + + +def test_build_model_edges_only_reference_declared_nodes(): + model = SyntheticModelCreator(4, 6).build_model() + node_names = {node["name"] for node in model["nodes"]} + + for source, target in model["edges"]: + assert source in node_names + assert target in node_names + + +def test_transfers_wrap_around_when_they_outnumber_inputs(): + """A transfer connects link i to link i+1, cycling back to the first link once it runs off the end — so more transfers than inputs must still produce valid edges rather than an index error.""" + inputs = 3 + model = SyntheticModelCreator(inputs, inputs + 1).build_model() + + wrapping_transfer = f"transfer_{inputs + 1}" + feeds_the_transfer = [source for source, target in model["edges"] if target == wrapping_transfer] + fed_by_the_transfer = [target for source, target in model["edges"] if source == wrapping_transfer] + + # The (inputs + 1)th transfer has no link_{inputs + 1} to start from, so it wraps back to the first link and feeds the second. + assert feeds_the_transfer == ["link_1"] + assert fed_by_the_transfer == ["link_2"] + + +def test_create_transfer_nodes_divides_by_input_count_and_rejects_zero_inputs(): + """Transfer routing is modulo the input count, so zero inputs is not a degenerate-but-valid model — it is an error, and it must surface as one rather than silently producing an empty model.""" + creator = SyntheticModelCreator(0, 1) + + with pytest.raises(ZeroDivisionError): + creator.build_model() + + +def test_get_model_summary_reports_the_counts_of_the_model_it_describes(): + inputs, transfers = 4, 7 + creator = SyntheticModelCreator(inputs, transfers) + model = creator.build_model() + + summary = creator.get_model_summary() + + assert f"inputs: {inputs}" in summary + assert f"Transfers: {transfers}" in summary + assert f"Total Nodes: {len(model['nodes'])}" in summary + assert f"Total Edges: {len(model['edges'])}" in summary + assert f"Parameters: {len(model['parameters'])}" in summary + assert model["timestepper"]["start"] in summary + assert model["timestepper"]["end"] in summary + + +def test_save_model_writes_json_that_round_trips_to_the_built_model(tmp_path): + creator = SyntheticModelCreator(2, 2) + destination = tmp_path / "model.json" + + returned_path = creator.save_model(str(destination)) + + assert returned_path == str(destination) + assert json.loads(destination.read_text()) == creator.build_model() + + +def test_create_synthetic_model_returns_the_model_without_writing_when_no_output_given(tmp_path): + model = create_synthetic_model(3, 1) + + assert model == SyntheticModelCreator(3, 1).build_model() + assert list(tmp_path.iterdir()) == [] + + +def test_create_synthetic_model_writes_the_returned_model_when_an_output_is_given(tmp_path): + destination = tmp_path / "written.json" + + model = create_synthetic_model(3, 1, output_file=str(destination)) + + assert json.loads(destination.read_text()) == model diff --git a/tests/unit/test_terminal_graph.py b/tests/unit/test_terminal_graph.py new file mode 100644 index 0000000..20e324f --- /dev/null +++ b/tests/unit/test_terminal_graph.py @@ -0,0 +1,161 @@ +"""UNIT tests for the terminal-plotting helpers in pywr_utils.model_runner — _display_terminal_graph and _draw_line. + +Both are pure functions over an in-memory DataFrame or grid: no PyWr model, no solver, no filesystem. They are tested here in isolation from the run machinery that normally feeds them. +""" +import pandas as pd +import pytest + +from pywr_utils.model_runner import _display_terminal_graph, _draw_line + + +def timing_frame(zones, transfers, setup_times): + return pd.DataFrame({"zones": zones, "transfers": transfers, "setup_time": setup_times}) + + +def test_empty_frame_reports_no_data_rather_than_raising(capsys): + _display_terminal_graph(pd.DataFrame(columns=["zones", "transfers", "setup_time"])) + + assert "No data to plot." in capsys.readouterr().out + + +# Each dataset is a permutation of the same four setup times against the same four transfer counts, chosen so the Pearson correlation is an exact, hand-checkable value that lands squarely inside one classification band. +@pytest.mark.parametrize( + "setup_times,expected_correlation,expected_label", + [ + ([1.0, 2.0, 3.0, 4.0], 1.0, "Strong positive correlation"), + ([2.0, 1.0, 4.0, 3.0], 0.6, "Moderate positive correlation"), + ([2.0, 4.0, 1.0, 3.0], 0.0, "Weak correlation"), + ([3.0, 4.0, 1.0, 2.0], -0.6, "Moderate negative correlation"), + ([4.0, 3.0, 2.0, 1.0], -1.0, "Strong negative correlation"), + ], +) +def test_overall_correlation_is_reported_and_classified(capsys, setup_times, expected_correlation, expected_label): + transfers = [1, 2, 3, 4] + frame = timing_frame([10] * len(transfers), transfers, setup_times) + + assert frame["transfers"].corr(frame["setup_time"]) == pytest.approx(expected_correlation) + + _display_terminal_graph(frame) + output = capsys.readouterr().out + + assert f"Overall correlation: {expected_correlation:.3f}" in output + assert expected_label in output + + +def test_single_zone_count_omits_the_per_zone_correlation_breakdown(capsys): + _display_terminal_graph(timing_frame([10, 10], [1, 2], [0.1, 0.2])) + + assert "Correlations by zone count:" not in capsys.readouterr().out + + +def test_each_zone_count_gets_its_own_legend_marker_and_correlation(capsys): + frame = timing_frame( + zones=[10, 10, 10, 20, 20, 20], + transfers=[1, 2, 3, 1, 2, 3], + setup_times=[0.1, 0.2, 0.3, 0.5, 0.4, 0.6], + ) + + _display_terminal_graph(frame) + output = capsys.readouterr().out + + legend_markers = [line.split("=")[0].strip() for line in output.splitlines() if line.strip().endswith("zones")] + assert len(legend_markers) == frame["zones"].nunique() + assert len(set(legend_markers)) == len(legend_markers), "each zone count needs a distinct marker" + + assert "Correlations by zone count:" in output + for zones, group in frame.groupby("zones"): + expected = group["transfers"].corr(group["setup_time"]) + assert f"{zones} zones: {expected:.3f}" in output + + for marker in legend_markers: + assert marker in output, "every legend marker should also appear plotted on the chart" + + +def test_reported_ranges_come_from_the_data(capsys): + frame = timing_frame([10, 10, 20], [5, 40, 25], [0.25, 1.5, 0.75]) + + _display_terminal_graph(frame) + output = capsys.readouterr().out + + assert "Transfer range: 5 - 40" in output + assert "Setup time range: 0.250s - 1.500s" in output + assert f"{len(frame)} models tested" in output + + +def test_single_point_is_plotted_without_a_correlation(capsys): + """One row means no range to scale against and no correlation to compute; the chart must still render rather than dividing by a zero span.""" + _display_terminal_graph(timing_frame([10], [7], [0.4])) + output = capsys.readouterr().out + + assert "Transfer range: 7 - 7" in output + assert "Overall correlation" not in output + + +def test_constant_setup_times_suppress_the_undefined_correlation(capsys): + """Zero variance in setup_time makes the Pearson correlation NaN — printing 'nan' would be worse than printing nothing.""" + frame = timing_frame([10, 10, 10], [1, 2, 3], [0.5, 0.5, 0.5]) + + assert pd.isna(frame["transfers"].corr(frame["setup_time"])) + + _display_terminal_graph(frame) + output = capsys.readouterr().out + + assert "Overall correlation" not in output + assert "nan" not in output.lower() + + +def blank_grid(width, height): + return [[" " for _ in range(width)] for _ in range(height)] + + +def test_draw_line_marks_a_diagonal_path_between_the_endpoints(): + width = height = 5 + grid = blank_grid(width, height) + + _draw_line(grid, 0, 0, 4, 4, "●", width, height) + + for i in range(5): + assert grid[i][i] == "·", f"expected the diagonal cell ({i}, {i}) to be on the drawn path" + assert grid[0][4] == " ", "cells off the path must be left blank" + + +def test_draw_line_marks_a_horizontal_path(): + width, height = 6, 3 + grid = blank_grid(width, height) + + _draw_line(grid, 1, 1, 4, 1, "●", width, height) + + assert "".join(grid[1]) == " ···· " + assert "".join(grid[0]).strip() == "" + + +def test_draw_line_does_not_overwrite_existing_plotted_points(): + """Line segments are connective tissue between real data points — a segment must never erase the marker of a point it passes through.""" + width = height = 5 + grid = blank_grid(width, height) + grid[2][2] = "●" + + _draw_line(grid, 0, 0, 4, 4, "■", width, height) + + assert grid[2][2] == "●" + + +def test_draw_line_clips_to_the_grid_instead_of_writing_out_of_bounds(): + """Negative coordinates must be skipped, not passed through to the list — Python would silently index from the end of the row and paint a segment on the wrong side of the chart.""" + width = height = 4 + grid = blank_grid(width, height) + + _draw_line(grid, 0, -3, 0, 1, "●", width, height) + + marked = {(x, y) for y in range(height) for x in range(width) if grid[y][x] != " "} + assert marked == {(0, 0), (0, 1)} + + +def test_draw_line_between_identical_points_marks_only_that_point(): + width = height = 3 + grid = blank_grid(width, height) + + _draw_line(grid, 1, 1, 1, 1, "●", width, height) + + assert grid[1][1] == "·" + assert sum(cell != " " for row in grid for cell in row) == 1 diff --git a/vulture_whitelist.py b/vulture_whitelist.py new file mode 100644 index 0000000..e69de29