Added tests to viz subcommand

Signed-off-by: Benjamín Guzmán <[email protected]>

Author: BenjaminGuzman

cmd/compose/alpha.go

@@ -34,7 +34,7 @@ func alphaCommand(p *ProjectOptions, backend api.Service) *cobra.Command {
 	cmd.AddCommand(
 		watchCommand(p, backend),
 		dryRunRedirectCommand(p),
-		vizCommand(p),
+		vizCommand(p, backend),
 	)
 	return cmd
 }

cmd/compose/viz.go

@@ -20,10 +20,9 @@ import (
 	"context"
 	"fmt"
 	"os"
-	"strconv"
 	"strings"
 
-	"github.com/compose-spec/compose-go/types"
+	"github.com/docker/compose/v2/pkg/api"
 	"github.com/spf13/cobra"
 )
 
@@ -35,10 +34,7 @@ type vizOptions struct {
 	indentationStr   string
 }
 
-// maps a service with the services it depends on
-type vizGraph map[*types.ServiceConfig][]*types.ServiceConfig
-
-func vizCommand(p *ProjectOptions) *cobra.Command {
+func vizCommand(p *ProjectOptions, backend api.Service) *cobra.Command {
 	opts := vizOptions{
 		ProjectOptions: p,
 	}
@@ -54,7 +50,7 @@ func vizCommand(p *ProjectOptions) *cobra.Command {
 			return err
 		}),
 		RunE: Adapt(func(ctx context.Context, args []string) error {
-			return runViz(ctx, &opts)
+			return runViz(ctx, backend, &opts)
 		}),
 	}
 
@@ -66,118 +62,26 @@ func vizCommand(p *ProjectOptions) *cobra.Command {
 	return cmd
 }
 
