diff --git a/defectdojo_api_generated/cli/commands/apis.py b/defectdojo_api_generated/cli/commands/apis.py index b9b3674..560fe3b 100644 --- a/defectdojo_api_generated/cli/commands/apis.py +++ b/defectdojo_api_generated/cli/commands/apis.py @@ -342,7 +342,12 @@ def _collect_command_values(parameters, getter, *, should_include): def _render_api_result(result: Any, instance: Any) -> None: - _render_result(result, json_mode=instance.json, jq_expression=instance.jq) + _render_result( + result, + json_mode=instance.json, + jq_expression=instance.jq, + max_records=getattr(instance, 'max_records', None), + ) def _create_standard_command_call(api_class: type, target_method: str, command_parameters): @@ -515,9 +520,13 @@ def _format_text_item(item: Any) -> str: return '\n'.join(lines) -def _render_result(result: Any, *, json_mode: bool, jq_expression: Optional[str]): +def _render_result(result: Any, *, json_mode: bool, jq_expression: Optional[str], max_records: Optional[int] = None): div = False + emitted = 0 for item in _iter_output_items(result): + if max_records is not None and emitted >= max_records: + break + item = _apply_jq(item, jq_expression) if json_mode: @@ -529,6 +538,8 @@ def _render_result(result: Any, *, json_mode: bool, jq_expression: Optional[str] div = True click.echo(_format_text_item(item)) + emitted += 1 + def make_api_command(api_class: type, command_name: str, target_method: str, *, parent_class: type): raw_parameters = list(_iter_command_parameters(api_class, target_method)) @@ -561,6 +572,15 @@ def make_api_command(api_class: type, command_name: str, target_method: str, *, default_getter=lambda item: item[1].default, ) + if target_method.endswith('_iterator'): + namespace['__annotations__']['max_records'] = Optional[int] + namespace['max_records'] = classyclick.Option( + '-m', + type=click.IntRange(min=1), + default=None, + help='Maximum number of records to emit.', + ) + return _build_command_class(parent_class, namespace) diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index 9396b1a..c83c7ba 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -84,6 +84,51 @@ def test_findings_list_accepts_limit_option(self): self.assertEqual(result.exit_code, 0) self.assertIn('limit: 1', result.output) + def test_findings_list_accepts_max_records_option(self): + runner = CliRunner() + with runner.isolated_filesystem(): + config_path = Path('config.toml') + config_path.write_text("host = 'https://example.com'\ntoken = 'token'\n") + + result = runner.invoke(CLI.click, ['--config', str(config_path), 'api', 'findings', 'list', '--help']) + + self.assertEqual(result.exit_code, 0) + self.assertIn('--max-records', result.output) + self.assertIn('-m', result.output) + + def test_findings_list_max_records_stops_output(self): + runner = CliRunner() + with runner.isolated_filesystem(): + config_path = Path('config.toml') + config_path.write_text("host = 'https://example.com'\ntoken = 'token'\n") + + with mock.patch.object( + FindingsApi, + 'list_iterator', + new=lambda self, **kwargs: [ + IteratorResult(result='first', page='page-1'), + IteratorResult(result='second', page='page-2'), + IteratorResult(result='third', page='page-3'), + ], + ): + result = runner.invoke( + CLI.click, + [ + '--config', + str(config_path), + 'api', + 'findings', + 'list', + '--limit', + '1', + '--max-records', + '2', + ], + ) + + self.assertEqual(result.exit_code, 0) + self.assertEqual(result.output.splitlines(), ['first', '', '---', '', 'second']) + def test_findings_list_prints_iterator_items_one_per_line(self): runner = CliRunner() with runner.isolated_filesystem():