Fix forwarding DOTNET_ROOT

src/Cli/dotnet/Commands/Test/TestCommand.cs

@@ -144,7 +144,9 @@
             convertedArgs.Add($"--testSessionCorrelationId:{testSessionCorrelationId}");
         }
 
-        int exitCode = new VSTestForwardingApp(convertedArgs).Execute();
+        string archArg = parseResult.ForwardedOptionValues<IReadOnlyCollection<string>>(TestCommandParser.GetCommand(), "--arch")?.SingleOrDefault() ?? null;
+
+        int exitCode = new VSTestForwardingApp(convertedArgs, archArg).Execute();
 
         // We run post processing also if execution is failed for possible partial successful result to post process.
         exitCode |= RunArtifactPostProcessingIfNeeded(testSessionCorrelationId, parseResult, FeatureFlag.Instance);
@@ -235,13 +237,14 @@
             }
         }
 
-        // Set DOTNET_PATH if it isn't already set in the environment as it is required
+        // Set DOTNET_ROOT if it isn't already set in the environment as it is required
         // by the testhost which uses the apphost feature (Windows only).
-        (bool hasRootVariable, string rootVariableName, string rootValue) = VSTestForwardingApp.GetRootVariable();
-        if (!hasRootVariable)
+        string archArg = result.ForwardedOptionValues<IReadOnlyCollection<string>>(TestCommandParser.GetCommand(), "--arch")?.SingleOrDefault() ?? null;
+        (bool setRootVariable, string rootVariableName, string rootValue) = VSTestForwardingApp.GetRootVariable(archArg);
+        if (!setRootVariable)
         {
             testCommand.EnvironmentVariable(rootVariableName, rootValue);
             VSTestTrace.SafeWriteTrace(() => $"Root variable set {rootVariableName}:{rootValue}");
         }
 
         VSTestTrace.SafeWriteTrace(() => $"Starting test using MSBuild with arguments '{testCommand.GetArgumentsToMSBuild()}' custom MSBuild path '{msbuildPath}' norestore '{noRestore}'");
@@ -272,10 +275,12 @@
             artifactsPostProcessArgs.Add($"--diag:{parseResult.GetValue(TestCommandParser.DiagOption)}");
         }
 
