Merge pull request #21 from thaim/add-all-flag

main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"runtime/debug"
@@ -42,9 +43,7 @@ func main() {
 			}
 			ids, err := Ec2id(name, client)
 			if err == nil && ids != nil {
-				for id := range ids {
-					fmt.Println(id)
-				}
+				printIds(os.Stdout, ids, all)
 			}
 			return err
 		},
@@ -72,3 +71,13 @@ func getVersion() string {
 
 	return i.Main.Version
 }
+
+func printIds(out io.Writer, ids []string, all bool) {
+	for _, id := range ids {
+		fmt.Fprintln(out, id)
+
+		if !all {
+			break
+		}
+	}
+}

main_test.go

@@ -0,0 +1,52 @@
+package main
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestMainPrintIds(t *testing.T) {
+	cases := []struct {
+		name    string
+		ids     []string
+		allFlag bool
+		expect  string
+	}{
+		{
+			name:    "print nothing with empty id",
+			ids:     nil,
+			allFlag: false,
+			expect:  "",
+		},
+		{
+			name:    "print single id with single id",
+			ids:     []string{"i-012345"},
+			allFlag: false,
+			expect:  "i-012345\n",
+		},
+		{
+			name:    "print single id with multiple ids and disabled all flag",
+			ids:     []string{"i-012345", "i-6789ab", "i-cdef01"},
+			allFlag: false,
+			expect:  "i-012345\n",
+		},
+		{
+			name:    "print multiple ids with multiple ids and enabled all flag",
+			ids:     []string{"i-012345", "i-6789ab", "i-cdef01"},
+			allFlag: true,
+			expect:  "i-012345\ni-6789ab\ni-cdef01\n",
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			var buf bytes.Buffer
+
+			printIds(&buf, tt.ids, tt.allFlag)
+
+			if bufString := buf.String(); bufString != tt.expect {
+				t.Errorf("expect %s, got id: %s", tt.expect, buf.String())
+			}
+		})
+	}
+}