small refactoring, tests work again

images.go

@@ -85,18 +85,61 @@ func (x *ImagesCommand) Execute(args []string) error {
 	if imagesCommand.Tree || imagesCommand.Dot {
 		var startImage *Image
 		if len(args) > 0 {
+			startImage, err = findStartImage(args[0], images)
+
+			if err != nil {
+				return err
+			}
+		}
+
+		// select the start image of the tree
+		var roots []Image
+		if startImage == nil {
+			roots = collectRoots(images)
+		} else {
+			startImage.ParentId = ""
+			roots = []Image{*startImage}
+		}
+
+		// build helper map (image -> children)
+		imagesByParent := collectChildren(images)
+
+		// filter images
+		if imagesCommand.OnlyLabelled {
+			*images, imagesByParent = filterImages(images, &imagesByParent)
+		}
+
+		if imagesCommand.Tree {
+			fmt.Print(jsonToTree(imagesCommand.NoTruncate, roots, imagesByParent))
+		}
+		if imagesCommand.Dot {
+			fmt.Print(jsonToDot(roots, imagesByParent))
+		}
+
+	} else if imagesCommand.Short {
+		fmt.Printf(jsonToShort(images))
+	} else {
+		return fmt.Errorf("Please specify either --dot, --tree, or --short")
+	}
+
+	return nil
+}
+
+func findStartImage(name string, images *[]Image) (*Image, error) {
+
+	var startImage *Image
 
 	// attempt to find the start image, which can be specified as an
 	// image ID or a repository name
-			startImageArg := args[0]
-			startImageRepo := args[0]
+	startImageArg := name
+	startImageRepo := name
 
 	// if tag is not defined, find by :latest tag
 	if strings.Index(startImageRepo, ":") == -1 {
 		startImageRepo = fmt.Sprintf("%s:latest", startImageRepo)
 	}
 
-		IMAGES:
+IMAGES:
 	for _, image := range *images {
 		// find by image id
 		if strings.Index(image.Id, startImageArg) == 0 {
@@ -114,51 +157,28 @@ func (x *ImagesCommand) Execute(args []string) error {
 	}
 
 	if startImage == nil {
-				return fmt.Errorf("Unable to find image %s = %s.", startImageArg, startImageRepo)
-			}
+		return nil, fmt.Errorf("Unable to find image %s = %s.", startImageArg, startImageRepo)
 	}
 
-		// select the start image of the tree
-		var roots []Image
-		if startImage == nil {
-			roots = collectRoots(images)
-		} else {
-			startImage.ParentId = ""
-			roots = []Image{*startImage}
-		}
-
-		// build helper map (image -> children)
-		var imagesByParent = make(map[string][]Image)
-		imagesByParent = collectChildren(images)
-
-		// image ids truncate
-		// initialize image informations
-
-		// filter images
-		if imagesCommand.OnlyLabelled {
-			*images, imagesByParent = filterImages(images, &imagesByParent)
-		}
+	return startImage, nil
+}
 
+func jsonToTree(noTrunc bool, images []Image, byParent map[string][]Image) string {
+	var buffer bytes.Buffer
+
+	jsonToText(&buffer, noTrunc, images, byParent, "")
+
+	return buffer.String()
+}
+
+func jsonToDot(roots []Image, byParent map[string][]Image) string {
 	var buffer bytes.Buffer
 
-		if imagesCommand.Tree {
-			jsonToText(&buffer, imagesCommand.NoTruncate, roots, imagesByParent, "")
-		}
-		if imagesCommand.Dot {
 	buffer.WriteString("digraph docker {\n")
-			imagesToDot(&buffer, roots, imagesByParent)
+	imagesToDot(&buffer, roots, byParent)
 	buffer.WriteString(" base [style=invisible]\n}\n")
-			buffer.String()
-		}
 
-		fmt.Print(buffer.String())
-	} else if imagesCommand.Short {
-		fmt.Printf(jsonToShort(images))
-	} else {
-		return fmt.Errorf("Please specify either --dot, --tree, or --short")
-	}
-
-	return nil
+	return buffer.String()
 }
 
 func collectChildren(images *[]Image) map[string][]Image {

images_test.go

@@ -56,7 +56,12 @@ func Test_Dot(t *testing.T) {
 
 	for _, dotTest := range dotTests {
 		im, _ := parseImagesJSON([]byte(dotTest.json))
-		result := jsonToDot(im)
+		byParent := collectChildren(im)
+		roots := collectRoots(im)
+
+		// TODO: test start image limiting
+
+		result := jsonToDot(roots, byParent)
 
 		for _, regexp := range allRegex {
 			if !regexp.MatchString(result) {
@@ -109,7 +114,16 @@ func Test_Tree(t *testing.T) {
 
 	for _, treeTest := range treeTests {
 		im, _ := parseImagesJSON([]byte(treeTest.json))
-		result := jsonToTree(im, treeTest.startImage, treeTest.noTrunc)
+		byParent := collectChildren(im)
+		var roots []Image
+		if len(treeTest.startImage) > 0 {
+			startImage, _ := findStartImage(treeTest.startImage, im)
+			startImage.ParentId = ""
+			roots = []Image{*startImage}
+		} else {
+			roots = collectRoots(im)
+		}
+		result := jsonToTree(treeTest.noTrunc, roots, byParent)
 
 		for _, regexp := range compileRegexps(t, treeTest.regexps) {
 			if !regexp.MatchString(result) {