Fix #173: Correctly set up test filtering and fail fast flags for users who explicitly pass `argv=` to the `main` function.

This resolves the issue that users of `tf.test.main` can't use bazel to filter tests, as described in https://github.com/abseil/abseil-py/issues/128#issuecomment-921993484

PiperOrigin-RevId: 397883687
Change-Id: Ica703a7c3c67ccca8973340d5812e385e29f1025

6 files changed, 65 insertions, 32 deletions

diff --git a/absl/testing/BUILD b/absl/testing/BUILD
index a54997e..ecfa8a2 100644
--- a/absl/testing/BUILD
+++ b/absl/testing/BUILD
@@ -134,6 +134,7 @@ py_test(
     deps = [
         ":_bazelize_command",
         ":absltest",
+        ":parameterized",
         "//absl/logging",
     ],
 )
diff --git a/absl/testing/absltest.py b/absl/testing/absltest.py
index 0708875..9d20033 100644
--- a/absl/testing/absltest.py
+++ b/absl/testing/absltest.py
@@ -2452,7 +2452,10 @@ def _setup_sharding(custom_loader=None):
 def _run_and_get_tests_result(argv, args, kwargs, xml_test_runner_class):
   # type: (MutableSequence[Text], Sequence[Any], MutableMapping[Text, Any], Type) -> unittest.TestResult
   # pylint: enable=line-too-long
-  """Executes a set of Python unit tests and returns the result."""
+  """Same as run_tests, except it returns the result instead of existing."""
+
+  # The entry from kwargs overrides argv.
+  argv = kwargs.pop('argv', argv)
 
   # Set up test filtering if requested in environment.
   _setup_filtering(argv)
@@ -2536,7 +2539,7 @@ def _run_and_get_tests_result(argv, args, kwargs, xml_test_runner_class):
 
   # Let unittest.TestProgram.__init__ do its own argv parsing, e.g. for '-v',
   # on argv, which is sys.argv without the command-line flags.
-  kwargs.setdefault('argv', argv)
+  kwargs['argv'] = argv
 
   try:
     test_program = unittest.TestProgram(*args, **kwargs)
@@ -2564,7 +2567,9 @@ def run_tests(argv, args, kwargs):  # pylint: disable=line-too-long
 
   Args:
     argv: sys.argv with the command-line flags removed from the front, i.e. the
-      argv with which app.run() has called __main__.main.
+      argv with which app.run() has called __main__.main. It is passed to
+      unittest.TestProgram.__init__(argv=), which does its own flag parsing. It
+      is ignored if kwargs contains an argv entry.
     args: Positional arguments passed through to unittest.TestProgram.__init__.
     kwargs: Keyword arguments passed through to unittest.TestProgram.__init__.
   """
diff --git a/absl/testing/tests/absltest_fail_fast_test.py b/absl/testing/tests/absltest_fail_fast_test.py
index 83a6c76..dbda0a1 100644
--- a/absl/testing/tests/absltest_fail_fast_test.py
+++ b/absl/testing/tests/absltest_fail_fast_test.py
@@ -24,22 +24,30 @@ import subprocess
 from absl import logging
 from absl.testing import _bazelize_command
 from absl.testing import absltest
+from absl.testing import parameterized
 
 
-class TestFailFastTest(absltest.TestCase):
+@parameterized.named_parameters(
+    ('use_argv', True),
+    ('no_argv', False),
+)
+class TestFailFastTest(parameterized.TestCase):
   """Integration tests: Runs a test binary with fail fast.
 
   This is done by setting the fail fast environment variable
   """
 
   def setUp(self):
+    super().setUp()
     self._test_name = 'absl/testing/tests/absltest_fail_fast_test_helper'
 
-  def _run_fail_fast(self, fail_fast):
+  def _run_fail_fast(self, fail_fast, use_argv):
     """Runs the py_test binary in a subprocess.
 
     Args:
       fail_fast: string, the fail fast value.
+      use_argv: bool, whether the test helper should use an explicit argv=
+          parameter when calling absltest.main.
 
     Returns:
       (stdout, exit_code) tuple of (string, int).
@@ -52,6 +60,8 @@ class TestFailFastTest(absltest.TestCase):
     additional_args = []
     if fail_fast is not None:
       env['TESTBRIDGE_TEST_RUNNER_FAIL_FAST'] = fail_fast
