[llvm] [llvm-lit] Add redirection handling for `env` command without args and write a lit test to check behavior with lit internal shell (PR #106629)

[via llvm-commits]

https://github.com/Harini0924 updated https://github.com/llvm/llvm-project/pull/106629

>From 5f5af8c7ff88b725452b4237dde643366953dbef Mon Sep 17 00:00:00 2001
From: Harini <harinidonthula at google.com>
Date: Thu, 29 Aug 2024 20:47:18 +0000
Subject: [PATCH] [llvm-lit] Update TestRunner.py for Redirection Handling and
 a test in lit

Added Redirection Handling: Enhanced TestRunner.py to correctly handle cases
where the env command is run without arguments. It now checks if there are
no arguments and, when redirection is specified, writes the environment
variables directly to the redirected file or a temporary file. Test for
Redirection Behavior: Created a new test in lit's internal shell to verify that
when the env command is executedwithout arguments and with redirection, it
properly outputs the environment variables to the specified file.
---
 llvm/utils/lit/lit/TestRunner.py              | 22 +++++++++++++++----
 .../shtest-env-positive/env-temp-redirect.txt |  7 ++++++
 llvm/utils/lit/tests/shtest-env-positive.py   | 12 +++++++---
 3 files changed, 34 insertions(+), 7 deletions(-)
 create mode 100755 llvm/utils/lit/tests/Inputs/shtest-env-positive/env-temp-redirect.txt

diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index 19f35fc7e212f3..2b3c6b8b9b750f 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -746,11 +746,25 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
                     env_str = "\n".join(
                         f"{key}={value}" for key, value in sorted(cmd_shenv.env.items())
                     )
-                    results.append(
-                        ShellCommandResult(
-                            j, env_str, "", 0, timeoutHelper.timeoutReached(), []
-                        )
+                    # Process redirections.
+                    stdin, stdout, stderr = processRedirects(
+                        j, default_stdin, cmd_shenv, opened_files
                     )
+                    if stdout != default_stdin:
+                        # Write directly to the redirected file (stdout).
+                        stdout.write(env_str)
+                        results.append(
+                            ShellCommandResult(
+                                j, "", "", 0, timeoutHelper.timeoutReached(), []
+                            )
+                        )
+                    else:
+                        # Capture the output for cases without redirection.
+                        results.append(
+                            ShellCommandResult(
+                                j, env_str, "", 0, timeoutHelper.timeoutReached(), []
+                            )
+                        )
                     return 0
             elif args[0] == "not":
                 not_args.append(args.pop(0))
diff --git a/llvm/utils/lit/tests/Inputs/shtest-env-positive/env-temp-redirect.txt b/llvm/utils/lit/tests/Inputs/shtest-env-positive/env-temp-redirect.txt
new file mode 100755
index 00000000000000..b2ad755e81b0ef
--- /dev/null
+++ b/llvm/utils/lit/tests/Inputs/shtest-env-positive/env-temp-redirect.txt
@@ -0,0 +1,7 @@
+## Test the env command with output redirection to a file.
+# RUN: rm -f %t
+# RUN: env > %t
+# RUN: FileCheck %s < %t
+
+# CHECK: BAR=2
+# CHECK: FOO=1
diff --git a/llvm/utils/lit/tests/shtest-env-positive.py b/llvm/utils/lit/tests/shtest-env-positive.py
index 863fbda8c5b6dc..8f5f8385872a65 100644
--- a/llvm/utils/lit/tests/shtest-env-positive.py
+++ b/llvm/utils/lit/tests/shtest-env-positive.py
@@ -7,7 +7,7 @@
 
 ## Test the env command's successful executions.
 
-# CHECK: -- Testing: 9 tests{{.*}}
+# CHECK: -- Testing: 10 tests{{.*}}
 
 # CHECK: PASS: shtest-env :: env-args-last-is-assign.txt ({{[^)]*}})
 # CHECK: env FOO=1
@@ -47,6 +47,12 @@
 # CHECK-NOT: # error:
 # CHECK: --
 
+# CHECK: PASS: shtest-env :: env-temp-redirect.txt ({{[^)]*}})
+# CHECK: env {{.*}}/env-temp-redirect.txt.tmp
+# CHECK: # executed command: env
+# CHECK-NOT: # error:
+# CHECK: --
+
 # CHECK: PASS: shtest-env :: env-u.txt ({{[^)]*}})
 # CHECK: env -u FOO | {{.*}}
 # CHECK: # executed command: env -u FOO
@@ -65,6 +71,6 @@
 # CHECK-NOT: # error:
 # CHECK: --
 
-# CHECK: Total Discovered Tests: 9
-# CHECK: Passed: 9 {{\([0-9]*\.[0-9]*%\)}}
+# CHECK: Total Discovered Tests: 10
+# CHECK: Passed: 10 {{\([0-9]*\.[0-9]*%\)}}
 # CHECK-NOT: {{.}}