From 1edb7d2cc0b61e1b1a8bca61a6eba69f0911df86 Mon Sep 17 00:00:00 2001 From: Krzysztof Parzyszek Date: Wed, 14 Jun 2023 14:58:46 -0700 Subject: [PATCH] [Test] Improve check for TVMError exception in test_cast The original code has ``` with pytest.raises(tvm.error.TVMError): s.astype("int") assert "Can't cast a handle to other types" in str(e) ``` The `astype` is expected to throw a TVMError exception, but if it does, the following assertion never gets executed. Add try/except to the test to verify that the expected exception is thrown. --- tests/python/unittest/test_tir_nodes.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/python/unittest/test_tir_nodes.py b/tests/python/unittest/test_tir_nodes.py index 7826b5960bb3..49816778f11f 100644 --- a/tests/python/unittest/test_tir_nodes.py +++ b/tests/python/unittest/test_tir_nodes.py @@ -118,9 +118,12 @@ def test_cast(): assert z.lanes == 4 s = tvm.tir.StringImm("s") - with pytest.raises(tvm.error.TVMError) as cm: - s.astype("int") - assert "Can't cast a handle to other types" in str(cm.execption) + with pytest.raises(tvm.error.TVMError): + try: + s.astype("int") + except Exception as e: + assert "Can't cast a handle to other types" in str(e) + raise def test_attr():