From 09fdd7a606ddde4c7fdb3c68188ae9c9e337020e Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Wed, 10 Apr 2024 10:41:43 -0700 Subject: [PATCH] fix the test Signed-off-by: Eli Uriegas --- tools/tests/test_generate_binary_build_matrix.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tools/tests/test_generate_binary_build_matrix.py b/tools/tests/test_generate_binary_build_matrix.py index 48e228cdb4..168d8a86ee 100644 --- a/tools/tests/test_generate_binary_build_matrix.py +++ b/tools/tests/test_generate_binary_build_matrix.py @@ -17,6 +17,7 @@ def matrix_compare_helper( rocm: bool, cpu: bool, reference_output_file: str, + build_python_only: bool = False, ) -> None: out = generate_build_matrix( package_type, @@ -27,6 +28,8 @@ def matrix_compare_helper( "enable" if cpu else "disable", "false", "false", + "enable" if build_python_only else "disable", + ) expected_json_filename = os.path.join(ASSETS_DIR, reference_output_file) @@ -37,6 +40,17 @@ def matrix_compare_helper( self.maxDiff = None self.assertEqual(out, expected) + def test_linux_wheel_cpu_python_only(self): + self.matrix_compare_helper( + package_type="wheel", + operating_system="linux", + cuda=True, + rocm=True, + cpu=True, + reference_output_file="build_matrix_linux_wheel_cuda.json", + build_python_only=True, + ) + def test_linux_wheel_cuda(self): self.matrix_compare_helper( package_type="wheel",