Add test rule for generic python tests. (#8940)

Fixed #8708

build_tools/benchmarks/CMakeLists.txt

@@ -6,4 +6,37 @@
 
 set(BENCHMARKS_TOOL_PYTHON_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
 
+# benchmark_tool_py_test()
+#
+# CMake function to test benchmark python tools.
+#
+# Parameters:
+# NAME: name of test
+# SRC: Test source file
+# ARGS: Command line arguments to the Python source file.
+# LABELS: Additional labels to apply to the test. The package path is added
+#     automatically.
+function(benchmark_tool_py_test)
+  cmake_parse_arguments(
+    _RULE
+    ""
+    "NAME;SRC"
+    "ARGS;LABELS"
+    ${ARGN}
+  )
+
+  iree_local_py_test(
+    NAME
+      "${_RULE_NAME}"
+    SRC
+      "${_RULE_SRC}"
+    ARGS
+      ${_RULE_ARGS}
+    LABELS
+      ${_RULE_LABELS}
+    PACKAGE_DIRS
+      ${BENCHMARKS_TOOL_PYTHON_DIR}
+  )
+endfunction()
+
 add_subdirectory(common)

build_tools/benchmarks/common/CMakeLists.txt

@@ -8,53 +8,37 @@
 # Tests
 ################################################################################
 
-if(IREE_BUILD_TESTS AND NOT ANDROID)
-
-iree_py_test(
+benchmark_tool_py_test(
   NAME
     linux_device_utils_test
-  SRCS
+  SRC
     "linux_device_utils_test.py"
 )
 
-iree_py_test(
+benchmark_tool_py_test(
   NAME
     common_arguments_test
-  SRCS
+  SRC
     "common_arguments_test.py"
 )
 
-iree_py_test(
+benchmark_tool_py_test(
   NAME
     benchmark_config_test
-  SRCS
+  SRC
     "benchmark_config_test.py"
 )
 
-iree_py_test(
+benchmark_tool_py_test(
   NAME
     benchmark_suite_test
-  SRCS
+  SRC
     "benchmark_suite_test.py"
 )
 
-iree_py_test(
+benchmark_tool_py_test(
   NAME
     benchmark_driver_test
-  SRCS
+  SRC
     "benchmark_driver_test.py"
 )
-
-# TODO(#8708): Temporary solution to fix python path for tests.
-set_property(TEST "build_tools/benchmarks/common/linux_device_utils_test"
-    APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
-set_property(TEST "build_tools/benchmarks/common/common_arguments_test"
-    APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
-set_property(TEST "build_tools/benchmarks/common/benchmark_config_test"
-    APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
-set_property(TEST "build_tools/benchmarks/common/benchmark_suite_test"
-    APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
-set_property(TEST "build_tools/benchmarks/common/benchmark_driver_test"
-    APPEND PROPERTY ENVIRONMENT "PYTHONPATH=${BENCHMARKS_TOOL_PYTHON_DIR}:$ENV{PYTHONPATH}")
-
-endif()

build_tools/cmake/iree_python.cmake

@@ -246,28 +246,29 @@ function(iree_py_library)
   endif()
 endfunction()
 
-# iree_py_test()
+# iree_local_py_test()
 #
-# CMake function to imitate Bazel's iree_py_test rule.
+# CMake function to run python test with provided python package paths.
 #
 # Parameters:
 # NAME: name of test
-# SRCS: Test source file (single file only, despite name)
+# SRC: Test source file
 # ARGS: Command line arguments to the Python source file.
 # LABELS: Additional labels to apply to the test. The package path is added
 #     automatically.
 # GENERATED_IN_BINARY_DIR: If present, indicates that the srcs have been
 #   in the CMAKE_CURRENT_BINARY_DIR.
-function(iree_py_test)
-  if(NOT IREE_BUILD_TESTS)
+# PACKAGE_DIRS: Python package paths to be added to PYTHONPATH.
+function(iree_local_py_test)
+  if(NOT IREE_BUILD_TESTS OR ANDROID)
     return()
   endif()
 
   cmake_parse_arguments(
     _RULE
     "GENERATED_IN_BINARY_DIR"
-    "NAME;SRCS"
-    "ARGS;LABELS"
+    "NAME;SRC"
+    "ARGS;LABELS;PACKAGE_DIRS"
     ${ARGN}
   )
 
@@ -290,17 +291,58 @@ function(iree_py_test)
     COMMAND
       "${IREE_SOURCE_DIR}/build_tools/cmake/run_test.${IREE_HOST_SCRIPT_EXT}"
       "${Python3_EXECUTABLE}"
-      "${CMAKE_CURRENT_SOURCE_DIR}/${_RULE_SRCS}"
+      "${CMAKE_CURRENT_SOURCE_DIR}/${_RULE_SRC}"
       ${_RULE_ARGS}
   )
 
+  list(APPEND _RULE_PACKAGE_DIRS "$ENV{PYTHONPATH}")
+  string(JOIN ":" _PYTHONPATH ${_RULE_PACKAGE_DIRS})
+
   set_property(TEST ${_NAME_PATH} PROPERTY LABELS "${_RULE_LABELS}")
   set_property(TEST ${_NAME_PATH} PROPERTY ENVIRONMENT
-      "PYTHONPATH=${IREE_BINARY_DIR}/iree/compiler/python:${IREE_BINARY_DIR}/runtime/bindings/python:$ENV{PYTHONPATH}"
+      "PYTHONPATH=${_PYTHONPATH}"
       "TEST_TMPDIR=${IREE_BINARY_DIR}/tmp/${_NAME}_test_tmpdir"
   )
   iree_add_test_environment_properties(${_NAME_PATH})
 
   # TODO(marbre): Find out how to add deps to tests.
   #               Similar to _RULE_DATA in iree_lit_test().
 endfunction()
+
+# iree_py_test()
+#
+# CMake function to imitate Bazel's iree_py_test rule.
+#
+# Parameters:
+# NAME: name of test
+# SRCS: Test source file (single file only, despite name)
+# ARGS: Command line arguments to the Python source file.
+# LABELS: Additional labels to apply to the test. The package path is added
+#     automatically.
+# GENERATED_IN_BINARY_DIR: If present, indicates that the srcs have been
+#   in the CMAKE_CURRENT_BINARY_DIR.
+function(iree_py_test)
+  cmake_parse_arguments(
+    _RULE
+    "GENERATED_IN_BINARY_DIR"
+    "NAME;SRCS"
+    "ARGS;LABELS"
+    ${ARGN}
+  )
+
+  iree_local_py_test(
+    NAME
+      "${_RULE_NAME}"
+    SRC
+      "${_RULE_SRCS}"
+    ARGS
+      ${_RULE_ARGS}
+    LABELS
+      ${_RULE_LABELS}
+    PACKAGE_DIRS
+      "${IREE_BINARY_DIR}/iree/compiler/python"
+      "${IREE_BINARY_DIR}/runtime/bindings/python"
+    GENERATED_IN_BINARY_DIR
+      "${_RULE_GENERATED_IN_BINARY_DIR}"
+  )
+endfunction()