fix: unit test error test_mixed_extension_setup, test_custom_raw_op_kernel_op_deprecated (#69832)

0x3878f · web-flow · commit 4cc893f6f941 · 2024-12-02T16:41:47.000+08:00
diff --git a/test/cpp_extension/test_mixed_extension_setup.py b/test/cpp_extension/test_mixed_extension_setup.py
@@ -103,10 +103,12 @@ def setUp(self):
         cur_dir = os.path.dirname(os.path.abspath(__file__))
         # install mixed custom_op and extension
         # compile, install the custom op egg into site-packages under background
+        site_dir = site.getsitepackages()[0]
         cmd = f'cd {cur_dir} && {sys.executable} mix_relu_and_extension_setup.py install'
+        if os.name != 'nt':
+            cmd += f' --install-lib={site_dir}'
         run_cmd(cmd)
 
-        site_dir = site.getsitepackages()[0]
         custom_egg_path = [
             x for x in os.listdir(site_dir) if 'mix_relu_extension' in x
         ]
diff --git a/test/deprecated/custom_op/test_custom_raw_op_kernel_op_deprecated.py b/test/deprecated/custom_op/test_custom_raw_op_kernel_op_deprecated.py
@@ -54,6 +54,9 @@ def setUpClass(cls):
         path = os.path.dirname(os.path.abspath(__file__))
         path = os.path.join(path, "custom_raw_op_kernel_op_setup.py")
         cmd = [sys.executable, path, "install", "--force"]
+        if os.name != 'nt':
+            install_lib = f"--install-lib={site.getsitepackages()[0]}"
+            cmd.append(install_lib)
         cmd = " ".join([shlex.quote(c) for c in cmd])
         os.environ['MODULE_NAME'] = MODULE_NAME
         assert os.system(cmd) == 0
