From e3308b105e6f9c1244ae8339cf76080ab3b7b837 Mon Sep 17 00:00:00 2001
From: Shubham <shbajpai@vmware.com>
Date: Thu, 5 May 2022 06:20:39 +0000
Subject: [PATCH] Add mocking and test for functions

Signed-off-by: Shubham <shbajpai@vmware.com>
---
 agent/registration/host_registrar.go      |  8 +--
 agent/registration/host_registrar_test.go | 61 +++++++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)
 create mode 100644 agent/registration/host_registrar_test.go

diff --git a/agent/registration/host_registrar.go b/agent/registration/host_registrar.go
index b934ae6ef..b3197a542 100644
--- a/agent/registration/host_registrar.go
+++ b/agent/registration/host_registrar.go
@@ -149,7 +149,7 @@ func (hr *HostRegistrar) getHostInfo() (infrastructurev1beta1.HostInfo, error) {
 	hostInfo.Architecture = runtime.GOARCH
 	hostInfo.OSName = runtime.GOOS
 
-	if distribution, err := getOperatingSystem(); err != nil {
+	if distribution, err := getOperatingSystem(ioutil.ReadFile); err != nil {
 		return hostInfo, errors.Wrap(err, "failed to get host operating system image")
 	} else {
 		hostInfo.OSImage = distribution
@@ -158,13 +158,13 @@ func (hr *HostRegistrar) getHostInfo() (infrastructurev1beta1.HostInfo, error) {
 }
 
 // getOperatingSystem gets the name of the current operating system image.
-func getOperatingSystem() (string, error) {
+func getOperatingSystem(f func(string) ([]byte, error)) (string, error) {
 	rex := regexp.MustCompile("(PRETTY_NAME)=(.*)")
 
-	bytes, err := ioutil.ReadFile("/etc/os-release")
+	bytes, err := f("/etc/os-release")
 	if err != nil && os.IsNotExist(err) {
 		// /usr/lib/os-release in stateless systems like Clear Linux
-		bytes, err = ioutil.ReadFile("/usr/lib/os-release")
+		bytes, err = f("/usr/lib/os-release")
 	}
 	if err != nil {
 		return "", fmt.Errorf("error opening file : %v", err)
diff --git a/agent/registration/host_registrar_test.go b/agent/registration/host_registrar_test.go
new file mode 100644
index 000000000..0bf246e0a
--- /dev/null
+++ b/agent/registration/host_registrar_test.go
@@ -0,0 +1,61 @@
+// Copyright 2021 VMware, Inc. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+// nolint: testpackage
+package registration
+
+import (
+	"fmt"
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+	"io/ioutil"
+	"os"
+)
+
+func getMockFile(targetOs string) ([]byte, error) {
+	out := fmt.Sprintf(`NAME="Ubuntu"
+VERSION="20.04.4 LTS (Focal Fossa)"
+ID=ubuntu
+ID_LIKE=debian
+PRETTY_NAME="%s"
+VERSION_ID="20.04"
+HOME_URL="https://www.ubuntu.com/"
+SUPPORT_URL="https://help.ubuntu.com/"
+BUG_REPORT_URL="https://bugs.launchpad.net/ubuntu/"
+PRIVACY_POLICY_URL="https://www.ubuntu.com/legal/terms-and-policies/privacy-policy"
+VERSION_CODENAME=focal
+UBUNTU_CODENAME=focal`, targetOs)
+	return []byte(out), nil
+}
+
+var _ = Describe("Host Registrar Tests", func() {
+
+	var (
+		targetOs = "Ubuntu 20.04.4 LTS"
+	)
+
+	Context("When the OS is detected", func() {
+		It("Should return the operating system for os following /etc/os-release", func() {
+			detectedOS, err := getOperatingSystem(func(string) ([]byte, error) { return getMockFile(targetOs) })
+			Expect(err).ShouldNot(HaveOccurred())
+			Expect(detectedOS).To(Equal("Ubuntu 20.04.4 LTS"))
+		})
+
+		It("Should return the operating system for os following /usr/lib/os-release", func() {
+			targetOs = "Clear Linux Initramfs"
+			detectedOS, err := getOperatingSystem(func(releaseFile string) ([]byte, error) {
+				if releaseFile == "/etc/os-release" {
+					return nil, os.ErrNotExist
+				}
+				return getMockFile(targetOs)
+			})
+			Expect(err).ShouldNot(HaveOccurred())
+			Expect(detectedOS).To(Equal("Clear Linux Initramfs"))
+		})
+
+		It("Should not error with real hostnamectl", func() {
+			_, err := getOperatingSystem(ioutil.ReadFile)
+			Expect(err).ShouldNot(HaveOccurred())
+		})
+	})
+})