From 77f3bfe16386dc2426023405d2834d67062c7a0e Mon Sep 17 00:00:00 2001
From: Cheuk Tse <cheuk@viam.com>
Date: Fri, 22 May 2026 20:05:38 -0400
Subject: [PATCH] bump + test

---
 manager.go      |  2 +-
 manager_test.go | 54 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/manager.go b/manager.go
index 2c06825f..ce6729b5 100644
--- a/manager.go
+++ b/manager.go
@@ -307,6 +307,7 @@ func (m *Manager) SubsystemUpdates(ctx context.Context) {
 					m.logger.Warn(err)
 				} else {
 					m.viamServerNeedsRestart = false
+					m.cache.MarkViamServerRunningVersion()
 				}
 				if m.viamAgentNeedsRestart {
 					m.Exit(fmt.Sprintf("A new version of %s has been installed", SubsystemName))
@@ -316,7 +317,6 @@ func (m *Manager) SubsystemUpdates(ctx context.Context) {
 				m.logger.Warnf("%s has NOT allowed a restart; will NOT restart", viamserver.SubsysName)
 			}
 		}
-		m.cache.MarkViamServerRunningVersion()
 		if err := m.viamServer.Start(ctx); err != nil {
 			m.logger.Warn(err)
 		}
diff --git a/manager_test.go b/manager_test.go
index 5abd6be1..1739340d 100644
--- a/manager_test.go
+++ b/manager_test.go
@@ -404,3 +404,57 @@ func TestSubsystemUpdatesViamServerRestart(t *testing.T) {
 		})
 	}
 }
+
+// TestSubsystemUpdatesMarksRunningVersion verifies that SubsystemUpdates calls
+// MarkViamServerRunningVersion iff the restart block successfully stopped
+// viam-server. Mark snapshots the installed CurrentVersion into runningVersion;
+// it is called before the new viam-server is started, so it reflects the
+// decision to run a version, not confirmation of a healthy start.
+func TestSubsystemUpdatesMarksRunningVersion(t *testing.T) {
+	for _, tc := range []struct {
+		name          string
+		updateReturns bool
+		stopErr       error
+		wantMark      bool
+	}{
+		{name: "no restart triggered; no mark", wantMark: false},
+		{name: "restart succeeds; marks running", updateReturns: true, wantMark: true},
+		{name: "restart stop fails; no mark", updateReturns: true, stopErr: errors.New("stop failed"), wantMark: false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			utils.MockAndCreateViamDirs(t)
+			logger := logging.NewTestLogger(t)
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			cfg := utils.DefaultConfiguration
+			fake := &fakeViamServer{
+				updateFn: func(context.Context, utils.AgentConfig) bool { return tc.updateReturns },
+				stopFn:   func(context.Context) error { return tc.stopErr },
+			}
+			m := &Manager{
+				logger:         logger,
+				cfg:            cfg,
+				globalCancel:   cancel,
+				viamServer:     fake,
+				networking:     networking.New(ctx, logger, cfg),
+				cache:          NewVersionCache(logger),
+				agentStartTime: time.Now(),
+			}
+			m.sysConfig = syscfg.New(ctx, logger, cfg, m.GetNetAppender, false, nil)
+
+			// Seed CurrentVersion so we can detect whether Mark copied it.
+			m.cache.ViamServer.CurrentVersion = "0.70.0"
+			test.That(t, m.cache.ViamServerRunningVersion(), test.ShouldEqual, "")
+
+			m.SubsystemUpdates(ctx)
+
+			want := ""
+			if tc.wantMark {
+				want = "0.70.0"
+			}
+			test.That(t, m.cache.ViamServerRunningVersion(), test.ShouldEqual, want)
+		})
+	}
+}