Fix tests

Author: justinchuby

src/onnx_ir/passes/_pass_infra_test.py

@@ -40,7 +40,7 @@ class PostconditionTest(unittest.TestCase):
 
     def test_ensures_called_with_result_model_not_input_model(self):
         """Test that ensures() is called with result.model, not the input model."""
-        
+
         class TestPass(_pass_infra.PassBase):
             def __init__(self):
                 self.ensures_called_with = None
@@ -56,8 +56,7 @@ def changes_input(self) -> bool:
             def call(self, model: ir.Model) -> _pass_infra.PassResult:
                 # Create a new model (different object)
                 new_model = ir.Model(
-                    graph=ir.Graph([], [], nodes=[]), 
-                    ir_version=model.ir_version
+                    graph=ir.Graph([], [], nodes=[]), ir_version=model.ir_version
                 )
                 return _pass_infra.PassResult(model=new_model, modified=True)
 
@@ -68,14 +67,14 @@ def ensures(self, model: ir.Model) -> None:
         pass_ = TestPass()
         input_model = ir.Model(graph=ir.Graph([], [], nodes=[]), ir_version=10)
         result = pass_(input_model)
-        
+
         # Verify that ensures was called with the result model, not the input model
         self.assertIs(pass_.ensures_called_with, result.model)
         self.assertIsNot(pass_.ensures_called_with, input_model)
 
     def test_ensures_called_with_result_model_in_place_pass(self):
         """Test that ensures() is called with result.model for in-place passes."""
-        
+
         class TestInPlacePass(_pass_infra.InPlacePass):
             def __init__(self):
                 self.ensures_called_with = None
@@ -91,7 +90,7 @@ def ensures(self, model: ir.Model) -> None:
         pass_ = TestInPlacePass()
         input_model = ir.Model(graph=ir.Graph([], [], nodes=[]), ir_version=10)
         result = pass_(input_model)
-        
+
         # For in-place passes, result.model should be the same as input_model
         self.assertIs(result.model, input_model)
         # Verify that ensures was called with the result model (which is the same as input)
@@ -100,7 +99,7 @@ def ensures(self, model: ir.Model) -> None:
 
     def test_postcondition_error_raised_when_ensures_fails(self):
         """Test that PostconditionError is raised when ensures() raises an exception."""
-        
+
         class TestPass(_pass_infra.PassBase):
             @property
             def in_place(self) -> bool:
@@ -119,16 +118,17 @@ def ensures(self, model: ir.Model) -> None:
 
         pass_ = TestPass()
         model = ir.Model(graph=ir.Graph([], [], nodes=[]), ir_version=10)
-        
-        with self.assertRaises(PostconditionError) as cm:
+
+        with self.assertRaisesRegex(
+            ir.passes.PostconditionError, "Post-condition for pass 'TestPass' failed"
+        ) as cm:
             pass_(model)
-        
-        self.assertIn("Post-condition for pass 'TestPass' failed", str(cm.exception))
+
         self.assertIsInstance(cm.exception.__cause__, ValueError)
 
     def test_postcondition_error_raised_when_ensures_raises_postcondition_error(self):
         """Test that PostconditionError is re-raised when ensures() raises PostconditionError."""
-        
+
         class TestPass(_pass_infra.PassBase):
             @property
             def in_place(self) -> bool:
@@ -143,35 +143,33 @@ def call(self, model: ir.Model) -> _pass_infra.PassResult:
 
             def ensures(self, model: ir.Model) -> None:
                 # Directly raise PostconditionError
-                raise PostconditionError("Direct postcondition error")
+                raise ir.passes.PostconditionError("Direct postcondition error")
 
         pass_ = TestPass()
         model = ir.Model(graph=ir.Graph([], [], nodes=[]), ir_version=10)
-        
-        with self.assertRaises(PostconditionError) as cm:
+
+        with self.assertRaisesRegex(
+            ir.passes.PostconditionError, "Direct postcondition error"
+        ):
             pass_(model)
-        
-        self.assertEqual(str(cm.exception), "Direct postcondition error")
 
     def test_ensures_receives_correct_model_when_pass_modifies_model(self):
         """Test a more complex scenario where the pass modifies the model structure."""
-        
+
         class ModelModifyingPass(_pass_infra.FunctionalPass):
             def __init__(self):
                 self.ensures_model_graph_nodes_count = None
 
             def call(self, model: ir.Model) -> _pass_infra.PassResult:
                 # Create a new model with additional nodes
-                new_graph = ir.Graph([], [], nodes=[])
-                # Add a dummy node to the graph to make it different
-                new_node = ir.Node(
-                    domain="", 
-                    op_type="Identity", 
-                    inputs=[], 
-                    outputs=[ir.Value(name="output", shape=ir.Shape([]), dtype=ir.DataType.FLOAT)],
-                    graph=new_graph
+                new_model = ir.Model(
+                    graph=ir.Graph(
+                        [],
+                        [output := ir.Value(name="output")],
+                        nodes=[ir.node("TestOp", [], outputs=[output])],
+                    ),
+                    ir_version=model.ir_version,
                 )
-                new_model = ir.Model(graph=new_graph, ir_version=model.ir_version)
                 return _pass_infra.PassResult(model=new_model, modified=True)
 
             def ensures(self, model: ir.Model) -> None:
@@ -180,8 +178,8 @@ def ensures(self, model: ir.Model) -> None:
 
         pass_ = ModelModifyingPass()
         input_model = ir.Model(graph=ir.Graph([], [], nodes=[]), ir_version=10)
-        result = pass_(input_model)
-        
+        _ = pass_(input_model)
+
         # The ensures method should see the modified model with 1 node
         self.assertEqual(pass_.ensures_model_graph_nodes_count, 1)