diff --git a/internal/agent/upgrade.go b/internal/agent/upgrade.go
index b1e6ad5e..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"
@@ -79,30 +80,14 @@ func Upgrade(
 		}
 	}
 
-	discoveredImage := ""
-	bus.Manager.Response(events.EventVersionImage, func(p *pluggable.Plugin, r *pluggable.EventResponse) {
-		discoveredImage = r.Data
-	})
-
-	_, err := bus.Manager.Publish(events.EventVersionImage, &events.VersionImagePayload{
-		Version: version,
-	})
-	if err != nil {
-		return err
-	}
-
-	registry, err := utils.OSRelease("IMAGE_REPO")
-	if err != nil {
-		fmt.Printf("Cant find IMAGE_REPO key under /etc/os-release\n")
-		return err
-	}
-
-	img := fmt.Sprintf("%s:%s", registry, version)
-	if discoveredImage != "" {
-		img = discoveredImage
-	}
-	if source != "" {
-		img = source
+	img := source
+	var err error
+	if img == "" {
+		img, err = determineUpgradeImage(version)
+		if err != nil {
+			fmt.Println(err.Error())
+			return err
+		}
 	}
 
 	if debug {
@@ -146,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
+}