linting + adding unit tests for one file

Author: rreddy2

nodestream/cli/commands/print_schema.py

@@ -2,7 +2,7 @@
 
 from ..operations import InitializeProject, PrintProjectSchema
 from .nodestream_command import NodestreamCommand
-from .shared_options import PROJECT_FILE_OPTION, MANY_PIPELINES_ARGUMENT
+from .shared_options import MANY_PIPELINES_ARGUMENT, PROJECT_FILE_OPTION
 
 
 class PrintSchema(NodestreamCommand):

nodestream/cli/operations/print_project_schema.py

@@ -29,11 +29,11 @@ async def perform(self, command: NodestreamCommand):
         if self.pipeline_names:
             schema = self.project.get_pipelines_schema(
                 pipeline_names=self.pipeline_names,
-                type_overrides_file=type_overrides_file
+                type_overrides_file=type_overrides_file,
             )
         else:
             schema = self.project.get_schema(type_overrides_file=type_overrides_file)
-        
+
         # Import all schema printers so that they can register themselves
         SchemaPrinter.import_all()
         printer = SchemaPrinter.from_name(self.format_string)

nodestream/project/project.py

@@ -291,7 +291,9 @@ def get_schema(self, type_overrides_file: Optional[Path] = None) -> Schema:
             schema.merge(overrides_schema)
         return schema
 
-    def get_pipelines_schema(self, pipeline_names: List[str], type_overrides_file: Optional[Path] = None) -> Schema:
+    def get_pipelines_schema(
+        self, pipeline_names: List[str], type_overrides_file: Optional[Path] = None
+    ) -> Schema:
         """Returns a `GraphSchema` representing only the specified pipelines.
 
         This method generates a schema from only the specified pipelines,
@@ -313,33 +315,36 @@ def get_pipelines_schema(self, pipeline_names: List[str], type_overrides_file: O
             targets_by_name=self.targets_by_name,
             storage_configuration=self.storage_configuration,
         )
-        
+
         pipelines_found = []
-        
+
         for scope in self.scopes_by_name.values():
             filtered_scope = PipelineScope(
                 name=scope.name,
                 config=scope.config,
                 pipeline_configuration=scope.pipeline_configuration,
             )
-            
+
             for pipeline_name in pipeline_names:
                 if pipeline_name in scope.pipelines_by_name:
                     filtered_scope.add_pipeline_definition(
                         scope.pipelines_by_name[pipeline_name]
                     )
                     pipelines_found.append(pipeline_name)
-            
+
             if filtered_scope.pipelines_by_name:
                 filtered_project.add_scope(filtered_scope)
-        
+
         if not pipelines_found:
             available_pipelines = [
-                name for scope in self.scopes_by_name.values() 
+                name
+                for scope in self.scopes_by_name.values()
                 for name in scope.pipelines_by_name.keys()
             ]
-            raise ValueError(f"None of the specified pipelines {pipeline_names} were found. Available pipelines: {available_pipelines}")
-        
+            raise ValueError(
+                f"None of the specified pipelines {pipeline_names} were found. Available pipelines: {available_pipelines}"
+            )
+
         schema = filtered_project.make_schema()
         if type_overrides_file is not None:
             overrides_schema = Schema.read_from_file(type_overrides_file)

tests/unit/cli/operations/test_print_project_schema.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import pytest
 
 from nodestream.cli.operations import PrintProjectSchema
@@ -42,3 +44,53 @@ async def test_print_project_schema_prints_schema_to_file(
     _, stdout, fileout = await ran_print_project_schema_operation("some/path")
     assert not stdout.called
     assert fileout.called
+
+
+@pytest.mark.asyncio
+async def test_print_project_schema_with_pipeline_filtering(
+    mocker, project_with_default_scope
+):
+    std_out = mocker.patch(
+        "nodestream.schema.printers.SchemaPrinter.print_schema_to_stdout"
+    )
+    project_with_default_scope.get_schema = mocker.Mock(return_value="full_schema")
+    project_with_default_scope.get_pipelines_schema = mocker.Mock(
+        return_value="filtered_schema"
+    )
+
+    operation = PrintProjectSchema(
+        project=project_with_default_scope,
+        format_string="plain",
+        pipeline_names=["test_pipeline"],
+    )
+    await operation.perform(mocker.Mock())
+
+    # Should call get_pipelines_schema, not get_schema
+    project_with_default_scope.get_pipelines_schema.assert_called_once_with(
+        pipeline_names=["test_pipeline"], type_overrides_file=None
+    )
+    project_with_default_scope.get_schema.assert_not_called()
+    std_out.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_print_project_schema_with_pipeline_filtering_and_overrides(
+    mocker, project_with_default_scope
+):
+    project_with_default_scope.get_pipelines_schema = mocker.Mock(
+        return_value="filtered_schema"
+    )
+
+    operation = PrintProjectSchema(
+        project=project_with_default_scope,
+        format_string="plain",
+        type_overrides_file="overrides.yaml",
+        pipeline_names=["test_pipeline", "another_pipeline"],
+    )
+    await operation.perform(mocker.Mock())
+
+    # Should call get_pipelines_schema with correct arguments
+    project_with_default_scope.get_pipelines_schema.assert_called_once_with(
+        pipeline_names=["test_pipeline", "another_pipeline"],
+        type_overrides_file=Path("overrides.yaml"),
+    )