Add functionality for executing free and iostat

This CL adds functionality for executing free and iostat shell
commands as well as the unit tests for them.

Change-Id: I14088ce8ca2db2193d4ad12e7dffe0436c070c8d
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/19450
Reviewed-by: Vaibhav Rustagi <vaibhavrustagi@google.com>
Reviewed-by: Dexter Rivera <riverade@google.com>
Tested-by: Vaibhav Rustagi <vaibhavrustagi@google.com>
Cloud-Build: Vaibhav Rustagi <vaibhavrustagi@google.com>
diff --git a/src/pkg/nodeprofiler/profiler/commands.go b/src/pkg/nodeprofiler/profiler/commands.go
index a41d779..4fbda71 100644
--- a/src/pkg/nodeprofiler/profiler/commands.go
+++ b/src/pkg/nodeprofiler/profiler/commands.go
@@ -12,28 +12,18 @@
 // structs to execute shell commands.
 type Command interface {
 	Name() string
-	Run(opts Options) (map[string][]string, error)
-}
-
-// Options stores the options that can be passed to a command and
-// to parsing functions.
-type Options struct {
-	// Delay specifies times between updates in seconds.
-	Delay int
-	// Count specifies number of updates.
-	Count int
-	// Titles specifies the titles to get values for.
-	Titles []string
+	Run() (map[string][]string, error)
 }
 
 // vmstat represents a vmstat command.
 type vmstat struct {
 	name string
-}
-
-// lscpu represents an lscpu command.
-type lscpu struct {
-	name string
+	// delay specifies times between updates in seconds.
+	delay int
+	// count specifies number of updates.
+	count int
+	// titles specifies the titles to get values for.
+	titles []string
 }
 
 // Name returns the name for vmstat command.
@@ -41,32 +31,42 @@
 	return v.name
 }
 
+func (v *vmstat) setDefaults() {
+	if v.delay == 0 {
+		v.delay = 1
+	}
+	if v.count == 0 {
+		v.count = 5
+	}
+}
+
 // Run executes the vmstat command, parses the output and returns it as
 // a map of titles to their values.