-func runViz(_ context.Context, opts *vizOptions) error {
+func runViz(ctx context.Context, backend api.Service, opts *vizOptions) error {
 	_, _ = fmt.Fprintln(os.Stderr, "viz command is EXPERIMENTAL")
 	project, err := opts.ToProject(nil)
 	if err != nil {
 		return err
 	}
 
 	// build graph
-	graph := make(vizGraph)
-	for i, serviceConfig := range project.Services {
-		serviceConfigPtr := &project.Services[i]
-		graph[serviceConfigPtr] = make([]*types.ServiceConfig, 0, len(serviceConfig.DependsOn))
-		for dependencyName := range serviceConfig.DependsOn {
-			// no error should be returned since dependencyName should exist
-			dependency, _ := project.GetService(dependencyName)
-			graph[serviceConfigPtr] = append(graph[serviceConfigPtr], &dependency)
-		}
-	}
+	graphStr, _ := backend.Viz(ctx, project, api.VizOptions{
+		IncludeNetworks:  opts.includeNetworks,
+		IncludePorts:     opts.includePorts,
+		IncludeImageName: opts.includeImageName,
+		Indentation:      opts.indentationStr,
+	})
 
-	// build graphviz graph
-	var graphBuilder strings.Builder
-	graphBuilder.WriteString("digraph " + project.Name + " {\n")
-	graphBuilder.WriteString(opts.indentationStr + "layout=dot;\n")
-	addNodes(&graphBuilder, graph, opts)
-	graphBuilder.WriteByte('\n')
-	addEdges(&graphBuilder, graph, opts)
-	graphBuilder.WriteString("}\n")
-
-	fmt.Println(graphBuilder.String())
+	fmt.Println(graphStr)
 
 	return nil
 }
 
-// addNodes adds the corresponding graphviz representation of all the nodes in the given graph to the graphBuilder
-// returns the same graphBuilder
-func addNodes(graphBuilder *strings.Builder, graph vizGraph, opts *vizOptions) *strings.Builder {
-	for serviceNode := range graph {
-		// write:
-		// "service name" [style="filled" label<<font point-size="15">service name</font>
-		graphBuilder.WriteString(opts.indentationStr)
-		writeQuoted(graphBuilder, serviceNode.Name)
-		graphBuilder.WriteString(" [style=\"filled\" label=<<font point-size=\"15\">")
-		graphBuilder.WriteString(serviceNode.Name)
-		graphBuilder.WriteString("</font>")
-
-		if opts.includeNetworks && len(serviceNode.Networks) > 0 {
-			graphBuilder.WriteString("<font point-size=\"10\">")
-			graphBuilder.WriteString("<br/><br/><b>Networks:</b>")
-			for _, networkName := range serviceNode.NetworksByPriority() {
-				graphBuilder.WriteString("<br/>")
-				graphBuilder.WriteString(networkName)
-			}
-			graphBuilder.WriteString("</font>")
-		}
-
-		if opts.includePorts && len(serviceNode.Ports) > 0 {
-			graphBuilder.WriteString("<font point-size=\"10\">")
-			graphBuilder.WriteString("<br/><br/><b>Ports:</b>")
-			for _, portConfig := range serviceNode.Ports {
-				graphBuilder.WriteString("<br/>")
-				if len(portConfig.HostIP) > 0 {
-					graphBuilder.WriteString(portConfig.HostIP)
-					graphBuilder.WriteByte(':')
-				}
-				graphBuilder.WriteString(portConfig.Published)
-				graphBuilder.WriteByte(':')
-				graphBuilder.WriteString(strconv.Itoa(int(portConfig.Target)))
-				graphBuilder.WriteString(" (")
-				graphBuilder.WriteString(portConfig.Protocol)
-				graphBuilder.WriteString(", ")
-				graphBuilder.WriteString(portConfig.Mode)
-				graphBuilder.WriteString(")")
-			}
-			graphBuilder.WriteString("</font>")
-		}
-
-		if opts.includeImageName {
-			graphBuilder.WriteString("<font point-size=\"10\">")
-			graphBuilder.WriteString("<br/><br/><b>Image:</b><br/>")
-			graphBuilder.WriteString(serviceNode.Image)
-			graphBuilder.WriteString("</font>")
-		}
-
-		graphBuilder.WriteString(">];\n")
-	}
-
-	return graphBuilder
-}
-
-// addEdges adds the corresponding graphviz representation of all edges in the given graph to the graphBuilder
-// returns the same graphBuilder
-func addEdges(graphBuilder *strings.Builder, graph vizGraph, opts *vizOptions) *strings.Builder {
-	for parent, children := range graph {
-		for _, child := range children {
-			graphBuilder.WriteString(opts.indentationStr)
-			writeQuoted(graphBuilder, parent.Name)
-			graphBuilder.WriteString(" -> ")
-			writeQuoted(graphBuilder, child.Name)
-			graphBuilder.WriteString(";\n")
-		}
-	}
-
-	return graphBuilder
-}
-
-// writeQuoted writes "str" to builder
-func writeQuoted(builder *strings.Builder, str string) {
-	builder.WriteByte('"')
-	builder.WriteString(str)
-	builder.WriteByte('"')
-}
-
 // preferredIndentationStr returns a single string given the indentation preference
 func preferredIndentationStr(size int, useSpace bool) (string, error) {
 	if size < 0 {

cmd/compose/viz_test.go

@@ -0,0 +1,92 @@
+/*
+   Copyright 2020 Docker Compose CLI authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package compose
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPreferredIndentationStr(t *testing.T) {
+	type args struct {
+		size     int
+		useSpace bool
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    string
+		wantErr bool
+	}{
+		{
+			name: "should return '\\t\\t'",
+			args: args{
+				size:     2,
+				useSpace: false,
+			},
+			want:    "\t\t",
+			wantErr: false,
+		},
+		{
+			name: "should return '    '",
+			args: args{
+				size:     4,
+				useSpace: true,
+			},
+			want:    "    ",
+			wantErr: false,
+		},
+		{
+			name: "should return ''",
+			args: args{
+				size:     0,
+				useSpace: false,
+			},
+			want:    "",
+			wantErr: false,
+		},
+		{
+			name: "should return ''",
+			args: args{
+				size:     0,
+				useSpace: true,
+			},
+			want:    "",
+			wantErr: false,
+		},
+		{
+			name: "should throw error because indentation size < 0",
+			args: args{
+				size:     -1,
+				useSpace: false,
+			},
+			want:    "",
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := preferredIndentationStr(tt.args.size, tt.args.useSpace)
+			if tt.wantErr && assert.NotNilf(t, err, fmt.Sprintf("preferredIndentationStr(%v, %v)", tt.args.size, tt.args.useSpace)) {
+				return
+			}
+			assert.Equalf(t, tt.want, got, "preferredIndentationStr(%v, %v)", tt.args.size, tt.args.useSpace)
+		})
+	}
+}

pkg/api/api.go

@@ -82,6 +82,19 @@ type Service interface {
 	DryRunMode(ctx context.Context, dryRun bool) (context.Context, error)
 	// Watch services' development context and sync/notify/rebuild/restart on changes
 	Watch(ctx context.Context, project *types.Project, services []string, options WatchOptions) error
+	// Viz generates a graphviz graph of the project services
+	Viz(ctx context.Context, project *types.Project, options VizOptions) (string, error)
+}
+
+type VizOptions struct {
+	// IncludeNetworks if true, network names a container is attached to should appear in the graph node
+	IncludeNetworks bool
+	// IncludePorts if true, ports a container exposes should appear in the graph node
+	IncludePorts bool
+	// IncludeImageName if true, name of the image used to create a container should appear in the graph node
+	IncludeImageName bool
+	// Indentation string to be used to indent graphviz code, e.g. "\t", "    "
+	Indentation string
 }
 
 // WatchOptions group options of the Watch API

pkg/api/proxy.go

@@ -53,6 +53,7 @@ type ServiceProxy struct {
 	WatchFn              func(ctx context.Context, project *types.Project, services []string, options WatchOptions) error
 	MaxConcurrencyFn     func(parallel int)
 	DryRunModeFn         func(ctx context.Context, dryRun bool) (context.Context, error)
+	VizFn                func(ctx context.Context, project *types.Project, options VizOptions) (string, error)
 	interceptors         []Interceptor
 }
 
@@ -93,6 +94,7 @@ func (s *ServiceProxy) WithService(service Service) *ServiceProxy {
 	s.WatchFn = service.Watch
 	s.MaxConcurrencyFn = service.MaxConcurrency
 	s.DryRunModeFn = service.DryRunMode
+	s.VizFn = service.Viz
 	return s
 }
 
@@ -323,6 +325,14 @@ func (s *ServiceProxy) Watch(ctx context.Context, project *types.Project, servic
 	return s.WatchFn(ctx, project, services, options)
 }
 
+// Viz implements Viz interface
+func (s *ServiceProxy) Viz(ctx context.Context, project *types.Project, options VizOptions) (string, error) {
+	if s.VizFn == nil {
+		return "", ErrNotImplemented
+	}
+	return s.VizFn(ctx, project, options)
+}
+
 func (s *ServiceProxy) MaxConcurrency(i int) {
 	s.MaxConcurrencyFn(i)
 }