+    if use_argv:
+      env['USE_ARGV'] = '1'
 
     proc = subprocess.Popen(
         args=([_bazelize_command.get_executable_path(self._test_name)]
@@ -65,8 +75,8 @@ class TestFailFastTest(absltest.TestCase):
     logging.info('output: %s', stdout)
     return stdout, proc.wait()
 
-  def test_no_fail_fast(self):
-    out, exit_code = self._run_fail_fast(None)
+  def test_no_fail_fast(self, use_argv):
+    out, exit_code = self._run_fail_fast(None, use_argv)
     self.assertEqual(1, exit_code)
     self.assertIn('class A test A', out)
     self.assertIn('class A test B', out)
@@ -74,8 +84,8 @@ class TestFailFastTest(absltest.TestCase):
     self.assertIn('class A test D', out)
     self.assertIn('class A test E', out)
 
-  def test_empty_fail_fast(self):
-    out, exit_code = self._run_fail_fast('')
+  def test_empty_fail_fast(self, use_argv):
+    out, exit_code = self._run_fail_fast('', use_argv)
     self.assertEqual(1, exit_code)
     self.assertIn('class A test A', out)
     self.assertIn('class A test B', out)
@@ -83,8 +93,8 @@ class TestFailFastTest(absltest.TestCase):
     self.assertIn('class A test D', out)
     self.assertIn('class A test E', out)
 
-  def test_fail_fast_1(self):
-    out, exit_code = self._run_fail_fast('1')
+  def test_fail_fast_1(self, use_argv):
+    out, exit_code = self._run_fail_fast('1', use_argv)
     self.assertEqual(1, exit_code)
     self.assertIn('class A test A', out)
     self.assertIn('class A test B', out)
@@ -92,8 +102,8 @@ class TestFailFastTest(absltest.TestCase):
     self.assertNotIn('class A test D', out)
     self.assertNotIn('class A test E', out)
 
-  def test_fail_fast_0(self):
-    out, exit_code = self._run_fail_fast('0')
+  def test_fail_fast_0(self, use_argv):
+    out, exit_code = self._run_fail_fast('0', use_argv)
     self.assertEqual(1, exit_code)
     self.assertIn('class A test A', out)
     self.assertIn('class A test B', out)
diff --git a/absl/testing/tests/absltest_fail_fast_test_helper.py b/absl/testing/tests/absltest_fail_fast_test_helper.py
index c684d12..dbef068 100644
--- a/absl/testing/tests/absltest_fail_fast_test_helper.py
+++ b/absl/testing/tests/absltest_fail_fast_test_helper.py
@@ -18,6 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import os
 import sys
 
 from absl.testing import absltest
@@ -44,4 +45,7 @@ class ClassA(absltest.TestCase):
 
 
 if __name__ == '__main__':
-  absltest.main()
+  if os.getenv('USE_ARGV', '0') == '1':
+    absltest.main(argv=sys.argv)
+  else:
+    absltest.main()
diff --git a/absl/testing/tests/absltest_filtering_test.py b/absl/testing/tests/absltest_filtering_test.py
index 469b121..8389404 100644
--- a/absl/testing/tests/absltest_filtering_test.py
+++ b/absl/testing/tests/absltest_filtering_test.py
@@ -28,8 +28,10 @@ from absl.testing import parameterized
 
 
 @parameterized.named_parameters(
-    ('as_env_variable', True),
-    ('as_commandline_args', False),
+    ('as_env_variable_use_argv', True, True),
+    ('as_env_variable_no_argv', True, False),
+    ('as_commandline_args_use_argv', False, True),
+    ('as_commandline_args_no_argv', False, True),
 )
 class TestFilteringTest(absltest.TestCase):
   """Integration tests: Runs a test binary with filtering.
@@ -39,15 +41,18 @@ class TestFilteringTest(absltest.TestCase):
   """
 
   def setUp(self):
+    super().setUp()
     self._test_name = 'absl/testing/tests/absltest_filtering_test_helper'
 
-  def _run_filtered(self, test_filter, use_env_variable):
+  def _run_filtered(self, test_filter, use_env_variable, use_argv):
     """Runs the py_test binary in a subprocess.
 
     Args:
       test_filter: string, the filter argument to use.
       use_env_variable: bool, pass the test filter as environment variable if
           True, otherwise pass as command line arguments.
+      use_argv: bool, whether the test helper should use an explicit argv=
+          parameter when calling absltest.main.
 
     Returns:
       (stdout, exit_code) tuple of (string, int).
@@ -63,6 +68,8 @@ class TestFilteringTest(absltest.TestCase):
         env['TESTBRIDGE_TEST_ONLY'] = test_filter
       elif test_filter:
         additional_args.extend(test_filter.split(' '))
+      if use_argv:
+        env['USE_ARGV'] = '1'
 
     proc = subprocess.Popen(
         args=([_bazelize_command.get_executable_path(self._test_name)]
@@ -76,34 +83,36 @@ class TestFilteringTest(absltest.TestCase):
     logging.info('output: %s', stdout)
     return stdout, proc.wait()
 
-  def test_no_filter(self, use_env_variable):
-    out, exit_code = self._run_filtered(None, use_env_variable)
+  def test_no_filter(self, use_env_variable, use_argv):
+    out, exit_code = self._run_filtered(None, use_env_variable, use_argv)
     self.assertEqual(1, exit_code)
     self.assertIn('class B test E', out)
 
-  def test_empty_filter(self, use_env_variable):
-    out, exit_code = self._run_filtered('', use_env_variable)
+  def test_empty_filter(self, use_env_variable, use_argv):
+    out, exit_code = self._run_filtered('', use_env_variable, use_argv)
     self.assertEqual(1, exit_code)
     self.assertIn('class B test E', out)
 
-  def test_class_filter(self, use_env_variable):
-    out, exit_code = self._run_filtered('ClassA', use_env_variable)
+  def test_class_filter(self, use_env_variable, use_argv):
+    out, exit_code = self._run_filtered('ClassA', use_env_variable, use_argv)
     self.assertEqual(0, exit_code)
     self.assertNotIn('class B', out)
 
-  def test_method_filter(self, use_env_variable):
-    out, exit_code = self._run_filtered('ClassB.testA', use_env_variable)
+  def test_method_filter(self, use_env_variable, use_argv):
+    out, exit_code = self._run_filtered('ClassB.testA', use_env_variable,
+                                        use_argv)
     self.assertEqual(0, exit_code)
     self.assertNotIn('class A', out)
     self.assertNotIn('class B test B', out)
 
-    out, exit_code = self._run_filtered('ClassB.testE', use_env_variable)
+    out, exit_code = self._run_filtered('ClassB.testE', use_env_variable,
+                                        use_argv)
     self.assertEqual(1, exit_code)
     self.assertNotIn('class A', out)
 
-  def test_multiple_class_and_method_filter(self, use_env_variable):
+  def test_multiple_class_and_method_filter(self, use_env_variable, use_argv):
     out, exit_code = self._run_filtered(
-        'ClassA.testA ClassA.testB ClassB.testC', use_env_variable)
+        'ClassA.testA ClassA.testB ClassB.testC', use_env_variable, use_argv)
     self.assertEqual(0, exit_code)
     self.assertIn('class A test A', out)
     self.assertIn('class A test B', out)
@@ -111,9 +120,9 @@ class TestFilteringTest(absltest.TestCase):
     self.assertIn('class B test C', out)
     self.assertNotIn('class B test A', out)
 
-  def test_not_found_filters(self, use_env_variable):
-    out, exit_code = self._run_filtered(
-        'NotExistedClass.not_existed_method', use_env_variable)
+  def test_not_found_filters(self, use_env_variable, use_argv):
+    out, exit_code = self._run_filtered('NotExistedClass.not_existed_method',
+                                        use_env_variable, use_argv)
     self.assertEqual(1, exit_code)
     self.assertIn("has no attribute 'NotExistedClass'", out)
 
diff --git a/absl/testing/tests/absltest_filtering_test_helper.py b/absl/testing/tests/absltest_filtering_test_helper.py
index aae0f92..9bc4568 100644
--- a/absl/testing/tests/absltest_filtering_test_helper.py
+++ b/absl/testing/tests/absltest_filtering_test_helper.py
@@ -18,6 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import os
 import sys
 
 from absl.testing import absltest
@@ -57,4 +58,7 @@ class ClassB(absltest.TestCase):
 
 
 if __name__ == '__main__':
-  absltest.main()
+  if os.getenv('USE_ARGV', '0') == '1':
+    absltest.main(argv=sys.argv)
+  else:
+    absltest.main()