From 58fb4027a17337c472e5d1c22cd3e66eb65447db Mon Sep 17 00:00:00 2001 From: Aliaksei Kharlap Date: Wed, 12 Aug 2026 21:12:26 +0300 Subject: [PATCH] refactor stock plugin to use utils package --- influxdata/library/plugin_library.json | 4 +- influxdata/stock_plugin/README.md | 63 +- influxdata/stock_plugin/manifest.toml | 8 +- influxdata/stock_plugin/requirements-dev.txt | 4 + influxdata/stock_plugin/requirements.txt | 1 + influxdata/stock_plugin/stock_plugin.py | 544 +++++++------ .../stock_plugin/stock_plugin.toml.example | 4 +- influxdata/stock_plugin/test_stock_plugin.py | 747 ++++++++++++++++++ 8 files changed, 1096 insertions(+), 279 deletions(-) create mode 100644 influxdata/stock_plugin/requirements-dev.txt create mode 100644 influxdata/stock_plugin/test_stock_plugin.py diff --git a/influxdata/library/plugin_library.json b/influxdata/library/plugin_library.json index fcddf40..9d9129f 100644 --- a/influxdata/library/plugin_library.json +++ b/influxdata/library/plugin_library.json @@ -380,8 +380,8 @@ "author": "InfluxData", "docs_file_link": "https://github.com/influxdata/influxdb3_plugins/blob/main/influxdata/stock_plugin/README.md", "required_plugins": [], - "required_libraries": ["yfinance", "pandas_market_calendars"], - "last_update": "2026-06-10", + "required_libraries": ["yfinance", "pandas_market_calendars", "influxdata-plugin-utils>=0.3.0"], + "last_update": "2026-08-12", "trigger_types_supported": ["scheduler"] }, { diff --git a/influxdata/stock_plugin/README.md b/influxdata/stock_plugin/README.md index 8cb0fcb..1dc4b97 100644 --- a/influxdata/stock_plugin/README.md +++ b/influxdata/stock_plugin/README.md @@ -20,12 +20,16 @@ This plugin includes a JSON metadata schema in its docstring that defines the su ### Optional parameters -| Parameter | Type | Default | Description | -|---------------|--------|--------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `database` | string | `stocks` | Target database for writes. Overrides the TOML `database` key if both are set. | -| `portfolio` | string | `AAPL:1\|MSFT:1\|GOOG:1` | Inline holdings: pipe-separated `SYMBOL:QUANTITY[:PORTFOLIO_NAME]` entries (e.g. `AAPL:10:401k\|MSFT:5:401k\|GOOG:2.5:brokerage`). Portfolio defaults to `main`. When omitted and no TOML config is found, falls back to the default shown. | -| `categories` | string | none | Inline category map: pipe-separated `PORTFOLIO:CATEGORY` entries (e.g. `401k:Retirement\|brokerage:Investment`). | -| `config_path` | string | `stock_plugin.toml` | Path to the TOML config file, relative to the InfluxDB plugin directory (or absolute). | +| Parameter | Type | Default | Description | +|-----------------------------|---------|--------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `database` | string | `stocks` | Target database for writes. | +| `portfolio` | string | `AAPL:1\|MSFT:1\|GOOG:1` | Inline holdings: pipe-separated `SYMBOL:QUANTITY[:PORTFOLIO_NAME]` entries (e.g. `AAPL:10:401k\|MSFT:5:401k\|GOOG:2.5:brokerage`). Portfolio defaults to `main`. When omitted and no TOML config is found, falls back to the default shown. | +| `categories` | string | none | Inline category map: pipe-separated `PORTFOLIO:CATEGORY` entries (e.g. `401k:Retirement\|brokerage:Investment`). | +| `config_path` | string | `stock_plugin.toml` | Path to the TOML config file, relative to the InfluxDB plugin directory (or absolute). The default file is loaded when it exists; an explicit `config_path` that does not exist is an error. | +| `write_during_closed_hours` | boolean | `true` | See the TOML table below. Also settable as a trigger argument. | +| `mutual_fund_check_time` | string | `18:00` | See the TOML table below. Also settable as a trigger argument. | +| `market_calendar` | string | `NYSE` | See the TOML table below. Also settable as a trigger argument. | +| `market_timezone` | string | `America/New_York` | See the TOML table below. Also settable as a trigger argument. | *If neither `portfolio` nor a TOML file with `[holdings.]` is provided, the plugin runs with the default holdings `AAPL:1|MSFT:1|GOOG:1` in the `main` portfolio.* @@ -33,15 +37,26 @@ This plugin includes a JSON metadata schema in its docstring that defines the su The TOML file is the recommended way to configure anything more than a handful of holdings. The plugin reads it from `config_path` (default: `/stock_plugin.toml`). -| Key | Type | Default | Description | -|-----------------------------|---------|----------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `database` | string | `stocks` | Target database for writes. | -| `write_during_closed_hours` | boolean | `true` | When `false`, stocks/ETFs are skipped outside the configured exchange's regular session (the calendar handles holidays and early closes). Mutual funds always follow their own daily check schedule. | -| `mutual_fund_check_time` | string | `"18:00"` | Time of day in `market_timezone` after which the plugin fetches mutual fund NAV. Mutual funds are fetched at most once per local calendar day, at the first tick at or after this time. Bootstrap exception: a mutual fund with no cached asset type is fetched on its first tick regardless of time. | -| `market_calendar` | string | `"NYSE"` | Exchange calendar used for the market-hours check. Any name accepted by [pandas_market_calendars](https://pandas-market-calendars.readthedocs.io/) (e.g. `NYSE`, `LSE`, `TSX`, `JPX`, `XETR`, `ASX`, `HKEX`). | -| `market_timezone` | string | `"America/New_York"` | IANA timezone for the exchange's local time. Used for `mutual_fund_check_time` comparisons and for resolving the "today" date the calendar consults. | -| `[portfolio_categories]` | table | empty | Maps portfolio name to category name. Portfolios not listed are uncategorized (omitted from `category_totals`). | -| `[holdings.]` | table | default holdings | Holdings for each portfolio. Each entry is `SYMBOL = quantity`. Fractional quantities supported. Quote symbols containing dots, for example `"VOD.L" = 10`. Duplicate same-symbol entries in one portfolio are aggregated. The portfolio name `_total` is reserved. When no `[holdings.*]` section is present, the plugin falls back to `AAPL:1\|MSFT:1\|GOOG:1`. | +Trigger arguments take precedence over TOML keys of the same name, so a TOML file can hold the full portfolio shape while a trigger argument overrides a single setting. + +Holdings and categories are the exception, because each is spelled differently per source: + +| Setting | Trigger argument | TOML | +|------------|---------------------------------------|--------------------------------| +| Holdings | `portfolio=AAPL:10:401k\|MSFT:5:401k` | `[holdings.401k]` tables | +| Categories | `categories=401k:Retirement` | `[portfolio_categories]` table | + +Each spelling is read only from its own source: a top-level `portfolio` or `categories` key in the TOML file is ignored, as is a trigger argument named `holdings` or `portfolio_categories`. When both sources are present, the trigger argument replaces the TOML tables entirely rather than merging with them. + +| Key | Type | Default | Description | +|-----------------------------|---------|----------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `database` | string | `stocks` | Target database for writes. | +| `write_during_closed_hours` | boolean | `true` | When `false`, stocks/ETFs are skipped outside the configured exchange's regular session (the calendar handles holidays and early closes). Mutual funds always follow their own daily check schedule. | +| `mutual_fund_check_time` | string | `"18:00"` | Time of day in `market_timezone` after which the plugin fetches mutual fund NAV. Mutual funds are fetched at most once per local calendar day, at the first tick at or after this time. Bootstrap exception: a mutual fund with no cached asset type is fetched on its first tick regardless of time. | +| `market_calendar` | string | `"NYSE"` | Exchange calendar used for the market-hours check. Any name accepted by [pandas_market_calendars](https://pandas-market-calendars.readthedocs.io/) (e.g. `NYSE`, `LSE`, `TSX`, `JPX`, `XETR`, `ASX`, `HKEX`). | +| `market_timezone` | string | `"America/New_York"` | IANA timezone for the exchange's local time. Used for `mutual_fund_check_time` comparisons and for resolving the "today" date the calendar consults. | +| `[portfolio_categories]` | table | empty | Maps portfolio name to category name. Portfolios not listed are uncategorized (omitted from `category_totals`). | +| `[holdings.]` | table | default holdings | Holdings for each portfolio. Each entry is `SYMBOL = quantity`. Fractional quantities supported; the quantity must be a finite number (`inf` and `nan` are rejected). Quote symbols containing dots, for example `"VOD.L" = 10`. Duplicate same-symbol entries in one portfolio are aggregated. The portfolio name `_total` is reserved. When no `[holdings.*]` section is present, the plugin falls back to `AAPL:1\|MSFT:1\|GOOG:1`. | The trigger spec is the source of truth for cadence. For example, `--trigger-spec "every:15m"` runs the plugin every 15 minutes. @@ -52,9 +67,11 @@ The trigger spec is the source of truth for cadence. For example, `--trigger-spe ## Software requirements - **InfluxDB 3 Core/Enterprise**: with the Processing Engine enabled. +- **Python 3.11 or higher** - **Python packages** (installed into the plugin venv): - `yfinance` — Yahoo Finance scraper for price data - - `pandas_market_calendars` — NYSE calendar for accurate market-hours and holiday gating + - `pandas_market_calendars` — exchange calendars for accurate market-hours and holiday gating + - `influxdata-plugin-utils>=0.3.0` — shared configuration, parsing, and write helpers ### Installation steps @@ -71,7 +88,7 @@ The trigger spec is the source of truth for cadence. For example, `--trigger-spe 2. Install required Python packages: ```bash - influxdb3 install package yfinance pandas_market_calendars + influxdb3 install package yfinance pandas_market_calendars influxdata-plugin-utils ``` 3. Copy `stock_plugin.toml.example` to `/stock_plugin.toml` and edit it with your holdings and categories. @@ -173,9 +190,9 @@ WHERE missing_symbols = 0; #### `process_scheduled_call(influxdb3_local, call_time, args)` -Entry point for the scheduled trigger. Resolves the plugin directory from `INFLUXDB3_PLUGIN_DIR`, normalizes `call_time` to UTC, and delegates to `_main` with the runtime-injected `LineBuilder` and the live `influxdb3_local`. All side-effecting work lives in `_main` so its logic can be reasoned about with injected dependencies. +Entry point for the scheduled trigger. Stamps one UTC timestamp for the whole run and delegates to `_main` with the runtime-injected `LineBuilder` and the live `influxdb3_local`. All side-effecting work lives in `_main` so its logic can be reasoned about with injected dependencies. -#### `_main(local, args, fetcher, line_builder_cls, plugin_dir, now_ns)` +#### `_main(local, args, fetcher, line_builder_cls, now_ns, task_id)` Drives the full plugin flow: @@ -185,9 +202,13 @@ Drives the full plugin flow: 4. Build carry-forward `HoldingRow`s for intentionally-skipped symbols whose last known price is cached. 5. Aggregate per-portfolio totals + a grand `_total` row. 6. Aggregate per-category totals across portfolios. -7. Emit line protocol via `LineBuilder` for `stock_holdings`, `portfolio_totals`, and `category_totals`. +7. Write `stock_holdings`, `portfolio_totals`, and `category_totals` as a single batched payload. 8. Log a single summary line. +#### `resolve_config(args)` + +Merges the TOML file with the trigger arguments and validates the result. The TOML path comes from `config_path`; relative paths resolve against the plugin directory (`PLUGIN_DIR`, `INFLUXDB3_PLUGIN_DIR`, or the `VIRTUAL_ENV` parent). Returns a `ResolvedConfig`, raising `ValueError` on any invalid value. + ### Measurements and fields #### `stock_holdings` @@ -241,7 +262,7 @@ When `write_during_closed_hours` is false, the plugin uses `pandas_market_calend **Solution:** Verify the ticker symbol and check whether Yahoo Finance exposes current price data for that instrument. -The plugin carries forward the last known price for intentionally skipped symbols, but it cannot value a new holding until the first successful fetch. +A symbol whose price is missing or not a finite number is counted as a fetch failure and reported in the summary log; the optional `previous_close`, `day_open`, `day_high`, and `day_low` fields are simply omitted when unusable. The plugin carries forward the last known price for intentionally skipped symbols, but it cannot value a new holding until the first successful fetch. ### Debugging tips diff --git a/influxdata/stock_plugin/manifest.toml b/influxdata/stock_plugin/manifest.toml index 559fe36..b973be4 100644 --- a/influxdata/stock_plugin/manifest.toml +++ b/influxdata/stock_plugin/manifest.toml @@ -2,7 +2,7 @@ manifest_schema_version = "1.2" [plugin] name = "stock_plugin" -version = "0.2.0" +version = "0.3.0" description = "Tracks stock, ETF, and mutual fund portfolio values from Yahoo Finance with market-hours gating and rollups." triggers = ["process_scheduled_call"] homepage = "https://www.influxdata.com/" @@ -16,4 +16,8 @@ exclude = [ [dependencies] database_version = ">=3.0.0" -python = ["yfinance>=0.2.40", "pandas_market_calendars>=4.4"] +python = [ + "yfinance>=0.2.40", + "pandas_market_calendars>=4.4", + "influxdata-plugin-utils>=0.3.0", +] diff --git a/influxdata/stock_plugin/requirements-dev.txt b/influxdata/stock_plugin/requirements-dev.txt new file mode 100644 index 0000000..cd2aedf --- /dev/null +++ b/influxdata/stock_plugin/requirements-dev.txt @@ -0,0 +1,4 @@ +pytest +yfinance>=0.2.40 +pandas_market_calendars>=4.4 +influxdata-plugin-utils>=0.3.0 diff --git a/influxdata/stock_plugin/requirements.txt b/influxdata/stock_plugin/requirements.txt index ac12b67..3b29187 100644 --- a/influxdata/stock_plugin/requirements.txt +++ b/influxdata/stock_plugin/requirements.txt @@ -1,2 +1,3 @@ yfinance>=0.2.40 pandas_market_calendars>=4.4 +influxdata-plugin-utils>=0.3.0 diff --git a/influxdata/stock_plugin/stock_plugin.py b/influxdata/stock_plugin/stock_plugin.py index 1973456..3c633f4 100644 --- a/influxdata/stock_plugin/stock_plugin.py +++ b/influxdata/stock_plugin/stock_plugin.py @@ -23,7 +23,31 @@ { "name": "config_path", "example": "stock_plugin.toml", - "description": "Path to TOML configuration file. Supports absolute paths or relative paths resolved from INFLUXDB3_PLUGIN_DIR, PLUGIN_DIR, VIRTUAL_ENV parent, or the plugin directory.", + "description": "Path to TOML configuration file. Absolute paths are used as-is; relative paths resolve against the plugin directory (PLUGIN_DIR, INFLUXDB3_PLUGIN_DIR, or the VIRTUAL_ENV parent). Defaults to stock_plugin.toml, loaded when present.", + "required": false + }, + { + "name": "write_during_closed_hours", + "example": "true", + "description": "Write rows even when the configured market is closed. When false, stocks and ETFs are skipped outside the exchange session. Defaults to true.", + "required": false + }, + { + "name": "mutual_fund_check_time", + "example": "18:00", + "description": "Time of day in market_timezone, as HH:MM, after which mutual fund NAV is fetched. Defaults to 18:00.", + "required": false + }, + { + "name": "market_calendar", + "example": "NYSE", + "description": "Exchange calendar governing the market-open check. Any name accepted by pandas_market_calendars. Defaults to NYSE.", + "required": false + }, + { + "name": "market_timezone", + "example": "America/New_York", + "description": "IANA timezone of the exchange, used for mutual_fund_check_time and for the local market day. Defaults to America/New_York.", "required": false } ] @@ -32,8 +56,7 @@ from __future__ import annotations -import os -import tomllib +import math import uuid from dataclasses import dataclass from datetime import datetime, timezone @@ -41,6 +64,27 @@ from typing import Optional from zoneinfo import ZoneInfo +from influxdata_plugin_utils.config import Validator, load_plugin_config, resolve_path +from influxdata_plugin_utils.parsing import ( + parse_bool, + parse_delimited_list, + parse_timestamp_ns, +) +from influxdata_plugin_utils.write import build_line, write_data + + +# Default holdings +DEFAULT_PORTFOLIO = "AAPL:1|MSFT:1|GOOG:1" + +DEFAULT_DATABASE = "stocks" + +# TOML config loaded when `config_path` is not set and the file exists +DEFAULT_CONFIG_FILE = "stock_plugin.toml" + +# Holdings and categories are spelled differently per medium, so each spelling is read only from its own medium. +INLINE_ARG_KEYS = ("portfolio", "categories") +TOML_TABLE_KEYS = ("holdings", "portfolio_categories") + def _parse_hhmm(s: str) -> tuple[int, int]: """Parse a HH:MM time string. Raises ValueError on bad input.""" @@ -57,6 +101,80 @@ def _parse_hhmm(s: str) -> tuple[int, int]: return hour, minute +def _parse_quantity(raw, where: str) -> float: + """Parse a holding quantity, rejecting non-finite values. + + `inf` and `nan` are valid TOML floats and valid input to float(), but they + cannot be written as a field, so they are rejected as configuration errors. + """ + try: + quantity = float(raw) + except (TypeError, ValueError) as e: + raise ValueError(f"invalid quantity {raw!r} {where}") from e + if not math.isfinite(quantity): + raise ValueError(f"invalid quantity {raw!r} {where}: must be a finite number") + return quantity + + +def _validated_hhmm(raw) -> str: + """Config cast: accept a HH:MM string, rejecting anything unparseable.""" + text = str(raw).strip() + _parse_hhmm(text) + return text + + +def _validated_market_calendar(raw) -> str: + """Config cast: accept an exchange name known to pandas_market_calendars.""" + name = str(raw).strip() + if not name: + raise ValueError("market_calendar must be a non-empty exchange name") + try: + import pandas_market_calendars as mcal + except ImportError as e: + raise ValueError( + "pandas_market_calendars is required to validate market_calendar; " + "install it in the plugin environment" + ) from e + try: + mcal.get_calendar(name) + except Exception as e: + raise ValueError( + f"market_calendar {name!r} is not accepted by " + f"pandas_market_calendars: {e}" + ) from e + return name + + +def _validated_market_timezone(raw) -> str: + """Config cast: accept an IANA timezone name.""" + name = str(raw).strip() + if not name: + raise ValueError("market_timezone must be a non-empty IANA timezone") + try: + ZoneInfo(name) + except Exception as e: + raise ValueError( + f"market_timezone {name!r} is not a valid IANA timezone: {e}" + ) from e + return name + + +VALIDATORS: list = [ + Validator( + "database", default=DEFAULT_DATABASE, + cast=lambda raw: str(raw).strip() or DEFAULT_DATABASE, + ), + Validator("portfolio", default="", cast=str), + Validator("categories", default="", cast=str), + Validator("write_during_closed_hours", default=True, cast=parse_bool), + Validator("mutual_fund_check_time", default="18:00", cast=_validated_hhmm), + Validator("market_calendar", default="NYSE", cast=_validated_market_calendar), + Validator( + "market_timezone", default="America/New_York", cast=_validated_market_timezone + ), +] + + def _normalize_quote_type(qt: Optional[str]) -> str: """Map yfinance fast_info.quote_type to our normalized asset_type tag.""" if not qt: @@ -182,10 +300,6 @@ class ResolvedConfig: market_timezone: str = "America/New_York" -# Default holdings -DEFAULT_PORTFOLIO = "AAPL:1|MSFT:1|GOOG:1" - - def parse_inline_portfolio(value: str) -> dict[str, list[Holding]]: """Parse the inline `portfolio=` trigger argument. @@ -200,16 +314,11 @@ def parse_inline_portfolio(value: str) -> dict[str, list[Holding]]: Raises ValueError on any malformed input. """ - if not value or not value.strip(): + entries = parse_delimited_list(value, sep="|") + if not entries: raise ValueError("portfolio argument is empty") result: dict[str, list[Holding]] = {} - for raw in value.split("|"): - raw = raw.strip() - if not raw: - raise ValueError( - f"invalid portfolio argument {value!r}: empty token " - f"(check for leading, trailing, or doubled '|')" - ) + for raw in entries: parts = raw.split(":") if len(parts) not in (2, 3): raise ValueError( @@ -218,12 +327,7 @@ def parse_inline_portfolio(value: str) -> dict[str, list[Holding]]: symbol = parts[0].strip().upper() if not symbol: raise ValueError(f"invalid holding spec {raw!r}: empty symbol") - try: - quantity = float(parts[1]) - except ValueError as e: - raise ValueError( - f"invalid quantity {parts[1]!r} in {raw!r}" - ) from e + quantity = _parse_quantity(parts[1], f"in {raw!r}") if len(parts) == 3 and parts[2].strip(): portfolio = parts[2].strip() else: @@ -242,23 +346,17 @@ def parse_inline_categories(value: str) -> dict[str, str]: Raises ValueError on any malformed input. """ - if not value or not value.strip(): + entries = parse_delimited_list(value, sep="|") + if not entries: raise ValueError("categories argument is empty") result: dict[str, str] = {} - for raw in value.split("|"): - raw = raw.strip() - if not raw: - raise ValueError( - f"invalid categories argument {value!r}: empty token " - f"(check for leading, trailing, or doubled '|')" - ) + for raw in entries: parts = raw.split(":") if len(parts) != 2: raise ValueError( f"invalid category spec {raw!r}: expected PORTFOLIO:CATEGORY" ) - portfolio_name = parts[0].strip() - category_name = parts[1].strip() + portfolio_name, category_name = parts[0].strip(), parts[1].strip() if not portfolio_name or not category_name: raise ValueError( f"invalid category spec {raw!r}: empty portfolio or category name" @@ -294,37 +392,10 @@ def _aggregate_duplicate_holdings( return result -def _validate_market_calendar(calendar_name: str) -> None: - """Fail fast if pandas_market_calendars does not know this calendar.""" - try: - import pandas_market_calendars as mcal - except ImportError as e: - raise ValueError( - "pandas_market_calendars is required to validate market_calendar; " - "install it in the plugin environment" - ) from e - - try: - mcal.get_calendar(calendar_name) - except Exception as e: - raise ValueError( - f"market_calendar {calendar_name!r} is not accepted by " - f"pandas_market_calendars: {e}" - ) from e - - -def load_toml_config(path: Path) -> tuple[dict, dict[str, list[Holding]]]: - """Load TOML config from `path`. - - Returns (top_level_data_dict, holdings_by_portfolio). The full raw - TOML top-level dict is returned so resolve_config can pick out - optional scalar keys (database, market_calendar, etc). - - Expected TOML shape: - database = "stocks" - write_during_closed_hours = true - mutual_fund_check_time = "18:00" +def parse_toml_holdings(holdings_section) -> dict[str, list[Holding]]: + """Parse the `[holdings.]` tables of the TOML config. + Expected shape: [holdings.401k] AAPL = 10 MSFT = 5 @@ -332,19 +403,12 @@ def load_toml_config(path: Path) -> tuple[dict, dict[str, list[Holding]]]: [holdings.brokerage] GOOG = 2.5 - Holdings may be empty (no [holdings.*] sections); the caller decides - how to handle that (e.g. fall back to a default portfolio). + Holdings may be empty (no [holdings.*] sections); the caller decides how to + handle that (e.g. fall back to a default portfolio). - Raises: - FileNotFoundError: path does not exist - tomllib.TOMLDecodeError: file is not valid TOML - ValueError: holdings is not a table, or a holding has an invalid quantity + Raises ValueError if holdings is not a table, or a holding has an invalid + quantity. """ - if not path.exists(): - raise FileNotFoundError(f"TOML config not found: {path}") - with open(path, "rb") as f: - data = tomllib.load(f) - holdings_section = data.get("holdings", {}) if not isinstance(holdings_section, dict): raise ValueError( f"[holdings] must be a table of portfolios; " @@ -365,13 +429,10 @@ def load_toml_config(path: Path) -> tuple[dict, dict[str, list[Holding]]]: f"ticker symbol, quote it in TOML, for example " f'"{dotted_hint}" = 1' ) - try: - qty = float(quantity) - except (TypeError, ValueError) as e: - raise ValueError( - f"invalid quantity {quantity!r} for symbol {symbol!r} " - f"in [holdings.{portfolio_name}]" - ) from e + qty = _parse_quantity( + quantity, + f"for symbol {symbol!r} in [holdings.{portfolio_name}]", + ) result.setdefault(portfolio_name, []).append( Holding( symbol=symbol.upper(), @@ -379,74 +440,81 @@ def load_toml_config(path: Path) -> tuple[dict, dict[str, list[Holding]]]: portfolio=portfolio_name, ) ) - return data, result + return result -def resolve_config_path(path: str, default_toml_path: Path) -> Path: - """Resolve TOML config path using the plugin-dir fallbacks used by plugins. +def resolve_toml_path(local, args: dict[str, str], task_id: str) -> Optional[Path]: + """Locate the TOML config file, or return None when there is none to load. - Absolute paths are used as-is. Relative paths are resolved from - INFLUXDB3_PLUGIN_DIR or PLUGIN_DIR when available, then VIRTUAL_ENV's - parent directory, and finally the supplied default TOML directory. + An explicit `config_path` must exist; the default file is loaded only when + present, so the plugin runs without any TOML config. """ - raw_path = Path(path) - if raw_path.is_absolute(): - return raw_path - - candidates: list[Path] = [] - if influxdb3_plugin_dir := os.environ.get("INFLUXDB3_PLUGIN_DIR"): - candidates.append(Path(influxdb3_plugin_dir)) - if plugin_dir := os.environ.get("PLUGIN_DIR"): - candidates.append(Path(plugin_dir)) - if virtual_env := os.environ.get("VIRTUAL_ENV"): - candidates.append(Path(virtual_env).parent) - candidates.append(default_toml_path.parent) - - for base in candidates: - candidate = base / raw_path - if candidate.exists(): - return candidate - return candidates[0] / raw_path - - -def resolve_config( - args: dict[str, str], default_toml_path: Path -) -> ResolvedConfig: - """Resolve final config from trigger args + TOML. + explicit = args.get("config_path") + try: + path = resolve_path(explicit or DEFAULT_CONFIG_FILE) + except ValueError as e: + if explicit: + raise ValueError(f"cannot resolve config_path {explicit!r}: {e}") from e + local.warn( + f"[{task_id}] stock_plugin: cannot resolve the plugin directory ({e}), " + f"so any {DEFAULT_CONFIG_FILE} is ignored; continuing with trigger arguments" + ) + return None + if path.exists(): + return path + if explicit: + raise ValueError(f"config_path was set but no TOML config found at {path}.") + return None - Precedence: - Holdings: inline `portfolio=` arg > TOML [holdings.*] > DEFAULT_PORTFOLIO - Database: `database=` arg > TOML `database` > default "stocks" - Config path: `config_path=` arg > default_toml_path. Relative paths - resolve from INFLUXDB3_PLUGIN_DIR, PLUGIN_DIR, VIRTUAL_ENV's parent, - or default_toml_path.parent. - Other TOML scalars (write_during_closed_hours, mutual_fund_check_time, - market_calendar, market_timezone) come only from TOML — there is no - inline-arg override for them. Defaults apply when the key is absent. +def resolve_config(local, args: dict[str, str], task_id: str) -> ResolvedConfig: + """Resolve final config from trigger args + TOML. - When no inline holdings and no TOML file are found, falls back to - DEFAULT_PORTFOLIO so the plugin runs without configuration. + Trigger arguments override TOML keys of the same name. Holdings come from the + inline `portfolio=` argument, else from the TOML `[holdings.*]` tables, else + from DEFAULT_PORTFOLIO, so the plugin runs without configuration; categories + follow the same order with `categories=` and `[portfolio_categories]`. - Raises ValueError if config_path is set but the file is missing, or if - any TOML scalar has an invalid value. + Raises ValueError if config_path is set but the file is missing, or if any + configuration value is invalid. """ - inline_portfolio = args.get("portfolio") - explicit_config_path = bool(args.get("config_path")) - config_path = resolve_config_path( - args.get("config_path") or str(default_toml_path), - default_toml_path, - ) + args = {key: value for key, value in (args or {}).items() if value not in (None, "")} toml_data: dict = {} - holdings: dict[str, list[Holding]] = {} - if inline_portfolio: - holdings = parse_inline_portfolio(inline_portfolio) - elif config_path.exists(): - toml_data, holdings = load_toml_config(config_path) - elif explicit_config_path: - raise ValueError("config_path was set but no TOML config found.") + toml_path = resolve_toml_path(local, args, task_id) + if toml_path: + loaded = load_plugin_config( + {"config_path": str(toml_path)}, + config_file_path_arg="config_path", + source="toml", + ) + toml_data = {key.lower(): value for key, value in loaded.as_dict().items()} + holdings_table = toml_data.get("holdings") + categories_table = toml_data.get("portfolio_categories") + settings = load_plugin_config( + { + **{ + key: value + for key, value in toml_data.items() + if key not in TOML_TABLE_KEYS + INLINE_ARG_KEYS + }, + **{ + key: value + for key, value in args.items() + if key not in TOML_TABLE_KEYS + }, + }, + validators=VALIDATORS, + config_file_path_arg="config_path", + source="args", + ) + config = {key.lower(): value for key, value in settings.as_dict().items()} + + if config["portfolio"]: + holdings = parse_inline_portfolio(config["portfolio"]) + else: + holdings = parse_toml_holdings({} if holdings_table is None else holdings_table) if not holdings: holdings = parse_inline_portfolio(DEFAULT_PORTFOLIO) @@ -457,66 +525,25 @@ def resolve_config( ) holdings = _aggregate_duplicate_holdings(holdings) - if args.get("database"): - database = args["database"] - elif toml_data.get("database"): - database = toml_data["database"] - else: - database = "stocks" - - write_closed = toml_data.get("write_during_closed_hours", True) - if not isinstance(write_closed, bool): - raise ValueError( - f"write_during_closed_hours must be true or false; got {write_closed!r}" - ) - - mf_check = toml_data.get("mutual_fund_check_time", "18:00") - if not isinstance(mf_check, str): - raise ValueError( - f"mutual_fund_check_time must be a HH:MM string; got {mf_check!r}" - ) - _parse_hhmm(mf_check) # validate format, raises ValueError on bad input - - market_calendar = toml_data.get("market_calendar", "NYSE") - if not isinstance(market_calendar, str) or not market_calendar.strip(): - raise ValueError( - f"market_calendar must be a non-empty exchange name string; got {market_calendar!r}" - ) - market_calendar = market_calendar.strip() - _validate_market_calendar(market_calendar) - - market_timezone = toml_data.get("market_timezone", "America/New_York") - if not isinstance(market_timezone, str) or not market_timezone.strip(): - raise ValueError( - f"market_timezone must be a non-empty IANA timezone string; got {market_timezone!r}" - ) - try: - ZoneInfo(market_timezone) - except Exception as e: - raise ValueError( - f"market_timezone {market_timezone!r} is not a valid IANA timezone: {e}" - ) from e - - # Resolve categories: trigger arg > TOML > {} - inline_categories = args.get("categories") - if inline_categories: - categories = parse_inline_categories(inline_categories) + if config["categories"]: + categories = parse_inline_categories(config["categories"]) else: - toml_cats = toml_data.get("portfolio_categories", {}) - if not isinstance(toml_cats, dict): + toml_categories = {} if categories_table is None else categories_table + if not isinstance(toml_categories, dict): raise ValueError( - f"[portfolio_categories] must be a TOML table; got {type(toml_cats).__name__}" + f"[portfolio_categories] must be a TOML table; " + f"got {type(toml_categories).__name__}" ) - categories = {str(k): str(v) for k, v in toml_cats.items()} + categories = {str(k): str(v) for k, v in toml_categories.items()} return ResolvedConfig( - database=database, + database=config["database"], holdings_by_portfolio=holdings, categories=categories, - write_during_closed_hours=write_closed, - mutual_fund_check_time=mf_check, - market_calendar=market_calendar, - market_timezone=market_timezone, + write_during_closed_hours=config["write_during_closed_hours"], + mutual_fund_check_time=config["mutual_fund_check_time"], + market_calendar=config["market_calendar"], + market_timezone=config["market_timezone"], ) @@ -627,15 +654,26 @@ def compute_category_totals( return rows +def _finite(value) -> Optional[float]: + """Coerce to float, treating None and non-finite values (NaN, inf) as missing.""" + if value is None: + return None + try: + number = float(value) + except (TypeError, ValueError): + return None + return number if math.isfinite(number) else None + + def fetch_quote(symbol: str) -> Quote: """Fetch current price + day OHL + previous close + currency for `symbol`. Uses yfinance.Ticker(symbol).fast_info — a single HTTP call that returns all the fields the plugin needs. Optional fields default to None when - yfinance returns None for them; the caller decides how to handle missing - fields when building line protocol. + yfinance returns None or a non-finite number for them; the caller decides + how to handle missing fields when building line protocol. - Raises ValueError if no last_price is returned. Raises whatever + Raises ValueError if no usable last_price is returned. Raises whatever yfinance raises on network/other failures (e.g. requests exceptions). yfinance is imported lazily so this module is importable / AST-parseable @@ -644,12 +682,15 @@ def fetch_quote(symbol: str) -> Quote: import yfinance as yf ticker = yf.Ticker(symbol) fi = ticker.fast_info - if fi.last_price is None: - raise ValueError(f"no last_price returned for {symbol}") + price = _finite(fi.last_price) + if price is None: + raise ValueError(f"no usable last_price returned for {symbol}") def _opt(attr: str) -> Optional[float]: - v = getattr(fi, attr, None) - return float(v) if v is not None else None + try: + return _finite(getattr(fi, attr, None)) + except (KeyError, AttributeError): + return None try: raw_currency = fi.currency @@ -663,7 +704,7 @@ def _opt(attr: str) -> Optional[float]: return Quote( symbol=symbol, - price=float(fi.last_price), + price=price, currency=str(raw_currency) if raw_currency else "USD", asset_type=_normalize_quote_type(raw_quote_type), previous_close=_opt("previous_close"), @@ -678,7 +719,6 @@ def _main( args: dict[str, str], fetcher, line_builder_cls, - plugin_dir: Path, now_ns: int, task_id: str, ) -> None: @@ -697,10 +737,9 @@ def _main( Per-symbol asset_type is auto-detected via yfinance.fast_info.quote_type on first fetch and cached (no TTL) so we don't re-query for type. """ - default_toml = plugin_dir / "stock_plugin.toml" try: - config = resolve_config(args, default_toml) - except (ValueError, OSError, tomllib.TOMLDecodeError) as e: + config = resolve_config(local, args, task_id) + except Exception as e: local.error(f"[{task_id}] stock_plugin: configuration error: {e}") return @@ -861,52 +900,57 @@ def _record_skip(holding: Holding, portfolio: str, reason: str) -> None: ) category_totals = compute_category_totals(totals, now_ns) - for r in successful: - lb = line_builder_cls("stock_holdings") - lb.tag("symbol", r.symbol) - lb.tag("portfolio", r.portfolio) - lb.tag("asset_type", r.asset_type) - if r.category: - lb.tag("category", r.category) - lb.float64_field("price", r.price) - lb.float64_field("quantity", r.quantity) - lb.float64_field("value", r.value) - lb.string_field("currency", r.currency) - if r.previous_close is not None: - lb.float64_field("previous_close", r.previous_close) - if r.day_open is not None: - lb.float64_field("day_open", r.day_open) - if r.day_high is not None: - lb.float64_field("day_high", r.day_high) - if r.day_low is not None: - lb.float64_field("day_low", r.day_low) - lb.time_ns(r.timestamp_ns) - local.write_to_db(config.database, lb) - - for t in totals: - lb = line_builder_cls("portfolio_totals") - lb.tag("portfolio", t.portfolio) - if t.category: - lb.tag("category", t.category) - lb.float64_field("value", t.value) - lb.int64_field("symbol_count", t.symbol_count) - lb.int64_field("missing_symbols", t.missing_symbols) - lb.int64_field("skipped_symbols", t.skipped_symbols) - lb.int64_field("carried_symbols", t.carried_symbols) - lb.time_ns(t.timestamp_ns) - local.write_to_db(config.database, lb) - - for c in category_totals: - lb = line_builder_cls("category_totals") - lb.tag("category", c.category) - lb.float64_field("value", c.value) - lb.int64_field("symbol_count", c.symbol_count) - lb.int64_field("portfolio_count", c.portfolio_count) - lb.int64_field("missing_symbols", c.missing_symbols) - lb.int64_field("skipped_symbols", c.skipped_symbols) - lb.int64_field("carried_symbols", c.carried_symbols) - lb.time_ns(c.timestamp_ns) - local.write_to_db(config.database, lb) + # Field types are inferred from the row values: prices are floats, counts ints. + # A None tag or field is omitted, as is an empty category. + lines: list = [] + try: + lines += [ + build_line( + line_builder_cls, + "stock_holdings", + tags={"symbol": r.symbol, "portfolio": r.portfolio, + "asset_type": r.asset_type, "category": r.category or None}, + fields={"price": r.price, "quantity": r.quantity, "value": r.value, + "currency": r.currency, "previous_close": r.previous_close, + "day_open": r.day_open, "day_high": r.day_high, + "day_low": r.day_low}, + time_ns=r.timestamp_ns, + ) + for r in successful + ] + lines += [ + build_line( + line_builder_cls, + "portfolio_totals", + tags={"portfolio": t.portfolio, "category": t.category or None}, + fields={"value": t.value, "symbol_count": t.symbol_count, + "missing_symbols": t.missing_symbols, + "skipped_symbols": t.skipped_symbols, + "carried_symbols": t.carried_symbols}, + time_ns=t.timestamp_ns, + ) + for t in totals + ] + lines += [ + build_line( + line_builder_cls, + "category_totals", + tags={"category": c.category}, + fields={"value": c.value, "symbol_count": c.symbol_count, + "portfolio_count": c.portfolio_count, + "missing_symbols": c.missing_symbols, + "skipped_symbols": c.skipped_symbols, + "carried_symbols": c.carried_symbols}, + time_ns=c.timestamp_ns, + ) + for c in category_totals + ] + write_data(local, lines, batch=True, retries=0, database=config.database) + except Exception as e: + local.error( + f"[{task_id}] stock_plugin: failed to write to {config.database}: {e}" + ) + return skipped_total = sum(skipped_by_portfolio.values()) parts = [f"fetched {len(successful)}/{total_symbols} symbols"] @@ -955,20 +999,14 @@ def process_scheduled_call(influxdb3_local, call_time, args): is unambiguous regardless of host timezone; the few-seconds offset from the scheduled tick boundary is negligible for a 15-minute polling cadence. - - The InfluxDB runtime executes plugins via exec(), so __file__ is not - defined. INFLUXDB3_PLUGIN_DIR is the env var the server itself uses - to locate the plugin directory. """ - plugin_dir = Path(os.environ.get("INFLUXDB3_PLUGIN_DIR", ".")) - now_ns = int(datetime.now(timezone.utc).timestamp() * 1_000_000_000) + now_ns = parse_timestamp_ns(datetime.now(timezone.utc), "datetime") task_id = uuid.uuid4().hex[:8] _main( local=influxdb3_local, args=args or {}, fetcher=fetch_quote, line_builder_cls=LineBuilder, # runtime-injected global - plugin_dir=plugin_dir, now_ns=now_ns, task_id=task_id, ) diff --git a/influxdata/stock_plugin/stock_plugin.toml.example b/influxdata/stock_plugin/stock_plugin.toml.example index 70c9aaa..aedd063 100644 --- a/influxdata/stock_plugin/stock_plugin.toml.example +++ b/influxdata/stock_plugin/stock_plugin.toml.example @@ -1,8 +1,10 @@ # Example stock_plugin configuration. # Copy to stock_plugin.toml in your InfluxDB plugin directory. +# The scalar keys below can also be passed as trigger arguments, which win over +# the values set here. The [holdings.*] and [portfolio_categories] tables have no +# trigger-argument spelling: pass portfolio=... and categories=... instead. # Target database for writes. Optional — defaults to "stocks". -# Overridden by --trigger-arguments database=... if set. database = "stocks" # Write rows even when US markets are closed? When false, stocks/ETFs diff --git a/influxdata/stock_plugin/test_stock_plugin.py b/influxdata/stock_plugin/test_stock_plugin.py new file mode 100644 index 0000000..f02e0cd --- /dev/null +++ b/influxdata/stock_plugin/test_stock_plugin.py @@ -0,0 +1,747 @@ +"""Unit and integration tests for the stock_plugin plugin.""" + +import json +import os +import sys +from collections import namedtuple +from datetime import datetime +from textwrap import dedent + +import pytest +from influxdata_plugin_utils import write as utils_write + +sys.path.insert(0, os.path.dirname(__file__)) +import stock_plugin as sp + +# Captured before the fixtures patch the module attribute. +REAL_IS_MARKET_OPEN = sp._is_market_open + + +def ns(iso: str) -> int: + return int(datetime.fromisoformat(iso).timestamp() * 1_000_000_000) + + +# 2023-11-15 is a Wednesday; the market timezone default is America/New_York. +BEFORE_NAV = ns("2023-11-15T15:00:00-05:00") +AFTER_NAV = ns("2023-11-15T19:00:00-05:00") +TODAY_LOCAL = "2023-11-15" + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class FakeCache: + def __init__(self, initial=None): + self.store = dict(initial or {}) + + def get(self, key, default=None, use_global=None): + return self.store.get(key, default) + + def put(self, key, value, ttl=None, use_global=None): + self.store[key] = value + + +class FakeLineBuilder: + def __init__(self, measurement): + self.measurement = measurement + self.tags = {} + self.fields = {} + self.timestamp = None + + def tag(self, key, value): + self.tags[key] = value + return self + + def int64_field(self, key, value): + self.fields[key] = f"{value}i" + return self + + def uint64_field(self, key, value): + self.fields[key] = f"{value}u" + return self + + def float64_field(self, key, value): + self.fields[key] = repr(float(value)) + return self + + def bool_field(self, key, value): + self.fields[key] = "true" if value else "false" + return self + + def string_field(self, key, value): + self.fields[key] = f'"{value}"' + return self + + def time_ns(self, timestamp_ns): + self.timestamp = timestamp_ns + return self + + def build(self): + line = self.measurement + if self.tags: + line += "," + ",".join(f"{k}={v}" for k, v in self.tags.items()) + line += " " + ",".join(f"{k}={v}" for k, v in self.fields.items()) + return f"{line} {self.timestamp}" + + +Record = namedtuple("Record", ["measurement", "tags", "fields", "timestamp"]) + + +def _parse_field(raw): + if raw.startswith('"'): + return raw[1:-1] + if raw[-1] in ("i", "u"): + return int(raw[:-1]) + return float(raw) + + +def _parse_lp(line): + head, fields_str, timestamp = line.rsplit(" ", 2) + parts = head.split(",") + tags = dict(kv.split("=", 1) for kv in parts[1:]) + fields = { + key: _parse_field(value) + for key, value in (kv.split("=", 1) for kv in fields_str.split(",")) + } + return Record(parts[0], tags, fields, int(timestamp)) + + +class FakeLocal: + def __init__(self, cache=None, write_error=None): + self.cache = FakeCache(cache) + self.write_error = write_error + self.writes = [] # one (database, [Record]) per write call + self.infos = [] + self.warns = [] + self.errors = [] + + def write_to_db(self, database, batch): + if self.write_error is not None: + raise self.write_error + lines = [_parse_lp(lp) for lp in batch.build().split("\n")] + self.writes.append((database, lines)) + + def write(self, batch): + raise AssertionError("writes must target an explicit database") + + def info(self, message): + self.infos.append(message) + + def warn(self, message): + self.warns.append(message) + + def error(self, message): + self.errors.append(message) + + +@pytest.fixture(autouse=True) +def plugin_dir(tmp_path, monkeypatch): + """Point config resolution at an empty directory, never the real plugin dir.""" + monkeypatch.setenv("PLUGIN_DIR", str(tmp_path)) + monkeypatch.delenv("INFLUXDB3_PLUGIN_DIR", raising=False) + monkeypatch.setattr(sp, "LineBuilder", FakeLineBuilder, raising=False) + monkeypatch.setattr(utils_write.time, "sleep", lambda _: None) + return tmp_path + + +@pytest.fixture(autouse=True) +def market(monkeypatch): + """Market state, so gating tests do not depend on the real calendar or clock.""" + state = {"open": True} + monkeypatch.setattr(sp, "_is_market_open", lambda *_: state["open"]) + return state + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def write_toml(plugin_dir, body, name=sp.DEFAULT_CONFIG_FILE): + path = plugin_dir / name + path.write_text(dedent(body)) + return path + + +def quote_fetcher(prices=None, asset_types=None, failures=(), currency="USD"): + """A fetcher returning deterministic quotes and raising for `failures`.""" + + def fetch(symbol): + if symbol in failures: + raise RuntimeError(f"no data for {symbol}") + return sp.Quote( + symbol=symbol, + price=(prices or {}).get(symbol, 100.0), + currency=currency, + asset_type=(asset_types or {}).get(symbol, "equity"), + previous_close=99.0, + day_open=None, + day_high=101.0, + day_low=None, + ) + + return fetch + + +def run(args=None, cache=None, now_ns=AFTER_NAV, fetcher=None, local=None): + local = local or FakeLocal(cache) + sp._main( + local=local, + args=args or {}, + fetcher=fetcher or quote_fetcher(), + line_builder_cls=FakeLineBuilder, + now_ns=now_ns, + task_id="test", + ) + return local + + +def records(local, measurement=None): + rows = [row for _, lines in local.writes for row in lines] + if measurement: + rows = [row for row in rows if row.measurement == measurement] + return rows + + +def resolve_config(args=None, local=None): + return sp.resolve_config(local or FakeLocal(), args or {}, "test") + + +def holdings_of(config): + return { + portfolio: [(h.symbol, h.quantity) for h in holdings] + for portfolio, holdings in config.holdings_by_portfolio.items() + } + + +# --------------------------------------------------------------------------- +# M1 — plugin metadata +# --------------------------------------------------------------------------- + + +def test_docstring_metadata_covers_every_supported_argument(): + header = json.loads(sp.__doc__) + assert header["plugin_type"] == ["scheduled"] + names = [arg["name"] for arg in header["scheduled_args_config"]] + validated = {validator.names[0] for validator in sp.VALIDATORS} + # every validated key is documented, plus the TOML path itself + assert set(names) == validated | {"config_path"} + for entry in header["scheduled_args_config"]: + assert set(entry) == {"name", "example", "description", "required"} + + +# --------------------------------------------------------------------------- +# M2 — configuration resolution +# --------------------------------------------------------------------------- + + +def test_defaults_apply_without_any_configuration(): + config = resolve_config({}) + assert config.database == "stocks" + assert config.write_during_closed_hours is True + assert config.mutual_fund_check_time == "18:00" + assert config.market_calendar == "NYSE" + assert config.market_timezone == "America/New_York" + assert holdings_of(config) == {"main": [("AAPL", 1.0), ("MSFT", 1.0), ("GOOG", 1.0)]} + assert config.categories == {} + + +def test_toml_supplies_holdings_categories_and_scalars(plugin_dir): + write_toml( + plugin_dir, + """ + database = "portfolio" + write_during_closed_hours = false + mutual_fund_check_time = "20:30" + market_calendar = "LSE" + market_timezone = "Europe/London" + + [portfolio_categories] + "401k" = "Retirement" + + [holdings.401k] + AAPL = 10 + "VOD.L" = 2.5 + """, + ) + config = resolve_config({}) + assert config.database == "portfolio" + assert config.write_during_closed_hours is False + assert config.mutual_fund_check_time == "20:30" + assert config.market_calendar == "LSE" + assert config.market_timezone == "Europe/London" + assert holdings_of(config) == {"401k": [("AAPL", 10.0), ("VOD.L", 2.5)]} + assert config.categories == {"401k": "Retirement"} + + +def test_trigger_arguments_override_toml_keys(plugin_dir): + write_toml( + plugin_dir, + """ + database = "from_toml" + write_during_closed_hours = false + market_calendar = "LSE" + + [holdings.401k] + AAPL = 10 + """, + ) + config = resolve_config( + { + "database": "from_args", + "write_during_closed_hours": "true", + "market_calendar": "NYSE", + } + ) + assert config.database == "from_args" + assert config.write_during_closed_hours is True + assert config.market_calendar == "NYSE" + # holdings still come from the file + assert holdings_of(config) == {"401k": [("AAPL", 10.0)]} + + +def test_inline_holdings_replace_the_toml_tables_but_keep_toml_scalars(plugin_dir): + write_toml( + plugin_dir, + """ + database = "portfolio" + + [portfolio_categories] + "401k" = "Retirement" + + [holdings.brokerage] + GOOG = 5 + """, + ) + config = resolve_config({"portfolio": "AAPL:2:401k"}) + assert holdings_of(config) == {"401k": [("AAPL", 2.0)]} + assert config.database == "portfolio" + assert config.categories == {"401k": "Retirement"} + + +def test_explicit_config_path_must_exist(plugin_dir): + write_toml(plugin_dir, '[holdings.main]\nAAPL = 1\n', name="custom.toml") + assert holdings_of(resolve_config({"config_path": "custom.toml"})) == { + "main": [("AAPL", 1.0)] + } + with pytest.raises(ValueError, match="no TOML config found"): + resolve_config({"config_path": "absent.toml"}) + + +@pytest.mark.parametrize( + "args, fragment", + [ + ({"mutual_fund_check_time": "25:00"}, "out of range"), + ({"mutual_fund_check_time": "noon"}, "expected HH:MM"), + ({"write_during_closed_hours": "maybe"}, "Invalid boolean"), + ({"market_timezone": "Mars/Olympus"}, "not a valid IANA timezone"), + ({"market_calendar": "NOPE"}, "not accepted by pandas_market_calendars"), + ({"portfolio": "AAPL:1:_total"}, "'_total' is reserved"), + ], +) +def test_invalid_configuration_is_rejected(args, fragment): + with pytest.raises(ValueError, match=fragment): + resolve_config(args) + + +def test_duplicate_symbols_in_one_portfolio_are_aggregated(): + config = resolve_config({"portfolio": "AAPL:2:401k|MSFT:1:401k|AAPL:3:401k"}) + assert holdings_of(config) == {"401k": [("AAPL", 5.0), ("MSFT", 1.0)]} + + +# --------------------------------------------------------------------------- +# M3 — inline argument parsing +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "value, expected", + [ + ("AAPL:10", {"main": [("AAPL", 10.0)]}), + ("AAPL:10|MSFT:5", {"main": [("AAPL", 10.0), ("MSFT", 5.0)]}), + ("aapl:2.5:401k", {"401k": [("AAPL", 2.5)]}), + ("AAPL:1||MSFT:2", {"main": [("AAPL", 1.0), ("MSFT", 2.0)]}), + (" AAPL:1 | MSFT:2 ", {"main": [("AAPL", 1.0), ("MSFT", 2.0)]}), + ("AAPL:-1", {"main": [("AAPL", -1.0)]}), + ], +) +def test_parse_inline_portfolio(value, expected): + parsed = sp.parse_inline_portfolio(value) + assert {p: [(h.symbol, h.quantity) for h in hs] for p, hs in parsed.items()} == expected + + +@pytest.mark.parametrize( + "value, fragment", + [ + ("", "is empty"), + ("AAPL", "expected SYMBOL:QUANTITY"), + ("AAPL:1:401k:extra", "expected SYMBOL:QUANTITY"), + (":1", "empty symbol"), + ("AAPL:many", "invalid quantity"), + ("AAPL:inf", "must be a finite number"), + ("AAPL:nan", "must be a finite number"), + ("AAPL:1e400", "must be a finite number"), + ], +) +def test_parse_inline_portfolio_rejects_malformed_input(value, fragment): + with pytest.raises(ValueError, match=fragment): + sp.parse_inline_portfolio(value) + + +@pytest.mark.parametrize( + "value, expected", + [ + ("401k:Retirement", {"401k": "Retirement"}), + ("401k:Retirement|brokerage:Investment", + {"401k": "Retirement", "brokerage": "Investment"}), + ("401k:Retirement||brokerage:Investment", + {"401k": "Retirement", "brokerage": "Investment"}), + ], +) +def test_parse_inline_categories(value, expected): + assert sp.parse_inline_categories(value) == expected + + +@pytest.mark.parametrize( + "value, fragment", + [ + ("|", "is empty"), + ("401k", "expected PORTFOLIO:CATEGORY"), + ("401k:", "empty portfolio or category name"), + (":Retirement", "empty portfolio or category name"), + ("401k:Retirement:extra", "expected PORTFOLIO:CATEGORY"), + # the message quotes the user's own token, not the parsed halves + ("401k : Retirement : extra", r"'401k : Retirement : extra'"), + ], +) +def test_parse_inline_categories_rejects_malformed_input(value, fragment): + with pytest.raises(ValueError, match=fragment): + sp.parse_inline_categories(value) + + +# --------------------------------------------------------------------------- +# M4 — TOML holdings parsing +# --------------------------------------------------------------------------- + + +def test_unquoted_dotted_symbol_reports_how_to_fix_it(plugin_dir): + write_toml(plugin_dir, "[holdings.main]\nVOD.L = 10\n") + with pytest.raises(ValueError, match=r'"VOD.L" = 1'): + resolve_config({}) + + +@pytest.mark.parametrize( + "body, fragment", + [ + ('holdings = "AAPL"', "must be a table of portfolios"), + ("[holdings.main]\nAAPL = nan\n", "must be a finite number"), + ('[holdings.main]\nAAPL = "ten"\n', "invalid quantity"), + ], +) +def test_invalid_toml_holdings_are_rejected(plugin_dir, body, fragment): + write_toml(plugin_dir, body) + with pytest.raises(ValueError, match=fragment): + resolve_config({}) + + +# --------------------------------------------------------------------------- +# M5 — market and mutual-fund gating +# --------------------------------------------------------------------------- + +WARM_EQUITY = {"asset_type:AAPL": "equity", "last_price:AAPL": 90.0} +WARM_FUND = {"asset_type:VFIAX": "mutualfund", "last_price:VFIAX": 40.0} + + +def test_closed_market_skips_equities_when_closed_hour_writes_are_disabled(market): + market["open"] = False + local = run( + args={"portfolio": "AAPL:2", "write_during_closed_hours": "false"}, + cache=WARM_EQUITY, + ) + assert records(local, "stock_holdings") == [] + total = records(local, "portfolio_totals")[0] + # value carried forward from the cached last price + assert total.fields["value"] == 180.0 + assert total.fields["skipped_symbols"] == 1 + assert total.fields["carried_symbols"] == 1 + assert total.fields["missing_symbols"] == 0 + + +def test_closed_market_still_fetches_when_closed_hour_writes_are_enabled(market): + market["open"] = False + local = run(args={"portfolio": "AAPL:2"}, cache=WARM_EQUITY) + assert [r.tags["symbol"] for r in records(local, "stock_holdings")] == ["AAPL"] + assert records(local, "portfolio_totals")[0].fields["skipped_symbols"] == 0 + + +@pytest.mark.parametrize( + "cache, now_ns, expected_skip", + [ + ({**WARM_FUND, "last_mf_date:VFIAX": TODAY_LOCAL}, AFTER_NAV, "already-today"), + (WARM_FUND, BEFORE_NAV, "too-early"), + ], +) +def test_mutual_fund_is_fetched_once_a_day_after_the_check_time( + cache, now_ns, expected_skip +): + local = run(args={"portfolio": "VFIAX:3"}, cache=cache, now_ns=now_ns) + assert records(local, "stock_holdings") == [] + assert records(local, "portfolio_totals")[0].fields["carried_symbols"] == 1 + assert expected_skip in local.infos[-1] + + +def test_mutual_fund_is_fetched_once_the_check_time_has_passed(): + local = run( + args={"portfolio": "VFIAX:3"}, + cache={**WARM_FUND, "last_mf_date:VFIAX": "2023-11-14"}, + fetcher=quote_fetcher(asset_types={"VFIAX": "mutualfund"}), + ) + assert [r.tags["symbol"] for r in records(local, "stock_holdings")] == ["VFIAX"] + assert local.cache.get("last_mf_date:VFIAX") == TODAY_LOCAL + + +def test_a_symbol_that_would_be_skipped_is_fetched_while_its_price_is_uncached(market): + market["open"] = False + local = run( + args={"portfolio": "AAPL:2", "write_during_closed_hours": "false"}, + cache={"asset_type:AAPL": "equity"}, + ) + assert [r.tags["symbol"] for r in records(local, "stock_holdings")] == ["AAPL"] + assert "cold-cache bootstrap" in local.infos[-1] + + +def test_asset_type_and_last_price_are_cached_for_later_runs(): + local = run(args={"portfolio": "AAPL:2"}, fetcher=quote_fetcher(prices={"AAPL": 12.5})) + assert local.cache.get("asset_type:AAPL") == "equity" + assert local.cache.get("last_price:AAPL") == 12.5 + assert local.cache.get("last_mf_date:AAPL") is None + + +@pytest.mark.parametrize( + "moment, expected", + [ + ("2023-11-15T10:00:00-05:00", True), # Wednesday, mid-session + ("2023-11-15T20:00:00-05:00", False), # Wednesday, after the close + ("2023-11-11T10:00:00-05:00", False), # Saturday + ("2023-11-23T10:00:00-05:00", False), # Thanksgiving + ], +) +def test_real_nyse_calendar_decides_whether_the_session_is_open(moment, expected): + from zoneinfo import ZoneInfo + + now_utc = datetime.fromisoformat(moment).astimezone(ZoneInfo("UTC")) + tz = ZoneInfo("America/New_York") + assert REAL_IS_MARKET_OPEN(now_utc, "NYSE", tz) is expected + + +# --------------------------------------------------------------------------- +# M6 — totals and category roll-ups +# --------------------------------------------------------------------------- + + +def test_totals_roll_up_per_portfolio_then_into_a_grand_total(): + local = run( + args={ + "portfolio": "AAPL:2:401k|MSFT:1:401k|GOOG:1:brokerage", + "categories": "401k:Retirement", + }, + fetcher=quote_fetcher(prices={"AAPL": 10.0, "MSFT": 20.0, "GOOG": 30.0}), + ) + totals = {r.tags["portfolio"]: r for r in records(local, "portfolio_totals")} + assert totals["401k"].fields["value"] == 40.0 + assert totals["401k"].tags["category"] == "Retirement" + assert totals["brokerage"].fields["value"] == 30.0 + assert "category" not in totals["brokerage"].tags + assert totals["_total"].fields["value"] == 70.0 + assert totals["_total"].fields["symbol_count"] == 3 + assert "category" not in totals["_total"].tags + + +def test_category_totals_exclude_uncategorized_portfolios_and_the_grand_total(): + local = run( + args={ + "portfolio": "AAPL:1:401k|MSFT:1:ira|GOOG:1:brokerage", + "categories": "401k:Retirement|ira:Retirement", + }, + fetcher=quote_fetcher(prices={"AAPL": 10.0, "MSFT": 20.0, "GOOG": 30.0}), + ) + rows = records(local, "category_totals") + assert [r.tags["category"] for r in rows] == ["Retirement"] + assert rows[0].fields["value"] == 30.0 + assert rows[0].fields["portfolio_count"] == 2 + assert rows[0].fields["symbol_count"] == 2 + + +# --------------------------------------------------------------------------- +# M7 — line protocol output +# --------------------------------------------------------------------------- + + +def test_a_run_emits_one_batched_write_with_a_single_timestamp(): + local = run(args={"portfolio": "AAPL:1:401k", "categories": "401k:Retirement", + "database": "portfolio"}) + assert len(local.writes) == 1 + database, lines = local.writes[0] + assert database == "portfolio" + assert [r.measurement for r in lines] == [ + "stock_holdings", + "portfolio_totals", + "portfolio_totals", + "category_totals", + ] + assert {r.timestamp for r in lines} == {AFTER_NAV} + + +def test_holding_row_carries_the_quote_and_omits_unavailable_fields(): + local = run(args={"portfolio": "AAPL:2"}, fetcher=quote_fetcher(prices={"AAPL": 10.0})) + row = records(local, "stock_holdings")[0] + assert row.tags == {"symbol": "AAPL", "portfolio": "main", "asset_type": "equity"} + assert row.fields == { + "price": 10.0, + "quantity": 2.0, + "value": 20.0, + "currency": "USD", + "previous_close": 99.0, + "day_high": 101.0, + } + + +def test_counts_are_written_as_integer_fields(): + local = run(args={"portfolio": "AAPL:1"}) + total = records(local, "portfolio_totals")[0] + for name in ("symbol_count", "missing_symbols", "skipped_symbols", "carried_symbols"): + assert isinstance(total.fields[name], int) + assert isinstance(total.fields["value"], float) + + +# --------------------------------------------------------------------------- +# M8 — failure handling +# --------------------------------------------------------------------------- + + +def test_a_fetch_failure_is_reported_and_counted_as_missing(): + local = run( + args={"portfolio": "AAPL:1|BOOM:1"}, fetcher=quote_fetcher(failures={"BOOM"}) + ) + assert [r.tags["symbol"] for r in records(local, "stock_holdings")] == ["AAPL"] + total = records(local, "portfolio_totals")[0] + assert total.fields["missing_symbols"] == 1 + assert total.fields["carried_symbols"] == 0 + assert any("failed to fetch BOOM" in w for w in local.warns) + assert "Failed: BOOM" in local.infos[-1] + + +def test_calendar_failure_aborts_only_when_closed_hour_writes_are_disabled(monkeypatch): + def boom(*_): + raise RuntimeError("calendar unavailable") + + monkeypatch.setattr(sp, "_is_market_open", boom) + + aborted = run(args={"portfolio": "AAPL:1", "write_during_closed_hours": "false"}) + assert aborted.writes == [] + assert any("skipping run" in e for e in aborted.errors) + + continued = run(args={"portfolio": "AAPL:1"}) + assert len(continued.writes) == 1 + assert any("continuing because" in w for w in continued.warns) + + +def test_a_configuration_error_is_logged_and_nothing_is_written(): + local = run(args={"portfolio": "AAPL:inf"}) + assert local.writes == [] + assert len(local.errors) == 1 + assert "configuration error" in local.errors[0] + assert "must be a finite number" in local.errors[0] + + +def test_a_row_that_cannot_be_built_is_logged_instead_of_raising(): + # quantity and price are both finite, but their product overflows to inf + local = run( + args={"portfolio": "AAPL:1e200"}, fetcher=quote_fetcher(prices={"AAPL": 1e200}) + ) + assert local.writes == [] + assert any("failed to write" in e and "not finite" in e for e in local.errors) + + +def test_a_write_failure_is_logged_instead_of_raising(): + local = FakeLocal(write_error=RuntimeError("database gone")) + run(args={"portfolio": "AAPL:1"}, local=local) + assert local.writes == [] + assert any("failed to write" in e and "database gone" in e for e in local.errors) + + +def test_process_scheduled_call_uses_the_runtime_line_builder(monkeypatch): + captured = {} + + def fake_main(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(sp, "_main", fake_main) + sp.process_scheduled_call(FakeLocal(), "ignored call_time", {"portfolio": "AAPL:1"}) + + assert captured["line_builder_cls"] is FakeLineBuilder + assert captured["fetcher"] is sp.fetch_quote + assert captured["args"] == {"portfolio": "AAPL:1"} + assert captured["now_ns"] > 0 + assert len(captured["task_id"]) == 8 + + +# --------------------------------------------------------------------------- +# M9 — quote fetching +# --------------------------------------------------------------------------- + + +class FakeFastInfo: + """Stand-in for yfinance fast_info, which raises KeyError for keys it cannot fill.""" + + def __init__(self, values): + self._values = values + + def __getattr__(self, name): + if name not in self._values: + raise KeyError(name) + return self._values[name] + + +@pytest.fixture +def fake_yfinance(monkeypatch): + def install(values): + module = type(sys)("yfinance") + module.Ticker = lambda symbol: type( + "Ticker", (), {"fast_info": FakeFastInfo(values)} + )() + monkeypatch.setitem(sys.modules, "yfinance", module) + + return install + + +def test_a_quote_survives_optional_fields_the_api_cannot_fill(fake_yfinance): + fake_yfinance({"last_price": 10.0, "currency": "GBP", "quote_type": "ETF"}) + quote = sp.fetch_quote("VOD.L") + assert (quote.price, quote.currency, quote.asset_type) == (10.0, "GBP", "etf") + assert (quote.previous_close, quote.day_open, quote.day_high, quote.day_low) == ( + None, + None, + None, + None, + ) + + +@pytest.mark.parametrize("last_price", [None, float("nan"), float("inf")]) +def test_an_unusable_last_price_fails_the_fetch(fake_yfinance, last_price): + fake_yfinance({"last_price": last_price}) + with pytest.raises(ValueError, match="no usable last_price"): + sp.fetch_quote("AAPL") + + +def test_non_finite_optional_fields_are_dropped(fake_yfinance): + fake_yfinance( + {"last_price": 10.0, "previous_close": float("nan"), "day_high": 11.0} + ) + quote = sp.fetch_quote("AAPL") + assert quote.previous_close is None + assert quote.day_high == 11.0 + # no currency or quote_type in the payload + assert (quote.currency, quote.asset_type) == ("USD", "other") \ No newline at end of file