-func (v *vmstat) Run(opts Options) (map[string][]string, error) {
-	// delay and count not set
-	if opts.Delay == 0 {
-		opts.Delay = 1
-	}
-	if opts.Count == 0 {
-		opts.Count = 5
-	}
-	interval := strconv.Itoa(opts.Delay)
-	count := strconv.Itoa(opts.Count)
+func (v *vmstat) Run() (map[string][]string, error) {
+	// if delay and count not set
+	v.setDefaults()
+	interval := strconv.Itoa(v.delay)
+	count := strconv.Itoa(v.count)
 	out, err := utils.RunCommand(v.Name(), "-n", interval, count)
 	if err != nil {
-		return nil, fmt.Errorf("failed to run vmstat command: %v", err)
+		return nil, fmt.Errorf("failed to run the command 'vmstat': %v", err)
 	}
 
 	s := string(out)
 	lines := strings.Split(strings.Trim(s, "\n"), "\n")
 	// ignore the first row in vmstat's output
 	lines = lines[1:]
-	titles := opts.Titles
+	titles := v.titles
 	// parse output by columns
 	output, err := utils.ParseColumns(lines, titles...)
 	return output, err
+}
 
+// lscpu represents an lscpu command.
+type lscpu struct {
+	name string
+	// titles specifies the titles to get values for.
+	titles []string
 }
 
 // Name returns the name for the lscpu command.
@@ -76,15 +76,94 @@
 
 // Run executes the lscpu command, parses the output and returns a
 // a map of title(s) to their values.
-func (l *lscpu) Run(opts Options) (map[string][]string, error) {
+func (l *lscpu) Run() (map[string][]string, error) {
 	out, err := utils.RunCommand(l.Name())
 	if err != nil {
-		return nil, fmt.Errorf("failed to run vmstat command: %v", err)
+		return nil, fmt.Errorf("failed to run the command 'lscpu': %v", err)
 	}
 	s := string(out)
 	lines := strings.Split(strings.Trim(s, "\n"), "\n")
-	titles := opts.Titles
+	titles := l.titles
 	// parse output by rows
-	output, err := utils.ParseRows(lines, titles...)
+	output, err := utils.ParseRows(lines, ":", titles...)
+	return output, err
+}
+
+// free represents a free command.
+type free struct {
+	name string
+	// titles specifies the titles to get values for.
+	titles []string
+}
+
+// Name returns the name for the free command.
+func (f *free) Name() string {
+	return f.name
+}
+
+// Run executes the free commands, parses the output and returns a
+// a map of title(s) to their values.
+func (f *free) Run() (map[string][]string, error) {
+	out, err := utils.RunCommand(f.Name(), "-m")
+	if err != nil {
+		return nil, fmt.Errorf("failed to run the command 'free': %v", err)
+	}
+
+	s := string(out)
+	lines := strings.Split(strings.Trim(s, "\n"), "\n")
+	titles := f.titles
+	// parse output by rows and columns
+	output, err := utils.ParseRowsAndColumns(lines, titles...)
+
+	return output, err
+}
+
+// iostat represents an iostat command
+type iostat struct {
+	name string
+	// flags specify the flags to be passed into the command.
+	flags string
+	// delay specifies times between updates in seconds.
+	delay int
+	// count specifies number of updates.
+	count int
+	// titles specifies the titles to get values for.
+	titles []string
+}
+
+// Name returns the name for the iostat command.
+func (i *iostat) Name() string {
+	return i.name
+}
+
+func (i *iostat) setDefaults() {
+	if i.delay == 0 {
+		i.delay = 1
+	}
+	if i.count == 0 {
+		i.count = 5
+	}
+}
+
+// Run executes the iostat commands, parses the output and returns a
+// a map of title(s) to their values.
+func (i *iostat) Run() (map[string][]string, error) {
+	// if delay and count not set
+	i.setDefaults()
+	interval := strconv.Itoa(i.delay)
+	count := strconv.Itoa(i.count)
+	out, err := utils.RunCommand(i.Name(), i.flags, interval, count)
+	if err != nil {
+		return nil, fmt.Errorf("failed to run the command 'iostat': %v", err)
+	}
+
+	s := string(out)
+	lines := strings.Split(strings.Trim(s, "\n"), "\n")
+	titles := i.titles
+	// ignore the first 2 lines in iostat's output so that the first line
+	// is column titles.
+	lines = lines[2:]
+	// parse output by rows and columns
+	output, err := utils.ParseColumns(lines, titles...)
 	return output, err
 }
diff --git a/src/pkg/nodeprofiler/profiler/commands_test.go b/src/pkg/nodeprofiler/profiler/commands_test.go
index fcb8b80..b03698e 100644
--- a/src/pkg/nodeprofiler/profiler/commands_test.go
+++ b/src/pkg/nodeprofiler/profiler/commands_test.go
@@ -10,17 +10,15 @@
 	tests := []struct {
 		name    string
 		fakeCmd Command
-		opts    Options
 		want    map[string][]string
 		wantErr bool
 	}{
 		{
-			name:    "vmstat",
-			fakeCmd: &vmstat{"testdata/vmstat.sh"},
-			opts: Options{
-				Delay:  1,
-				Count:  3,
-				Titles: []string{"us", "st", "sy"},
+			name: "vmstat",
+			fakeCmd: &vmstat{
+				name:   "testdata/vmstat.sh",
+				count:  3,
+				titles: []string{"us", "st", "sy"},
 			},
 			want: map[string][]string{
 				"us": {"1", "2", "7"},
@@ -29,21 +27,45 @@
 			},
 		},
 		{
-			name:    "lscpu",
-			fakeCmd: &lscpu{"testdata/lscpu.sh"},
-			opts: Options{
-				Titles: []string{"CPU(s)"},
+			name: "lscpu",
+			fakeCmd: &lscpu{
+				name:   "testdata/lscpu.sh",
+				titles: []string{"CPU(s)"},
 			},
 			want: map[string][]string{
 				"CPU(s)": {"8"},
 			},
 		},
 		{
-			name:    "no titles",
-			fakeCmd: &vmstat{"testdata/vmstat.sh"},
-			opts: Options{
-				Delay: 1,
-				Count: 2,
+			name: "free",
+			fakeCmd: &free{
+				name: "testdata/free.sh",
+				titles: []string{"Mem:used", "Mem:total",
+					"Swap:used", "Swap:total"},
+			},
+			want: map[string][]string{
+				"Mem:used":   {"13"},
+				"Mem:total":  {"14520"},
+				"Swap:used":  {"0"},
+				"Swap:total": {"0"},
+			},
+		},
+		{
+			name: "iostat",
+			fakeCmd: &iostat{
+				name:   "testdata/iostat.sh",
+				flags:  "xdz",
+				titles: []string{"%util"},
+			},
+			want: map[string][]string{
+				"%util": {"5.59", "0.00"},
+			},
+		},
+		{
+			name: "no titles",
+			fakeCmd: &vmstat{
+				name:  "testdata/vmstat.sh",
+				count: 2,
 			},
 			want: map[string][]string{
 				"r":  {"3", "1"},
@@ -55,12 +77,11 @@
 			},
 		},
 		{
-			name:    "spaced rows",
-			fakeCmd: &vmstat{"testdata/vmstat.sh"},
-			opts: Options{
-				Delay:  1,
-				Count:  4,
-				Titles: []string{"us", "st", "sy"},
+			name: "spaced rows",
+			fakeCmd: &vmstat{
+				name:   "testdata/vmstat.sh",
+				count:  4,
+				titles: []string{"us", "st", "sy"},
 			},
 			want: map[string][]string{
 				"us": {"7", "3", "1", "1"},
@@ -68,16 +89,24 @@
 				"st": {"0", "0", "0", "0"},
 			},
 		},
+		{
+			name: "illegal argument",
+			fakeCmd: &vmstat{
+				name:  "testdata/vmstat.sh",
+				count: -4,
+			},
+			wantErr: true,
+		},
 	}
 
 	for _, test := range tests {
-		got, err := test.fakeCmd.Run(test.opts)
+		got, err := test.fakeCmd.Run()
 
 		if gotErr := err != nil; gotErr != test.wantErr {
-			t.Fatalf("Run(%v) err %q, wantErr %v", test.opts, err, test.wantErr)
+			t.Fatalf("Run() err %v, wantErr %t", err, test.wantErr)
 		}
 		if diff := cmp.Diff(test.want, got); diff != "" {
-			t.Errorf("Ran Run(%v), but got mismatch between got and want (-got, +want): \n diff %s", test.opts, diff)
+			t.Errorf("Ran Run(), but got mismatch between got and want (-got, +want): \n diff %s", diff)
 		}
 	}
 }
diff --git a/src/pkg/nodeprofiler/profiler/testdata/free.sh b/src/pkg/nodeprofiler/profiler/testdata/free.sh
new file mode 100755
index 0000000..6637b3b
--- /dev/null
+++ b/src/pkg/nodeprofiler/profiler/testdata/free.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+
+main() {
+    cat <<EOF
+                  total        used        free      shared  buff/cache   available
+    Mem:          14520          13       14481           0          25       14506
+    Swap:             0           0           0
+EOF
+}
+
+main "$#"
\ No newline at end of file
diff --git a/src/pkg/nodeprofiler/profiler/testdata/iostat.sh b/src/pkg/nodeprofiler/profiler/testdata/iostat.sh
new file mode 100755
index 0000000..093bfa9
--- /dev/null
+++ b/src/pkg/nodeprofiler/profiler/testdata/iostat.sh
@@ -0,0 +1,21 @@
+#!/bin/bash
+
+main() {
+    cat <<EOF
+    Linux 5.4.109-26092-g9d947a4eeb73 (penguin)     07/09/2021      _x86_64_        (8 CPU)
+
+    Device            r/s     w/s     rkB/s     wkB/s   rrqm/s   wrqm/s  %rrqm  %wrqm r_await w_await aqu-sz rareq-sz wareq-sz  svctm  %util
+    vdb              0.01    0.60      0.86     21.39     0.00     0.20   0.24  25.16    8.82 1503.09   0.90    95.89    35.81  92.21   5.59
+    vda              0.00    0.00      0.04      0.00     0.00     0.00   2.73   0.00    3.08    0.00   0.00    62.55     0.00   2.20   0.00
+
+    Device            r/s     w/s     rkB/s     wkB/s   rrqm/s   wrqm/s  %rrqm  %wrqm r_await w_await aqu-sz rareq-sz wareq-sz  svctm  %util
+
+    Device            r/s     w/s     rkB/s     wkB/s   rrqm/s   wrqm/s  %rrqm  %wrqm r_await w_await aqu-sz rareq-sz wareq-sz  svctm  %util
+
+    Device            r/s     w/s     rkB/s     wkB/s   rrqm/s   wrqm/s  %rrqm  %wrqm r_await w_await aqu-sz rareq-sz wareq-sz  svctm  %util
+
+    Device            r/s     w/s     rkB/s     wkB/s   rrqm/s   wrqm/s  %rrqm  %wrqm r_await w_await aqu-sz rareq-sz wareq-sz  svctm  %util
+EOF
+}
+
+main "$#"
\ No newline at end of file
diff --git a/src/pkg/nodeprofiler/profiler/testdata/vmstat.sh b/src/pkg/nodeprofiler/profiler/testdata/vmstat.sh
index 125066c..6385b14 100755
--- a/src/pkg/nodeprofiler/profiler/testdata/vmstat.sh
+++ b/src/pkg/nodeprofiler/profiler/testdata/vmstat.sh
@@ -1,11 +1,9 @@
 #!/bin/bash
 
 main () {
-    if [[ "$#" -ne 3 ]]; then
-     echo "command not called with 3 arguments" >&2 return 1
-    fi
     if [[ $3 -lt 0 ]]; then
-        echo "$3 is not a valid argument" >&2 return 1
+        echo "$3 is not a valid argument" 1>&2 
+        return 1
     fi
     case "$3" in
         "3")
@@ -17,31 +15,6 @@
         2  0      0 14827096      0  25608    0    0     0     0 5283 8037  7  3 90  0  0
 EOF
         ;;
-        "0")
-        cat <<EOF
-        procs -----------memory---------- ---swap-- -----io---- -system-- ------cpu-----
-        r  b   swpd   free   buff  cache   si   so    bi    bo   in   cs us sy id wa st
-        9  0      0 14828152      0  25740    0    0     5     5   69   98  1  0 90  9  0
-EOF
-        ;;
-        "+5")
-        cat <<EOF
-        procs -----------memory---------- ---swap-- -----io---- -system-- ------cpu-----
-        r  b   swpd   free   buff  cache   si   so    bi    bo   in   cs us sy id wa st
-        3  0      0 14828112      0  25740    0    0     3     6   85  121  1  0 93  6  0
-        1  0      0 14828112      0  25740    0    0     0     0  854 1098  1  1 98  0  0
-        1  0      0 14828112      0  25740    0    0     0     0 1012 1399  2  1 98  0  0
-        1  0      0 14828112      0  25740    0    0     0     0 2991 4478  5  2 92  0  0
-        2  0      0 14828112      0  25740    0    0     0     0 7698 8623 18  6 75  0  1
-EOF
-        ;;
-        "1")
-        cat <<EOF
-        procs -----------memory---------- ---swap-- -----io---- -system-- ------cpu-----
-        r  b   swpd   free   buff  cache   si   so    bi    bo   in   cs us sy id wa st
-        2  0      0 14827724      0  25608    0    0     1     6   10   37  1  0 96  2  0
-EOF
-        ;;
         "4")
         cat <<EOF
         procs -----------memory---------- ---swap-- -----io---- -system-- ------cpu-----