+        string archArg = parseResult.ForwardedOptionValues<IReadOnlyCollection<string>>(TestCommandParser.GetCommand(), "--arch")?.SingleOrDefault() ?? null;
+
         try
         {
-            return new VSTestForwardingApp(artifactsPostProcessArgs).Execute();
+            return new VSTestForwardingApp(artifactsPostProcessArgs, archArg).Execute();
         }
         finally
         {
             if (Directory.Exists(expectedArtifactDirectory))

src/Cli/dotnet/Commands/Test/VSTestForwardingApp.cs

@@ -3,6 +3,7 @@
 
 #nullable disable
 
+using System.Reflection;
 using Microsoft.DotNet.Cli.Utils;
 
 namespace Microsoft.DotNet.Cli.Commands.Test;
@@ -11,11 +12,11 @@ public class VSTestForwardingApp : ForwardingApp
 {
     private const string VstestAppName = "vstest.console.dll";
 
-    public VSTestForwardingApp(IEnumerable<string> argsToForward)
+    public VSTestForwardingApp(IEnumerable<string> argsToForward, string targetArchitecture)
         : base(GetVSTestExePath(), argsToForward)
     {
-        (bool hasRootVariable, string rootVariableName, string rootValue) = GetRootVariable();
-        if (!hasRootVariable)
+        (bool setRootVariable, string rootVariableName, string rootValue) = GetRootVariable(targetArchitecture);
+        if (!setRootVariable)
         {
             WithEnvironmentVariable(rootVariableName, rootValue);
             VSTestTrace.SafeWriteTrace(() => $"Root variable set {rootVariableName}:{rootValue}");
@@ -38,14 +39,57 @@ private static string GetVSTestExePath()
         return Path.Combine(AppContext.BaseDirectory, VstestAppName);
     }
 
-    internal static (bool hasRootVariable, string rootVariableName, string rootValue) GetRootVariable()
+    internal static (bool setRootVariable, string rootVariableName, string rootValue) GetRootVariable(string targetArchitecture)
     {
-        string rootVariableName = Environment.Is64BitProcess ? "DOTNET_ROOT" : "DOTNET_ROOT(x86)";
-        bool hasRootVariable = Environment.GetEnvironmentVariable(rootVariableName) != null;
-        string rootValue = hasRootVariable ? null : Path.GetDirectoryName(new Muxer().MuxerPath);
+        string[] rootVariables = [];
+        var processArchitecture = RuntimeInformation.ProcessArchitecture;
+        if (string.IsNullOrWhiteSpace(targetArchitecture) || processArchitecture != Enum.Parse<Architecture>(targetArchitecture, ignoreCase: true))
+        {
+            // User specified the --arch parameter but it is different from current process architecture, so we won't set anything
+            // to not break child processes by setting DOTNET_ROOT on them;
+            return (setRootVariable: false, null, null);
+        }
+
+        // Get variables to pick up from the current environment in the order in which they are inspected by the process.
+        rootVariables = GetRootVariablesToInspect(processArchitecture);
+
+        string rootPath = null;
+        foreach (var variable in rootVariables)
+        {
+            rootPath = Environment.GetEnvironmentVariable(variable);
+            if (!string.IsNullOrWhiteSpace(rootPath))
+            {
+                break;
+            }
+        }
+
+        if (!string.IsNullOrWhiteSpace(rootPath))
+        {
+            // Root path for this process was already set externally don't do anything.
+            return (setRootVariable: false, null, null);
+        }
+
+        // VSTest only accepts the two variants below, and the apphost does fallback to DOTNET_ROOT in all architectures, so we pass the
+        // architecture non-specific env variable.
+        string rootVariableName = Environment.Is64BitProcess ? "VSTEST_WINAPPHOST_DOTNET_ROOT" : "VSTEST_WINAPPHOST_DOTNET_ROOT(x86)";
 
         // We rename env variable to support --arch switch that relies on DOTNET_ROOT/DOTNET_ROOT(x86)
         // We provide VSTEST_WINAPPHOST_ only in case of testhost*.exe removing VSTEST_WINAPPHOST_ prefix and passing as env vars.
-        return (hasRootVariable, $"VSTEST_WINAPPHOST_{rootVariableName}", rootValue);
+        return (setRootVariable: true, rootVariableName, rootPath);
+
+        static string[] GetRootVariablesToInspect(Architecture processArchitecture)
+        {
+            switch (processArchitecture)
+            {
+                case Architecture.X86:
+                    return ["DOTNET_ROOT_X86", "DOTNET_ROOT(x86)"];
+
+                case Architecture.X64:
+                    return ["DOTNET_ROOT_X64", "DOTNET_ROOT"];
+
+                default:
+                    return [$"DOTNET_ROOT_{processArchitecture.ToString().ToUpperInvariant()}"];
+            }
+        }
     }
 }

src/Cli/dotnet/Commands/VSTest/VSTestCommand.cs

@@ -29,7 +29,8 @@ public static int Run(ParseResult parseResult)
             args.Add($"--testSessionCorrelationId:{testSessionCorrelationId}");
         }
 
-        VSTestForwardingApp vsTestforwardingApp = new(args);
+        string archArg = parseResult.ForwardedOptionValues<IReadOnlyCollection<string>>(TestCommandParser.GetCommand(), "--arch")?.SingleOrDefault() ?? null;
+        VSTestForwardingApp vsTestforwardingApp = new(args, archArg);
 
         int exitCode = vsTestforwardingApp.Execute();
 

src/Cli/dotnet/Parser.cs

@@ -341,7 +341,7 @@ public override void Write(HelpContext context)
             }
             else if (command.Name.Equals(VSTestCommandParser.GetCommand().Name))
             {
-                new VSTestForwardingApp(helpArgs).Execute();
+                new VSTestForwardingApp(helpArgs, targetArchitecture: null).Execute();
             }
             else if (command.Name.Equals(FormatCommandParser.GetCommand().Name))
             {

test/dotnet.Tests/CommandTests/MSBuild/GivenDotnetVsTestForwardingApp.cs

@@ -10,7 +10,7 @@ public class GivenDotnetVsTestForwardingApp
         [Fact]
         public void ItRunsVsTestApp()
         {
-            new VSTestForwardingApp(new string[0])
+            new VSTestForwardingApp([], targetArchitecture: null)
                 .GetProcessStartInfo().Arguments.Should().EndWith("vstest.console.dll");
         }
 
@@ -23,7 +23,7 @@ public void ItCanUseEnvironmentVariableToForceCustomPathToVsTestApp()
             try
             {
                 Environment.SetEnvironmentVariable(vsTestConsolePath, dummyPath);
-                new VSTestForwardingApp(new string[0])
+                new VSTestForwardingApp([], targetArchitecture: null)
                     .GetProcessStartInfo().Arguments.Should().EndWith("vstest.custom.console.dll");
             }
             finally