Merge pull request #43 from compose-viz/dev

Fix service root KeyError

README.md

@@ -76,7 +76,9 @@ You need to install [Graphviz](https://graphviz.org/download/) to generate graph
 
 #### Using `pip`
 
-`pip install compose-viz`
+```
+pip install compose-viz
+```
 
 #### Using `.whl`
 
@@ -116,10 +118,10 @@ Check out the result [here](https://github.com/compose-viz/compose-viz/blob/main
 ### Options
 
 | Option                            | Description                                                                                                                                                                         |
-| ----------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `-o, --output-filename` | Output filename for the generated visualization file. [default: compose-viz]                                                                                                        |
-| `-m, --format`          | Output format for the generated visualization file. See [supported formats](https://github.com/compose-viz/compose-viz/blob/main/compose_viz/models/viz_formats.py). [default: png] |
-| `-r, --root-service`    | Root of the service tree (convenient for large compose yamls)                                                                                                                       |
+| --------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `-o, --output-filename FILENAME`  | Output filename for the generated visualization file. [default: compose-viz]                                                                                                        |
+| `-m, --format FORMAT`             | Output format for the generated visualization file. See [supported formats](https://github.com/compose-viz/compose-viz/blob/main/compose_viz/models/viz_formats.py). [default: png] |
+| `-r, --root-service SERVICE_NAME` | Root of the service tree (convenient for large compose yamls)                                                                                                                       |
 | `-v, --version`                   | Show the version of compose-viz.                                                                                                                                                    |
 | `--help`                          | Show help and exit.                                                                                                                                                                 |
 

compose_viz/parser.py

@@ -26,13 +26,14 @@ class Parser:
         return service_depends_on
 
     @staticmethod
-    def compile_dependencies(service_name: str, compose_data: spec.ComposeSpecification) -> List[str]:
-        assert compose_data.services
+    def compile_dependencies(service_name: str, services: Dict[Any, spec.Service], file_path: str) -> List[str]:
+        assert service_name in services, f"Service '{service_name}' not found in given compose file: '{file_path}'"
+
         dependencies = []
-        for dependency in Parser._unwrap_depends_on(compose_data.services[service_name].depends_on):
+        for dependency in Parser._unwrap_depends_on(services[service_name].depends_on):
             if dependency:
                 dependencies.append(dependency)
-                dependencies.extend(Parser.compile_dependencies(dependency, compose_data))
+                dependencies.extend(Parser.compile_dependencies(dependency, services, file_path))
         return dependencies
 
     def parse(self, file_path: str, root_service: Optional[str] = None) -> Compose:
@@ -47,10 +48,11 @@ class Parser:
 
         assert compose_data.services is not None, "No services found, aborting."
 
-        root_dependencies: List[str] = list()
+        root_dependencies: List[str] = []
         if root_service:
-            root_dependencies = Parser.compile_dependencies(root_service, compose_data)
+            root_dependencies = Parser.compile_dependencies(root_service, compose_data.services, file_path)
             root_dependencies.append(root_service)
+            root_dependencies = list(set(root_dependencies))
 
         for service_name, service_data in compose_data.services.items():
             service_name = str(service_name)

tests/test_cli.py

@@ -26,6 +26,7 @@ runner = CliRunner()
         "tests/ymls/volumes/docker-compose.yml",
         "examples/full-stack-node-app/docker-compose.yml",
         "examples/non-normative/docker-compose.yml",
+        "examples/voting-app/docker-compose.yml",
     ],
 )
 def test_cli(test_file_path: str) -> None:

tests/test_root_service.py

@@ -0,0 +1,32 @@
+import os
+
+from typer.testing import CliRunner
+
+from compose_viz import cli
+
+runner = CliRunner()
+
+
+def test_root_service() -> None:
+    input_path = "examples/voting-app/docker-compose.yml"
+    output_filename = "compose-viz-test"
+    default_format = "png"
+    result = runner.invoke(cli.app, ["-r", "vote", "-o", output_filename, input_path])
+
+    assert result.exit_code == 0
+    assert f"Successfully parsed {input_path}\n" in result.stdout
+    assert os.path.exists(f"{output_filename}.{default_format}")
+
+    os.remove(f"{output_filename}.{default_format}")
+
+
+def test_root_service_key_error() -> None:
+    input_path = "examples/voting-app/docker-compose.yml"
+    output_filename = "compose-viz-test"
+    default_format = "png"
+    result = runner.invoke(cli.app, ["-r", "not_exist_service", "-o", output_filename, input_path])
+
+    assert result.exit_code == 1
+    assert result.exception is not None
+    assert result.exception.args[0] == f"Service 'not_exist_service' not found in given compose file: '{input_path}'"
+    assert not os.path.exists(f"{output_filename}.{default_format}")