From 5f5af8c7ff88b725452b4237dde643366953dbef Mon Sep 17 00:00:00 2001
From: Harini <harinidonthula@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 19f35fc7e212f..2b3c6b8b9b750 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 0000000000000..b2ad755e81b0e
--- /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 863fbda8c5b6d..8f5f8385872a6 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: {{.}}