From 9b027c88051797ada7350098ef190191b953014a Mon Sep 17 00:00:00 2001
From: Richard Levasseur <rlevasseur@google.com>
Date: Wed, 3 Jan 2024 14:08:24 -0800
Subject: [PATCH] python: fix bazel py_test testSmoke for Python 3.11+

The test was failing because the `b.py` file couldn't be imported because, starting with Python 3.11, the `PYTHONSAFEPATH` environment variable (set by the bootstrap startup) is respected. This setting inhibits the default Python behavior of adding the main script's directory to sys.path.

To fix, use `py_library.imports` to explicitly add the necessary directory to `sys.path`.

Fixes https://github.com/bazelbuild/bazel/issues/20660

Closes #20738.

PiperOrigin-RevId: 595498775
Change-Id: I0c9521b210fe9e2c40692727fc87d41995e0968a
---
 src/test/py/bazel/py_test.py   | 5 +++--
 src/test/py/bazel/test_base.py | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/test/py/bazel/py_test.py b/src/test/py/bazel/py_test.py
index 83e45a16d103f6..a456bd877267a5 100644
--- a/src/test/py/bazel/py_test.py
+++ b/src/test/py/bazel/py_test.py
@@ -28,8 +28,9 @@ def createSimpleFiles(self):
         'a/BUILD',
         [
             'py_binary(name="a", srcs=["a.py"], deps=[":b"])',
-            'py_library(name="b", srcs=["b.py"])',
-        ])
+            'py_library(name="b", srcs=["b.py"], imports=["."])',
+        ],
+    )
 
     self.ScratchFile(
         'a/a.py',
diff --git a/src/test/py/bazel/test_base.py b/src/test/py/bazel/test_base.py
index ead5fcc4b9c753..49a4a6506fb1ff 100644
--- a/src/test/py/bazel/test_base.py
+++ b/src/test/py/bazel/test_base.py
@@ -562,7 +562,7 @@ def RunProgram(
         ]
 
         if not allow_failure:
-          self.AssertExitCode(exit_code, 0, stderr_lines)
+          self.AssertExitCode(exit_code, 0, stderr_lines, stdout_lines)
 
         return exit_code, stdout_lines, stderr_lines