From 62eee9ffbfea95d69df71708f93fa792076c60ab Mon Sep 17 00:00:00 2001
From: Mauro Morales <mauro.morales@spectrocloud.com>
Date: Fri, 30 Jun 2023 11:37:50 +0300
Subject: [PATCH] Refactor

Signed-off-by: Dimitris Karakasilis <dimitris@spectrocloud.com>
---
 internal/agent/upgrade.go | 49 +++++++++++++++++++++++++--------------
 1 file changed, 31 insertions(+), 18 deletions(-)

diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index b6996438..458a7887 100644
--- a/internal/agent/upgrade.go
+++ b/internal/agent/upgrade.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+
 	"github.com/Masterminds/semver/v3"
 	events "github.com/kairos-io/kairos-sdk/bus"
 	"github.com/kairos-io/kairos-sdk/collector"
@@ -80,28 +81,13 @@ func Upgrade(
 	}
 
 	img := source
-
-	if img == "" {
-		bus.Manager.Response(events.EventVersionImage, func(p *pluggable.Plugin, r *pluggable.EventResponse) {
-			img = r.Data
-		})
-	}
-
-	_, err := bus.Manager.Publish(events.EventVersionImage, &events.VersionImagePayload{
-		Version: version,
-	})
-	if err != nil {
-		return err
-	}
-
+	var err error
 	if img == "" {
-		registry, err := utils.OSRelease("IMAGE_REPO")
+		img, err = determineUpgradeImage(version)
 		if err != nil {
-			fmt.Printf("Can't find IMAGE_REPO key under /etc/os-release\n")
+			fmt.Println(err.Error())
 			return err
 		}
-
-		img = fmt.Sprintf("%s:%s", registry, version)
 	}
 
 	if debug {
@@ -145,3 +131,30 @@ func Upgrade(
 
 	return upgradeAction.Run()
 }
+
+// determineUpgradeImage asks the provider plugin for an image or constructs
+// it using version and data from /etc/os-release
+func determineUpgradeImage(version string) (string, error) {
+	var img string
+	bus.Manager.Response(events.EventVersionImage, func(p *pluggable.Plugin, r *pluggable.EventResponse) {
+		img = r.Data
+	})
+
+	_, err := bus.Manager.Publish(events.EventVersionImage, &events.VersionImagePayload{
+		Version: version,
+	})
+	if err != nil {
+		return "", err
+	}
+
+	if img != "" {
+		return img, nil
+	}
+
+	registry, err := utils.OSRelease("IMAGE_REPO")
+	if err != nil {
+		return "", fmt.Errorf("can't find IMAGE_REPO key under /etc/os-release %w", err)
+	}
+
+	return fmt.Sprintf("%s:%s", registry, version), nil
+}