fix bug per review comment

NVIDIA · Jan 2, 2024 · d461e3f · d461e3f
1 parent 52ae8b0
commit d461e3f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/tests/test_logistic_regression.py b/python/tests/test_logistic_regression.py
@@ -1141,7 +1141,8 @@ def test_compat_one_label(
                 blor_model = blor.fit(bdf)
                 assert False, "There should be a java exception"
             except Py4JJavaError as e:
-                assert spark_v34_msg or spark_v33_msg in e.java_exception.getMessage()
+                java_msg = e.java_exception.getMessage()
+                assert spark_v34_msg in java_msg or spark_v33_msg in java_msg
 
             return