cli: add --max-records (-m) cap for iterator list output (#72)

Author: fopina-ci

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)
 
 

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():