From c92715f833576c22510d5927175484e8649a2aa7 Mon Sep 17 00:00:00 2001
From: "James O. D. Hunt" <james.o.hunt@intel.com>
Date: Thu, 18 Jul 2019 10:43:51 +0100
Subject: [PATCH] tests: Add test for getMemory

Add a new test for `getMemory()`. This necessitated making
the memory file a `var` to allow the tests to modify it.

Signed-off-by: James O. D. Hunt <james.o.hunt@intel.com>
---
 agent.go      |   3 +-
 agent_test.go | 134 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 136 insertions(+), 1 deletion(-)

diff --git a/agent.go b/agent.go
index 199cf37943..29d4a9d096 100644
--- a/agent.go
+++ b/agent.go
@@ -42,7 +42,6 @@ import (
 
 const (
 	procCgroups = "/proc/cgroups"
-	meminfo     = "/proc/meminfo"
 
 	bashPath         = "/bin/bash"
 	shPath           = "/bin/sh"
@@ -50,6 +49,8 @@ const (
 )
 
 var (
+	meminfo = "/proc/meminfo"
+
 	// cgroup fs is mounted at /sys/fs when systemd is the init process
 	sysfsDir                     = "/sys"
 	cgroupPath                   = sysfsDir + "/fs/cgroup"
diff --git a/agent_test.go b/agent_test.go
index a7b53cbf81..bf57ebd8c2 100644
--- a/agent_test.go
+++ b/agent_test.go
@@ -7,6 +7,7 @@
 package main
 
 import (
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path"
@@ -25,8 +26,17 @@ import (
 const (
 	testExecID      = "testExecID"
 	testContainerID = "testContainerID"
+	testFileMode    = os.FileMode(0640)
 )
 
+func createFileWithPerms(file, contents string, perms os.FileMode) error {
+	return ioutil.WriteFile(file, []byte(contents), perms)
+}
+
+func createFile(file, contents string) error {
+	return createFileWithPerms(file, contents, testFileMode)
+}
+
 func skipUnlessRoot(t *testing.T) {
 	if os.Getuid() != 0 {
 		t.Skip("Test disabled as requires root user")
@@ -487,3 +497,127 @@ func TestAddGuestHooks(t *testing.T) {
 	assert.True(len(spec.Hooks.Poststop) == 1)
 	assert.True(strings.Contains(spec.Hooks.Poststop[0].Path, "executable"))
 }
+
+func TestGetMemory(t *testing.T) {
+	assert := assert.New(t)
+
+	dir, err := ioutil.TempDir("", "")
+	assert.NoError(err)
+	defer os.RemoveAll(dir)
+
+	file := filepath.Join(dir, "meminfo")
+
+	savedMeminfo := meminfo
+	defer func() {
+		meminfo = savedMeminfo
+	}()
+
+	// Override the file
+	meminfo = file
+
+	type testData struct {
+		contents       string
+		expectedResult string
+		createFile     bool
+		expectError    bool
+	}
+
+	memKB := 13
+	memKBStr := fmt.Sprintf("%d", memKB)
+
+	entry := fmt.Sprintf("MemTotal:      %d\n", memKB)
+	tooManyFieldsEntry := fmt.Sprintf("MemTotal: foo:     %d\n", memKB)
+	noNumericEntry := fmt.Sprintf("MemTotal: \n")
+
+	data := []testData{
+		{
+			"",
+			"",
+			false,
+			true,
+		},
+		{
+			"",
+			"",
+			true,
+			true,
+		},
+		{
+			"hello",
+			"",
+			true,
+			true,
+		},
+		{
+			"hello:",
+			"",
+			true,
+			true,
+		},
+		{
+			"hello: world",
+			"",
+			true,
+			true,
+		},
+		{
+			"hello: world:",
+			"",
+			true,
+			true,
+		},
+		{
+			"MemTotal:      ",
+			"",
+			true,
+			true,
+		},
+		{
+			"MemTotal:",
+			"",
+			true,
+			true,
+		},
+		{
+			tooManyFieldsEntry,
+			"",
+			true,
+			true,
+		},
+		{
+			noNumericEntry,
+			"",
+			true,
+			true,
+		},
+		{
+			entry,
+			memKBStr,
+			true,
+			false,
+		},
+	}
+
+	for i, d := range data {
+		msg := fmt.Sprintf("test[%d]: %+v\n", i, d)
+
+		if d.createFile {
+			err := createFile(file, d.contents)
+			assert.NoError(err, msg)
+			defer os.Remove(file)
+		} else {
+			// Ensure it does not exist
+			os.Remove(file)
+		}
+
+		mem, err := getMemory()
+		if d.expectError {
+			assert.Error(err, msg)
+			continue
+		}
+
+		assert.NoError(err, msg)
+
+		assert.Equal(d.expectedResult, mem, msg)
+	}
+}