From 5ab1d790d8ed7e2d06a9461e683ceaea793e909a Mon Sep 17 00:00:00 2001 From: juliet Date: Fri, 10 Jul 2026 14:45:50 -0400 Subject: [PATCH] Add osprey-stress CLI (closes #324) (#367) Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: Cassidy James Blaede --- CHANGELOG.md | 2 + .../event_stream/EventStream.module.css | 28 ++ osprey_worker/pyproject.toml | 1 + osprey_worker/src/osprey/worker/stress/cli.py | 452 ++++++++++++++++++ .../osprey/worker/stress/tests/test_cli.py | 270 +++++++++++ 5 files changed, 753 insertions(+) create mode 100644 osprey_worker/src/osprey/worker/stress/cli.py create mode 100644 osprey_worker/src/osprey/worker/stress/tests/test_cli.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0715b42..3dc19cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ For more information about each release including git tags and artifacts, see [R - Per-action health metrics in the executor ([#191](https://github.com/roostorg/osprey/pull/191) by [@cmttt](https://github.com/cmttt)) - Option to suppress cached errors to reduce metric bloat ([#180](https://github.com/roostorg/osprey/pull/180) by [@lithium-powered](https://github.com/lithium-powered)) - Experimental asyncio-native worker with metrics and engine/coordinator improvements ([#341](https://github.com/roostorg/osprey/pull/341) by [@cmttt](https://github.com/cmttt)) +- `osprey-stress` CLI: closed-loop stress harness that produces synthetic events at a configurable rate, observes their `ExecutionResult`s on the output topic, and reports drop rate and p50/p95/p99 latency, exiting non-zero on threshold breach so it can gate CI on pipeline health ([#367](https://github.com/roostorg/osprey/pull/367) by [@julietshen](https://github.com/julietshen), closes [#324](https://github.com/roostorg/osprey/issues/324)) ### Changed @@ -38,6 +39,7 @@ For more information about each release including git tags and artifacts, see [R - Tolerate malformed URI escapes in `EntityWithPopover` UI component ([#377](https://github.com/roostorg/osprey/pull/377) by [@julietshen](https://github.com/julietshen)) - Add retention limits to Kafka topics to prevent unbounded disk growth ([#249](https://github.com/roostorg/osprey/pull/249) by [@VINODvoid](https://github.com/VINODvoid)) - Fix failed UDF query ([#233](https://github.com/roostorg/osprey/pull/233) by [@chimosky](https://github.com/chimosky)) +- Force a visible scrollbar on the event-stream virtualized list, which was relying on the OS auto-hide default and was easy to miss ([#367](https://github.com/roostorg/osprey/pull/367) by [@julietshen](https://github.com/julietshen)) ## [1.0.1] - 2026-02-27 diff --git a/osprey_ui/src/components/event_stream/EventStream.module.css b/osprey_ui/src/components/event_stream/EventStream.module.css index 9ec36af..14a37f0 100644 --- a/osprey_ui/src/components/event_stream/EventStream.module.css +++ b/osprey_ui/src/components/event_stream/EventStream.module.css @@ -18,6 +18,34 @@ outline: none; } +/* react-virtualized renders a scrolling div with class `ReactVirtualized__List`. + On macOS the default scrollbar auto-hides until you're actively scrolling, + which makes it hard to skim long result sets — flagged in PR #367. Force a + thin always-visible bar in the theme colors. */ +.virtualizedList :global(.ReactVirtualized__List) { + scrollbar-width: thin; + scrollbar-color: var(--text-light-secondary) var(--background-secondary); +} + +.virtualizedList :global(.ReactVirtualized__List)::-webkit-scrollbar { + width: 10px; + height: 10px; +} + +.virtualizedList :global(.ReactVirtualized__List)::-webkit-scrollbar-track { + background: var(--background-secondary); +} + +.virtualizedList :global(.ReactVirtualized__List)::-webkit-scrollbar-thumb { + background-color: var(--text-light-secondary); + border-radius: 5px; + border: 2px solid var(--background-secondary); +} + +.virtualizedList :global(.ReactVirtualized__List)::-webkit-scrollbar-thumb:hover { + background-color: var(--text-light-primary); +} + .iconText { margin-right: 8px; } diff --git a/osprey_worker/pyproject.toml b/osprey_worker/pyproject.toml index 588311a..5bfc135 100644 --- a/osprey_worker/pyproject.toml +++ b/osprey_worker/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ ] [project.scripts] osprey-cli = "osprey.worker.lib.cli:cli" +osprey-stress = "osprey.worker.stress.cli:main" [project.entry-points."osprey_plugin"] stdlib = "osprey.worker._stdlibplugin" diff --git a/osprey_worker/src/osprey/worker/stress/cli.py b/osprey_worker/src/osprey/worker/stress/cli.py new file mode 100644 index 0000000..146f1c1 --- /dev/null +++ b/osprey_worker/src/osprey/worker/stress/cli.py @@ -0,0 +1,452 @@ +"""Stress harness CLI entry point. + +Usage: + + osprey-stress run \\ + --events 10000 --rate 1000 \\ + --threshold-drop-rate 0.01 --threshold-p95-ms 500 \\ + --report json + +Run an end-to-end stress test against a live Osprey worker. Produces N events +to the input topic at rate R, consumes the resulting `ExecutionResult`s from +the output topic, and reports drop rate + latency. Exits non-zero on threshold +breach so it can be wired into CI as a gate. + +The `measure` subcommand is reserved for when #236 (jetstream input stream +plugin) lands — it will let the same measurement layer run in open-loop mode +against any external input source. Today it prints a stub message. +""" + +from __future__ import annotations + +import argparse +import logging +import sys +import threading +import time +from collections.abc import Hashable, Mapping +from dataclasses import dataclass +from typing import Callable, Optional, Sequence, TextIO, cast + +from kafka import KafkaConsumer +from kafka.structs import TopicPartition +from osprey.worker.stress.consumer import Consumer, ConsumerConfig +from osprey.worker.stress.producer import Producer, ProducerConfig, _action_id_for +from osprey.worker.stress.reporter import ( + EXIT_INTERNAL_ERROR, + Thresholds, + compute_report, + exit_code_for, +) + +logger = logging.getLogger(__name__) + +DEFAULT_BOOTSTRAP = 'localhost:9092' +DEFAULT_INPUT_TOPIC = 'osprey.actions_input' +DEFAULT_OUTPUT_TOPIC = 'osprey.execution_results' + + +@dataclass(frozen=True) +class TopicSnapshot: + """High-water-mark offsets for a topic at one point in time.""" + + topic: str + end_offsets: Mapping[int, int] + + @property + def total(self) -> int: + return sum(self.end_offsets.values()) + + +def probe_topic_head( + bootstrap_servers: list[str], + topic: str, + *, + consumer_factory: Callable[..., KafkaConsumer] = KafkaConsumer, +) -> TopicSnapshot: + """Snapshot a topic's current end-offsets without joining a consumer group. + + Used to surface drain progress and "prior runs left backlog" situations: the + worker's processing rate is observable as the rate at which the output + topic's head advances vs the rate the input topic's head advances. A topic + that doesn't exist yet returns an empty snapshot rather than raising. + """ + consumer = consumer_factory( + bootstrap_servers=bootstrap_servers, + group_id=None, + client_id='osprey-stress-probe', + consumer_timeout_ms=1000, + ) + try: + partitions = consumer.partitions_for_topic(topic) + if not partitions: + return TopicSnapshot(topic=topic, end_offsets={}) + tps = [TopicPartition(topic, p) for p in partitions] + consumer.assign(tps) + end_offsets = consumer.end_offsets(tps) + return TopicSnapshot(topic=topic, end_offsets={tp.partition: off for tp, off in end_offsets.items()}) + finally: + try: + consumer.close(autocommit=False) + except Exception: + # Probe is best-effort. A noisy close() during teardown would mask + # the snapshot value the caller needs. Log so it's discoverable + # without becoming user-visible noise. + logger.debug('probe_topic_head: consumer.close() failed for topic %s', topic, exc_info=True) + + +class ProgressReporter: + """Periodic stderr progress lines for an in-flight stress run. + + Decoupled from Consumer/Producer so they don't carry observability + concerns. The constructor takes accessor callables for the live values it + needs to print, so the same reporter works against the real consumer or a + fake in tests. + """ + + def __init__( + self, + *, + interval_seconds: float, + target_events: int, + get_produced: Callable[[], int], + get_matched: Callable[[], int], + probe_input: Callable[[], int], + probe_output: Callable[[], int], + output: TextIO = sys.stderr, + ): + self._interval = interval_seconds + self._target = target_events + self._get_produced = get_produced + self._get_matched = get_matched + self._probe_input = probe_input + self._probe_output = probe_output + self._output = output + self._stop = threading.Event() + self._thread: Optional[threading.Thread] = None + self._initial_input = 0 + self._initial_output = 0 + self._start_monotonic = 0.0 + + def start(self) -> None: + if self._thread is not None: + raise RuntimeError('ProgressReporter already started') + self._start_monotonic = time.monotonic() + # Probe baselines on the calling thread so the first periodic emit can + # diff against them without racing the worker's startup. Best-effort: + # if the broker hiccups now we'd rather start with zero baselines + # (consumer is already running, the user wants the run to proceed) + # than abort cmd_run from inside this constructor-adjacent path. + try: + self._initial_input = self._probe_input() + except Exception: + logger.debug('ProgressReporter.start: initial input probe failed', exc_info=True) + self._initial_input = 0 + try: + self._initial_output = self._probe_output() + except Exception: + logger.debug('ProgressReporter.start: initial output probe failed', exc_info=True) + self._initial_output = 0 + self._thread = threading.Thread(target=self._run, name='stress-progress', daemon=True) + self._thread.start() + + def stop(self) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=self._interval + 1.0) + + def _run(self) -> None: + # Event.wait returns True when set, False on timeout; we use it as an + # interruptible sleep so stop() takes effect within `interval` seconds. + while not self._stop.wait(self._interval): + try: + self._emit() + except Exception: + # A failed probe (broker hiccup, transient timeout) shouldn't + # tear down the in-flight stress run. Skip the tick but log so + # repeated failures are diagnosable. + logger.debug('ProgressReporter tick failed; skipping', exc_info=True) + + def _emit(self) -> None: + elapsed = max(time.monotonic() - self._start_monotonic, 0.001) + produced = self._get_produced() + matched = self._get_matched() + in_delta = self._probe_input() - self._initial_input + out_delta = self._probe_output() - self._initial_output + match_rate = matched / elapsed + print( + f'[osprey-stress] t={elapsed:4.0f}s ' + f'produced={produced}/{self._target} matched={matched} ' + f'match_rate={match_rate:5.0f}/s ' + f'in_topic_Δ={in_delta} out_topic_Δ={out_delta}', + file=self._output, + ) + + +def _positive_int(s: str) -> int: + value = int(s) + if value <= 0: + raise argparse.ArgumentTypeError(f'must be > 0, got {value}') + return value + + +def _positive_float(s: str) -> float: + value = float(s) + if value <= 0: + raise argparse.ArgumentTypeError(f'must be > 0, got {value}') + return value + + +def _nonneg_float(s: str) -> float: + value = float(s) + if value < 0: + raise argparse.ArgumentTypeError(f'must be >= 0, got {value}') + return value + + +def _add_common_kafka_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + '--bootstrap-servers', + default=DEFAULT_BOOTSTRAP, + help=f'Comma-separated Kafka bootstrap servers (default: {DEFAULT_BOOTSTRAP}).', + ) + parser.add_argument( + '--input-topic', + default=DEFAULT_INPUT_TOPIC, + help=f'Kafka topic to produce events to (default: {DEFAULT_INPUT_TOPIC}).', + ) + parser.add_argument( + '--output-topic', + default=DEFAULT_OUTPUT_TOPIC, + help=f'Kafka topic to read execution results from (default: {DEFAULT_OUTPUT_TOPIC}).', + ) + + +def _add_threshold_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + '--threshold-drop-rate', + type=float, + default=None, + help='Exit non-zero if observed drop rate exceeds this (0.0–1.0).', + ) + parser.add_argument( + '--threshold-p95-ms', + type=float, + default=None, + help='Exit non-zero if observed p95 round-trip latency exceeds this (ms).', + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog='osprey-stress', + description=( + 'Produce synthetic events to Osprey, observe results, report drop ' + 'rate and latency. Useful for validating dependency bumps, ' + 'measuring throughput regressions, and gating CI on pipeline health.' + ), + ) + subparsers = parser.add_subparsers(dest='command', required=True) + + run = subparsers.add_parser( + 'run', + help='Produce synthetic events and measure their round-trip (closed-loop).', + ) + run.add_argument('--events', type=_positive_int, default=1000, help='Number of events to produce.') + run.add_argument('--rate', type=_positive_float, default=100.0, help='Events per second.') + run.add_argument( + '--drain-seconds', + type=_nonneg_float, + default=30.0, + help='Max wall-clock time to wait for in-flight events to be evaluated after producer finishes.', + ) + run.add_argument( + '--report', + choices=('human', 'json'), + default='human', + help='Report format on stdout.', + ) + run.add_argument( + '--verbose', + action='store_true', + help='Emit periodic progress lines to stderr (live throughput, topic deltas).', + ) + run.add_argument( + '--verbose-interval-seconds', + type=_positive_float, + default=2.0, + help='Seconds between verbose progress lines (default: 2.0). Ignored without --verbose.', + ) + _add_common_kafka_args(run) + _add_threshold_args(run) + + measure = subparsers.add_parser( + 'measure', + help='Measure-only mode (open-loop). Stub until #236 lands.', + ) + _add_common_kafka_args(measure) + _add_threshold_args(measure) + measure.add_argument('--duration', type=_positive_float, default=60.0) + measure.add_argument('--report', choices=('human', 'json'), default='human') + + return parser + + +def _bootstrap_list(s: str) -> list[str]: + return [host.strip() for host in s.split(',') if host.strip()] + + +def cmd_run(args: argparse.Namespace) -> int: + run_id = ProducerConfig.make_run_id() + expected_action_ids = frozenset(_action_id_for(run_id, n) for n in range(args.events)) + bootstrap_list = _bootstrap_list(args.bootstrap_servers) + + print( + f'[osprey-stress] run_id={run_id} events={args.events} rate={args.rate:.0f}/s ' + f'bootstrap={args.bootstrap_servers} input={args.input_topic} output={args.output_topic}', + file=sys.stderr, + ) + + # Snapshot input + output topic heads BEFORE the producer starts. This + # baseline is the regression-across-runs signal: if input topic head is + # already high (relative to expected throughput), the worker is still + # working through events from a prior run, and the numbers from THIS run + # will be misleading until that backlog drains. + try: + input_baseline = probe_topic_head(bootstrap_list, args.input_topic) + output_baseline = probe_topic_head(bootstrap_list, args.output_topic) + except Exception as e: + print(f'[osprey-stress] could not probe topic heads ({e}); skipping baseline', file=sys.stderr) + input_baseline = TopicSnapshot(args.input_topic, {}) + output_baseline = TopicSnapshot(args.output_topic, {}) + print( + f'[osprey-stress] baseline: input_topic head={input_baseline.total} output_topic head={output_baseline.total}', + file=sys.stderr, + ) + + consumer = Consumer( + ConsumerConfig( + bootstrap_servers=bootstrap_list, + topic=args.output_topic, + group_id=f'osprey-stress-{run_id}', + action_id_filter=expected_action_ids, + max_runtime_seconds=(args.events / args.rate) + args.drain_seconds, + stop_when_filter_complete=True, + ) + ) + consumer.start() + # Give the consumer time to join the group and get a partition assignment + # before the producer starts emitting. Without this, results from the + # earliest events can be missed. + time.sleep(2.0) + if consumer.error is not None: + print(f'[osprey-stress] consumer failed to start: {consumer.error}', file=sys.stderr) + return EXIT_INTERNAL_ERROR + + producer = Producer( + ProducerConfig( + bootstrap_servers=bootstrap_list, + topic=args.input_topic, + events=args.events, + rate_per_second=args.rate, + run_id=run_id, + ) + ) + + reporter: Optional[ProgressReporter] = None + if args.verbose: + reporter = ProgressReporter( + interval_seconds=args.verbose_interval_seconds, + target_events=args.events, + get_produced=lambda: len(producer.produced), + get_matched=lambda: len(consumer.consumed), + probe_input=lambda: probe_topic_head(bootstrap_list, args.input_topic).total, + probe_output=lambda: probe_topic_head(bootstrap_list, args.output_topic).total, + ) + reporter.start() + + start_wall = time.monotonic() + producer.start() + producer.wait() + if producer.error is not None: + consumer.stop() + consumer.wait(timeout=5) + if reporter is not None: + reporter.stop() + print(f'[osprey-stress] producer failed: {producer.error}', file=sys.stderr) + return EXIT_INTERNAL_ERROR + + print('[osprey-stress] producer done, draining consumer...', file=sys.stderr) + consumer.wait() # bounded by max_runtime_seconds set above + end_wall = time.monotonic() + if reporter is not None: + reporter.stop() + + if consumer.error is not None: + print(f'[osprey-stress] consumer error: {consumer.error}', file=sys.stderr) + return EXIT_INTERNAL_ERROR + + # Post-run snapshot: how far did each topic actually advance? If + # output_topic advanced by far less than `args.events`, the worker is the + # bottleneck (regardless of what the report says about THIS run's matches). + try: + input_after = probe_topic_head(bootstrap_list, args.input_topic) + output_after = probe_topic_head(bootstrap_list, args.output_topic) + in_delta = input_after.total - input_baseline.total + out_delta = output_after.total - output_baseline.total + print( + f'[osprey-stress] topic deltas: input +{in_delta} (expected ~{args.events}) output +{out_delta}', + file=sys.stderr, + ) + if out_delta < args.events: + shortfall = args.events - out_delta + print( + f'[osprey-stress] worker emitted {shortfall} fewer results than events sent — ' + f'pipeline backpressure; next run will start with this backlog still in flight', + file=sys.stderr, + ) + except Exception as e: + print(f'[osprey-stress] could not probe post-run topic heads ({e})', file=sys.stderr) + + # cast: `dict[int, float]` is a valid `Mapping[Hashable, float]` at runtime + # (int is Hashable), but mypy treats `Mapping` keys as invariant so the + # auto-conversion doesn't fly. The cast is honest. + report = compute_report( + produced=cast(Mapping[Hashable, float], producer.produced), + consumed=cast(Mapping[Hashable, float], consumer.consumed), + duration_seconds=end_wall - start_wall, + thresholds=Thresholds(drop_rate=args.threshold_drop_rate, p95_ms=args.threshold_p95_ms), + ) + + if args.report == 'json': + print(report.to_json()) + else: + print(report.to_human()) + + return exit_code_for(report) + + +def cmd_measure(args: argparse.Namespace) -> int: + print( + '[osprey-stress] `measure` is not yet implemented. It will activate once ' + '#236 (jetstream input stream plugin) lands so this CLI can run the ' + 'measurement layer against externally-produced events. ' + 'Use `osprey-stress run` for synthetic closed-loop testing.', + file=sys.stderr, + ) + return EXIT_INTERNAL_ERROR + + +def main(argv: Sequence[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + if args.command == 'run': + return cmd_run(args) + if args.command == 'measure': + return cmd_measure(args) + parser.print_help(sys.stderr) + return EXIT_INTERNAL_ERROR + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/osprey_worker/src/osprey/worker/stress/tests/test_cli.py b/osprey_worker/src/osprey/worker/stress/tests/test_cli.py new file mode 100644 index 0000000..95b9530 --- /dev/null +++ b/osprey_worker/src/osprey/worker/stress/tests/test_cli.py @@ -0,0 +1,270 @@ +import io +import time +from typing import Callable +from unittest.mock import MagicMock + +import pytest +from kafka.structs import TopicPartition +from osprey.worker.stress.cli import ( + EXIT_INTERNAL_ERROR, + ProgressReporter, + TopicSnapshot, + build_parser, + cmd_measure, + main, + probe_topic_head, +) + + +def _wait_until(predicate: Callable[[], bool], deadline_seconds: float = 5.0, sleep_step: float = 0.01) -> None: + """Block until `predicate()` returns True or fail after `deadline_seconds`. + + Used by the ProgressReporter tests to avoid fixed sleep() windows that flake + on busy CI runners. + """ + deadline = time.monotonic() + deadline_seconds + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(sleep_step) + raise AssertionError(f'predicate did not become true within {deadline_seconds}s') + + +class TestParser: + def test_run_minimal_args_parse(self) -> None: + args = build_parser().parse_args(['run']) + assert args.command == 'run' + assert args.events == 1000 + assert args.rate == 100.0 + assert args.report == 'human' + + def test_run_threshold_args_parse(self) -> None: + args = build_parser().parse_args( + [ + 'run', + '--events', + '50', + '--rate', + '10', + '--threshold-drop-rate', + '0.05', + '--threshold-p95-ms', + '750', + '--report', + 'json', + ] + ) + assert args.events == 50 + assert args.rate == 10.0 + assert args.threshold_drop_rate == 0.05 + assert args.threshold_p95_ms == 750.0 + assert args.report == 'json' + + def test_run_kafka_args_parse(self) -> None: + args = build_parser().parse_args( + [ + 'run', + '--bootstrap-servers', + 'kafka-a:9092,kafka-b:9092', + '--input-topic', + 'my.input', + '--output-topic', + 'my.output', + ] + ) + assert args.bootstrap_servers == 'kafka-a:9092,kafka-b:9092' + assert args.input_topic == 'my.input' + assert args.output_topic == 'my.output' + + def test_no_command_requires_one(self) -> None: + with pytest.raises(SystemExit): + build_parser().parse_args([]) + + @pytest.mark.parametrize( + 'argv', + [ + ['run', '--events', '0'], + ['run', '--events', '-5'], + ['run', '--rate', '0'], + ['run', '--rate', '-1.5'], + ['run', '--drain-seconds', '-0.1'], + ['measure', '--duration', '0'], + ['measure', '--duration', '-10'], + ], + ) + def test_rejects_non_positive_numeric_args(self, argv: list[str]) -> None: + # Regression: previously the parser accepted 0 / negatives, which led + # to ZeroDivisionError later when computing max_runtime_seconds. + with pytest.raises(SystemExit): + build_parser().parse_args(argv) + + +class TestMeasureStub: + def test_measure_returns_internal_error_for_now(self, capsys: pytest.CaptureFixture[str]) -> None: + args = build_parser().parse_args(['measure']) + rc = cmd_measure(args) + assert rc == EXIT_INTERNAL_ERROR + captured = capsys.readouterr() + assert '#236' in captured.err + + def test_main_dispatches_to_measure(self) -> None: + rc = main(['measure']) + assert rc == EXIT_INTERNAL_ERROR + + +class TestVerboseFlagParsing: + def test_verbose_defaults_off(self) -> None: + args = build_parser().parse_args(['run']) + assert args.verbose is False + assert args.verbose_interval_seconds == 2.0 + + def test_verbose_flag_sets_true(self) -> None: + args = build_parser().parse_args(['run', '--verbose']) + assert args.verbose is True + + def test_verbose_interval_custom(self) -> None: + args = build_parser().parse_args(['run', '--verbose', '--verbose-interval-seconds', '0.5']) + assert args.verbose_interval_seconds == 0.5 + + def test_verbose_interval_rejects_non_positive(self) -> None: + with pytest.raises(SystemExit): + build_parser().parse_args(['run', '--verbose-interval-seconds', '0']) + + +def _fake_kafka_consumer(partitions: set[int], end_offsets: dict[int, int]) -> MagicMock: + """Build a MagicMock that mimics enough of `KafkaConsumer` for `probe_topic_head`.""" + fake = MagicMock() + fake.partitions_for_topic.return_value = partitions + fake.end_offsets.side_effect = lambda tps: {tp: end_offsets[tp.partition] for tp in tps} + return fake + + +class TestProbeTopicHead: + def test_sums_partition_offsets(self) -> None: + fake = _fake_kafka_consumer({0, 1, 2}, {0: 100, 1: 250, 2: 75}) + snapshot = probe_topic_head(['localhost:9092'], 'osprey.actions_input', consumer_factory=lambda **_: fake) + assert snapshot.topic == 'osprey.actions_input' + assert snapshot.end_offsets == {0: 100, 1: 250, 2: 75} + assert snapshot.total == 425 + fake.close.assert_called_once() + + def test_empty_when_topic_missing(self) -> None: + # partitions_for_topic returns None for nonexistent topics — the probe + # tolerates this so a brand-new cluster doesn't crash the CLI. + fake = MagicMock() + fake.partitions_for_topic.return_value = None + snapshot = probe_topic_head(['localhost:9092'], 'no.such.topic', consumer_factory=lambda **_: fake) + assert snapshot.total == 0 + assert snapshot.end_offsets == {} + fake.assign.assert_not_called() + + def test_close_failure_is_swallowed(self) -> None: + fake = _fake_kafka_consumer({0}, {0: 10}) + fake.close.side_effect = RuntimeError('broker went away') + # Should not raise — probe is best-effort. + snapshot = probe_topic_head(['localhost:9092'], 't', consumer_factory=lambda **_: fake) + assert snapshot.total == 10 + + def test_assigns_topic_partitions(self) -> None: + fake = _fake_kafka_consumer({0, 1}, {0: 5, 1: 8}) + probe_topic_head(['localhost:9092'], 'my.topic', consumer_factory=lambda **_: fake) + assigned: list[TopicPartition] = fake.assign.call_args.args[0] + assert {tp.partition for tp in assigned} == {0, 1} + assert all(tp.topic == 'my.topic' for tp in assigned) + + +class TestProgressReporter: + def test_emits_periodic_progress_with_topic_deltas(self) -> None: + # Counters the reporter pulls from each tick. Producer climbs, matched + # lags behind, topics advance — the line should include all four signals. + produced = [0] + matched = [0] + input_offset = [1000] + output_offset = [500] + + def tick(p: int, m: int, in_o: int, out_o: int) -> None: + produced[0] = p + matched[0] = m + input_offset[0] = in_o + output_offset[0] = out_o + + sink = io.StringIO() + reporter = ProgressReporter( + interval_seconds=0.05, + target_events=100, + get_produced=lambda: produced[0], + get_matched=lambda: matched[0], + probe_input=lambda: input_offset[0], + probe_output=lambda: output_offset[0], + output=sink, + ) + reporter.start() + # Initial baseline taken in start() = (input=1000, output=500). + # First tick: in_delta=100, out_delta=50; wait for that to appear. + tick(50, 30, 1100, 550) + _wait_until(lambda: 'in_topic_Δ=100' in sink.getvalue()) + # Second tick: in_delta=200, out_delta=100. Once that line lands the + # reporter has emitted at least twice and we have the state we need. + tick(100, 70, 1200, 600) + _wait_until(lambda: 'in_topic_Δ=200' in sink.getvalue()) + reporter.stop() + + lines = [ln for ln in sink.getvalue().splitlines() if ln] + assert lines, 'expected at least one progress line' + last = lines[-1] + assert 'produced=' in last + assert 'matched=' in last + assert 'in_topic_Δ=' in last + assert 'out_topic_Δ=' in last + + def test_double_start_raises(self) -> None: + reporter = ProgressReporter( + interval_seconds=1.0, + target_events=1, + get_produced=lambda: 0, + get_matched=lambda: 0, + probe_input=lambda: 0, + probe_output=lambda: 0, + output=io.StringIO(), + ) + reporter.start() + try: + with pytest.raises(RuntimeError): + reporter.start() + finally: + reporter.stop() + + def test_probe_failure_does_not_kill_reporter(self) -> None: + # Transient broker errors during a probe shouldn't tear down the run. + calls = [0] + + def flaky_probe() -> int: + calls[0] += 1 + if calls[0] % 2 == 0: + raise RuntimeError('transient kafka error') + return calls[0] + + sink = io.StringIO() + reporter = ProgressReporter( + interval_seconds=0.05, + target_events=10, + get_produced=lambda: 0, + get_matched=lambda: 0, + probe_input=flaky_probe, + probe_output=lambda: 0, + output=sink, + ) + reporter.start() + # Half the probes raise; the reporter must keep ticking past failures. + # Wait until at least one successful tick has emitted output. + _wait_until(lambda: bool(sink.getvalue())) + reporter.stop() + + +class TestTopicSnapshot: + def test_total_sums_partition_offsets(self) -> None: + snap = TopicSnapshot(topic='t', end_offsets={0: 10, 1: 20, 2: 5}) + assert snap.total == 35 + + def test_total_zero_for_empty(self) -> None: + assert TopicSnapshot(topic='t', end_offsets={}).total == 0 -- 2.51.2