From d43ad333b8847b596b3e1205e85635f0c513aa08 Mon Sep 17 00:00:00 2001
From: Masood Malekghassemi <soltanmm@users.noreply.github.com>
Date: Thu, 17 Sep 2015 14:50:36 -0700
Subject: [PATCH] Make load_tests protocol tests run via py.test

---
 src/python/grpcio_test/grpc_test/conftest.py | 60 ++++++++++++++++++++
 tools/run_tests/run_python.sh                | 10 ----
 2 files changed, 60 insertions(+), 10 deletions(-)
 create mode 100644 src/python/grpcio_test/grpc_test/conftest.py

diff --git a/src/python/grpcio_test/grpc_test/conftest.py b/src/python/grpcio_test/grpc_test/conftest.py
new file mode 100644
index 00000000000..357320ec64b
--- /dev/null
+++ b/src/python/grpcio_test/grpc_test/conftest.py
@@ -0,0 +1,60 @@
+import types
+import unittest
+
+import pytest
+
+
+class LoadTestsSuiteCollector(pytest.Collector):
+
+  def __init__(self, name, parent, suite):
+    super(LoadTestsSuiteCollector, self).__init__(name, parent=parent)
+    self.suite = suite
+    self.obj = suite
+
+  def collect(self):
+    collected = []
+    for case in self.suite:
+      if isinstance(case, unittest.TestCase):
+        collected.append(LoadTestsCase(case.id(), self, case))
+      elif isinstance(case, unittest.TestSuite):
+        collected.append(
+            LoadTestsSuiteCollector('suite_child_of_mine', self, case))
+    return collected
+
+  def reportinfo(self):
+    return str(self.suite)
+
+
+class LoadTestsCase(pytest.Function):
+
+  def __init__(self, name, parent, item):
+    super(LoadTestsCase, self).__init__(name, parent, callobj=self._item_run)
+    self.item = item
+
+  def _item_run(self):
+    result = unittest.TestResult()
+    self.item(result)
+    if result.failures:
+      test_method, trace = result.failures[0]
+      pytest.fail(trace, False)
+    elif result.errors:
+      test_method, trace = result.errors[0]
+      pytest.fail(trace, False)
+    elif result.skipped:
+      test_method, reason = result.skipped[0]
+      pytest.skip(reason)
+
+
+def pytest_pycollect_makeitem(collector, name, obj):
+  if name == 'load_tests' and isinstance(obj, types.FunctionType):
+    suite = unittest.TestSuite()
+    loader = unittest.TestLoader()
+    pattern = '*'
+    try:
+      # Check that the 'load_tests' object is actually a callable that actually
+      # accepts the arguments expected for the load_tests protocol.
+      suite = obj(loader, suite, pattern)
+    except Exception as e:
+      return None
+    else:
+      return LoadTestsSuiteCollector(name, collector, suite)
diff --git a/tools/run_tests/run_python.sh b/tools/run_tests/run_python.sh
index 977b02fd949..e2135be04c5 100755
--- a/tools/run_tests/run_python.sh
+++ b/tools/run_tests/run_python.sh
@@ -40,14 +40,4 @@ export DYLD_LIBRARY_PATH=$ROOT/libs/$CONFIG
 export PATH=$ROOT/bins/$CONFIG:$ROOT/bins/$CONFIG/protobuf:$PATH
 source "python"$PYVER"_virtual_environment"/bin/activate
 
-# TODO(atash): These tests don't currently run under py.test and thus don't
-# appear under the coverage report. Find a way to get these tests to work with
-# py.test (or find another tool or *something*) that's acceptable to the rest of
-# the team...
-"python"$PYVER -m grpc_test._core_over_links_base_interface_test
-"python"$PYVER -m grpc_test._crust_over_core_over_links_face_interface_test
-"python"$PYVER -m grpc_test.beta._face_interface_test
-"python"$PYVER -m grpc_test.framework._crust_over_core_face_interface_test
-"python"$PYVER -m grpc_test.framework.core._base_interface_test
-
 "python"$PYVER $GRPCIO_TEST/setup.py test -a "-n8 --cov=grpc --junitxml=./report.xml --timeout=300"