[APP-8439] Ensure VIAM_MODULE_ROOT is set for local modules (viamrobotics#5246)

Co-authored-by: Daniel Botros <[email protected]>

Author: danielbotros

config/module.go

@@ -180,8 +180,8 @@ func (m Module) SyntheticPackage() (PackageConfig, error) {
 	return PackageConfig{Name: name, Package: name, Type: PackageTypeModule, Version: m.LocalVersion}, nil
 }
 
-// exeDir returns the parent directory for the unpacked module.
-func (m Module) exeDir(packagesDir string) (string, error) {
+// ExeDir returns the parent directory for the unpacked module.
+func (m Module) ExeDir(packagesDir string) (string, error) {
 	if !m.NeedsSyntheticPackage() {
 		return filepath.Dir(m.ExePath), nil
 	}
@@ -218,7 +218,7 @@ func (m Module) EvaluateExePath(packagesDir string) (string, error) {
 	if !filepath.IsAbs(path) {
 		return "", fmt.Errorf("expected ExePath to be absolute path, got %s", path)
 	}
-	exeDir, err := m.exeDir(packagesDir)
+	exeDir, err := m.ExeDir(packagesDir)
 	if err != nil {
 		return "", err
 	}
@@ -270,7 +270,7 @@ func (m *Module) FirstRun(
 ) error {
 	logger = logger.Sublogger("first_run").WithFields("module", m.Name)
 
-	unpackedModDir, err := m.exeDir(localPackagesDir)
+	unpackedModDir, err := m.ExeDir(localPackagesDir)
 	if err != nil {
 		return err
 	}

config/module_test.go

@@ -40,7 +40,7 @@ func TestInternalMeta(t *testing.T) {
 		})
 
 		t.Run("meta.json present", func(t *testing.T) {
-			exeDir, err := mod.exeDir(packagesDir)
+			exeDir, err := mod.ExeDir(packagesDir)
 			test.That(t, err, test.ShouldBeNil)
 			manifest := JSONManifest{Entrypoint: "entry"}
 			testWriteJSON(t, filepath.Join(exeDir, "meta.json"), manifest)
@@ -91,7 +91,7 @@ func TestSyntheticModule(t *testing.T) {
 	})
 
 	t.Run("syntheticPackageExeDir", func(t *testing.T) {
-		dir, err := modNeedsSynthetic.exeDir(tmp)
+		dir, err := modNeedsSynthetic.ExeDir(tmp)
 		test.That(t, err, test.ShouldBeNil)
 		test.That(t, dir, test.ShouldEqual, filepath.Join(tmp, "data/module/synthetic--"))
 	})

module/modmanager/manager.go

@@ -1024,13 +1024,14 @@ func (mgr *Manager) FirstRun(ctx context.Context, conf config.Module) error {
 			return err
 		}
 	}
-	env := getFullEnvironment(conf, dataDir, mgr.viamHomeDir)
+	env := getFullEnvironment(conf, pkgsDir, dataDir, mgr.viamHomeDir)
 
 	return conf.FirstRun(ctx, pkgsDir, dataDir, env, mgr.logger)
 }
 
 func getFullEnvironment(
 	cfg config.Module,
+	packagesDir string,
 	dataDir string,
 	viamHomeDir string,
 ) map[string]string {
@@ -1042,8 +1043,18 @@ func getFullEnvironment(
 	if cfg.Type == config.ModuleTypeRegistry {
 		environment["VIAM_MODULE_ID"] = cfg.ModuleID
 	}
-	// Overwrite the base environment variables with the module's environment variables (if specified)
+
+	// For local modules, we set VIAM_MODULE_ROOT to the parent directory of the unpacked module.
 	// VIAM_MODULE_ROOT is filled out by app.viam.com in cloud robots.
+	if cfg.Type == config.ModuleTypeLocal {
+		moduleRoot, err := cfg.ExeDir(packagesDir)
+		// err should never not be nil since we are working with local modules
+		if err == nil {
+			environment["VIAM_MODULE_ROOT"] = moduleRoot
+		}
+	}
+
+	// Overwrite the base environment variables with the module's environment variables (if specified)
 	for key, value := range cfg.Environment {
 		environment[key] = value
 	}

module/modmanager/manager_test.go

@@ -160,17 +160,18 @@ func TestModManagerFunctions(t *testing.T) {
 
 			test.That(t, mod.process.Stop(), test.ShouldBeNil)
 
-			modEnv := mod.getFullEnvironment(viamHomeTemp)
+			modEnv := mod.getFullEnvironment(viamHomeTemp, filepath.Join(viamHomeTemp, "packages"))
 			test.That(t, modEnv["VIAM_HOME"], test.ShouldEqual, viamHomeTemp)
 			test.That(t, modEnv["VIAM_MODULE_DATA"], test.ShouldEqual, "module-data-dir")
 			test.That(t, modEnv["VIAM_MODULE_ID"], test.ShouldEqual, "new:york")
 			test.That(t, modEnv["SMART"], test.ShouldEqual, "MACHINES")
 
-			// Test that VIAM_MODULE_ID is unset for local modules
+			// Test that VIAM_MODULE_ID is unset and VIAM_MODULE_ROOT is set correctly for local modules
 			mod.cfg.Type = config.ModuleTypeLocal
-			modEnv = mod.getFullEnvironment(viamHomeTemp)
+			modEnv = mod.getFullEnvironment(viamHomeTemp, filepath.Join(viamHomeTemp, "packages"))
 			_, ok = modEnv["VIAM_MODULE_ID"]
 			test.That(t, ok, test.ShouldBeFalse)
+			test.That(t, modEnv["VIAM_MODULE_ROOT"], test.ShouldNotBeEmpty)
 
 			// Make a copy of addr and client to test that connections are properly remade
 			oldAddr := mod.addr

module/modmanager/module.go

@@ -190,7 +190,7 @@ func (m *module) startProcess(
 	if err != nil {
 		return err
 	}
-	moduleEnvironment := m.getFullEnvironment(viamHomeDir)
+	moduleEnvironment := m.getFullEnvironment(viamHomeDir, packagesDir)
 	// Prefer VIAM_MODULE_ROOT as the current working directory if present but fallback to the directory of the exepath
 	moduleWorkingDirectory, ok := moduleEnvironment["VIAM_MODULE_ROOT"]
 	if !ok {
@@ -399,8 +399,8 @@ func (m *module) cleanupAfterCrash(mgr *Manager) {
 	}
 }
 
-func (m *module) getFullEnvironment(viamHomeDir string) map[string]string {
-	return getFullEnvironment(m.cfg, m.dataDir, viamHomeDir)
+func (m *module) getFullEnvironment(viamHomeDir, packagesDir string) map[string]string {
+	return getFullEnvironment(m.cfg, packagesDir, m.dataDir, viamHomeDir)
 }
 
 func (m *module) getFTDCName() string {