Combine rhel like distros logic

Author: Ricky Cousins

distros/rhel/rocky.go distros/rhel/rhel.godistros/rhel/rocky.go renamed to distros/rhel/rhel.go

@@ -7,9 +7,9 @@ import (
 	utils "github.com/Puppet-Finland/updates-exporter/distros"
 )
 
-type Rocky struct{}
+type Rhel struct{}
 
-func (Rocky) GetSecurityUpdates() int {
+func (Rhel) GetSecurityUpdates() int {
 	cmd := exec.Command("sh", "-c", "dnf updateinfo list --sec-severity=Critical --sec-severity=Important --all | wc -l")
 	out, err := cmd.Output()
 	if err != nil {
@@ -19,7 +19,7 @@ func (Rocky) GetSecurityUpdates() int {
 	return utils.ParseUpdateCount(string(out))
 }
 
-func (Rocky) GetTotalUpdates() int {
+func (Rhel) GetTotalUpdates() int {
 	cmd := exec.Command("sh", "-c", "dnf updateinfo list --all | wc -l")
 	out, err := cmd.Output()
 	if err != nil {
@@ -29,8 +29,11 @@ func (Rocky) GetTotalUpdates() int {
 	return utils.ParseUpdateCount(string(out))
 }
 
-func (Rocky) GetRebootRequired() bool {
+func (Rhel) GetRebootRequired() bool {
 	cmd := exec.Command("needs-restarting", "-r")
-	err := cmd.Run()
-	return err != nil
+	if err := cmd.Run(); err != nil {
+		log.Printf("Error %v", err)
+		return false
+	}
+	return true
 }

distros/utils.go

@@ -9,7 +9,12 @@ import (
 	"strings"
 )
 
-var rockyReleaseFile = "/etc/rocky-release"
+var rhelReleases = []string{
+	"/etc/rocky-release",
+	"/etc/almalinux-release",
+	"/etc/redhat-release",
+}
+
 var osReleaseFile = "/etc/os-release"
 
 func ParseUpdateCount(out string) int {
@@ -22,8 +27,10 @@ func GetLinuxDistro() string {
 		return "unknown"
 	}
 
-	if _, err := os.Stat(rockyReleaseFile); err == nil {
-		return "rocky"
+	for _, releaseFile := range rhelReleases {
+		if _, err := os.Stat(releaseFile); err == nil {
+			return "rhel"
+		}
 	}
 
 	out, err := exec.Command("sh", "-c", fmt.Sprintf("cat %s", osReleaseFile)).Output()
@@ -35,8 +42,6 @@ func GetLinuxDistro() string {
 	switch {
 	case strings.Contains(s, "ubuntu"):
 		return "ubuntu"
-	case strings.Contains(s, "alma"):
-		return "alma"
 	case strings.Contains(s, "rhel"), strings.Contains(s, "centos"), strings.Contains(s, "fedora"):
 		return "rhel"
 	default:

distros/utils_test.go

@@ -6,14 +6,7 @@ import (
 )
 
 func getRelease(r string, t *testing.T) {
-	osReleaseFile = "/tmp/os-release"
 
-	os.WriteFile(osReleaseFile, []byte(r), 0644)
-	got := GetLinuxDistro()
-	expected := r
-	if got != expected {
-		t.Errorf("Expected %s, got %s", expected, got)
-	}
 }
 
 func TestParseUpdateCount(t *testing.T) {
@@ -26,17 +19,31 @@ func TestParseUpdateCount(t *testing.T) {
 }
 
 func TestGetDistros(t *testing.T) {
-	getRelease("alma", t)
-	getRelease("rhel", t)
-	getRelease("ubuntu", t)
-}
-
-func TestGetDistrosRocky(t *testing.T) {
-	rockyReleaseFile = "/tmp/rocky-release"
+	osReleaseFile = "/tmp/os-release"
+	rhelReleases = []string{
+		"/tmp/rocky-release",
+	}
+	os.Remove(rhelReleases[0])
 
-	os.WriteFile(rockyReleaseFile, []byte("rocky"), 0644)
+	os.WriteFile(osReleaseFile, []byte("ubuntu"), 0644)
 	got := GetLinuxDistro()
-	expected := "rocky"
+	expected := "ubuntu"
+	if got != expected {
+		t.Errorf("Expected %s, got %s", expected, got)
+	}
+
+	os.WriteFile(osReleaseFile, []byte("fedora"), 0644)
+	got = GetLinuxDistro()
+	expected = "rhel"
+	if got != expected {
+		t.Errorf("Expected %s, got %s", expected, got)
+	}
+
+	if err := os.WriteFile(rhelReleases[0], []byte("test"), 0644); err != nil {
+		t.Errorf("Error creating %s", rhelReleases[0])
+	}
+
+	got = GetLinuxDistro()
 	if got != expected {
 		t.Errorf("Expected %s, got %s", expected, got)
 	}

main.go

@@ -39,8 +39,8 @@ func getDistro() distros.Distro {
 	switch distros.GetLinuxDistro() {
 	case "ubuntu":
 		return ubuntu.Ubuntu{}
-	case "rocky", "alma":
-		return rhel.Rocky{}
+	case "rhel":
+		return rhel.Rhel{}
 	default:
 		return nil
 	}