From 94bda43d873ec96333d47ae2b86c6fa322c51db6 Mon Sep 17 00:00:00 2001
From: Zhanlue Yang <jim19930609@gmail.com>
Date: Thu, 29 Sep 2022 20:38:18 +0800
Subject: [PATCH] [build] Support executing manually-specified cpp tests for
 run_tests.py (#6206)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Issue: #

After this PR, we're able to run user-specified cpp tests via `python
tests/run_tests.py -cpp -k xxx`.

Supports C-API tests, AOT tests and regular tests.

1. C-API tests: `python3 tests/run_tests.py --cpp -k CapiAotTest.*Field`
![2022-09-29 18-11-55
的屏幕截图](https://user-images.githubusercontent.com/22334008/193004897-49a424e9-fb1d-4b7d-80eb-c0efed23dd3c.png)

2. AOT tests: `python3 tests/run_tests.py --cpp -k LlvmAotTest.Cuda.*`
![2022-09-29 18-13-30
的屏幕截图](https://user-images.githubusercontent.com/22334008/193005217-9ebeebc0-20fc-449a-9f0a-7753e2e14af5.png)

3. Regular cpp tests: `python3 tests/run_tests.py --cpp -k Scalarize.*`
![2022-09-29 18-14-25
的屏幕截图](https://user-images.githubusercontent.com/22334008/193005385-35d7a55e-0d52-46d8-a292-5cad8fece754.png)
---
 tests/run_tests.py | 37 +++++++++++++++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/tests/run_tests.py b/tests/run_tests.py
index 6321f69d341c3..f53a3773f83e2 100644
--- a/tests/run_tests.py
+++ b/tests/run_tests.py
@@ -1,8 +1,10 @@
 import argparse
 import atexit
+import copy
 import os
 import pdb
 import platform
+import re
 import shutil
 import subprocess
 import sys
@@ -61,7 +63,7 @@ def _test_cpp_aot(test_filename, build_dir, test_info):
     return exclude_tests_cmd
 
 
-def _test_cpp():
+def _test_cpp(test_keys=None):
     curr_dir = os.path.dirname(os.path.abspath(__file__))
     build_dir = os.path.join(curr_dir, '../build')
     cpp_test_filename = 'taichi_cpp_tests'
@@ -73,6 +75,37 @@ def _test_cpp():
     capi_tests_exe_path = os.path.join(build_dir, capi_test_filename)
     cpp_tests_exe_path = os.path.join(build_dir, cpp_test_filename)
 
+    # Run manually specified C++ tests only, for example:
+    # "python3 tests/run_tests.py --cpp -k Scalarize.*"
+    if test_keys:
+        # Search AOT tests
+        aot_test_cases = copy.copy(__aot_test_cases)
+        for cpp_test_name, (_, _) in __aot_test_cases.items():
+            name_match = re.match(test_keys, cpp_test_name, re.I)
+            if name_match is None:
+                aot_test_cases.pop(cpp_test_name, None)
+        if aot_test_cases:
+            _test_cpp_aot(cpp_test_filename, build_dir, aot_test_cases)
+
+        # Search CAPI tests
+        capi_aot_test_cases = copy.copy(__capi_aot_test_cases)
+        for cpp_test_name, (_, _) in __capi_aot_test_cases.items():
+            name_match = re.match(test_keys, cpp_test_name, re.I)
+            if name_match is None:
+                capi_aot_test_cases.pop(cpp_test_name, None)
+        if capi_aot_test_cases:
+            _test_cpp_aot(capi_test_filename, build_dir, capi_aot_test_cases)
+
+        # Search Cpp tests
+        _run_cpp_test(cpp_test_filename, build_dir,
+                      f"--gtest_filter={test_keys}")
+
+        _run_cpp_test(capi_test_filename, build_dir,
+                      f"--gtest_filter={test_keys}")
+
+        return
+
+    # Regular C++ tests
     if os.path.exists(capi_tests_exe_path):
         # Run C-API test cases
         exclude_tests_cmd = _test_cpp_aot(capi_test_filename, build_dir,
@@ -336,7 +369,7 @@ def size_of_dir(dir):
         os.environ['TI_OFFLINE_CACHE'] = '0'
 
     if args.cpp:
-        _test_cpp()
+        _test_cpp(args.keys)
         return
 
     for _ in range(run_count):