Add Tests for PR#78 (#79)

This PR adds tests for #78

### Commits:
- Add unit tests for postcondition checking bug fix

Tests verify that ensures() is called with result.model instead of input
model,
covering the fix in PassBase.__call__ where postconditions should check
the
transformed model rather than the original input model.
- Add unit tests for postcondition checking bug fix

Tests verify that ensures() is called with result.model instead of input
model,
covering the fix in PassBase.__call__ where postconditions should check
the
transformed model rather than the original input model.

Co-authored-by: codecov-ai[bot] <156709835+codecov-ai[bot]@users.noreply.github.com>

Author: codecov-ai[bot]

src/onnx_ir/passes/_pass_infra_test.py

@@ -35,5 +35,156 @@ def call(self, model: ir.Model) -> _pass_infra.PassResult:
         self.assertIs(result_1.model, result_2.model)
 
 
+class PostconditionTest(unittest.TestCase):
+    """Test that postconditions are checked on the result model, not the input model."""
+
+    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
+
+            @property
+            def in_place(self) -> bool:
+                return False  # Not in-place to create a new model
+
+            @property
+            def changes_input(self) -> bool:
+                return False
+
+            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
+                )
+                return _pass_infra.PassResult(model=new_model, modified=True)
+
+            def ensures(self, model: ir.Model) -> None:
+                # Record which model ensures was called with
+                self.ensures_called_with = model
+
+        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
+
+            def call(self, model: ir.Model) -> _pass_infra.PassResult:
+                # In-place pass returns the same model
+                return _pass_infra.PassResult(model=model, modified=True)
+
+            def ensures(self, model: ir.Model) -> None:
+                # Record which model ensures was called with
+                self.ensures_called_with = model
+
+        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)
+        self.assertIs(pass_.ensures_called_with, result.model)
+        self.assertIs(pass_.ensures_called_with, input_model)
+
+    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:
+                return True
+
+            @property
+            def changes_input(self) -> bool:
+                return True
+
+            def call(self, model: ir.Model) -> _pass_infra.PassResult:
+                return _pass_infra.PassResult(model=model, modified=False)
+
+            def ensures(self, model: ir.Model) -> None:
+                # Simulate a postcondition failure
+                raise ValueError("Postcondition failed")
+
+        pass_ = TestPass()
+        model = ir.Model(graph=ir.Graph([], [], nodes=[]), ir_version=10)
+        
+        with self.assertRaises(PostconditionError) 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:
+                return True
+
+            @property
+            def changes_input(self) -> bool:
+                return True
+
+            def call(self, model: ir.Model) -> _pass_infra.PassResult:
+                return _pass_infra.PassResult(model=model, modified=False)
+
+            def ensures(self, model: ir.Model) -> None:
+                # Directly raise PostconditionError
+                raise PostconditionError("Direct postcondition error")
+
+        pass_ = TestPass()
+        model = ir.Model(graph=ir.Graph([], [], nodes=[]), ir_version=10)
+        
+        with self.assertRaises(PostconditionError) as cm:
+            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=new_graph, ir_version=model.ir_version)
+                return _pass_infra.PassResult(model=new_model, modified=True)
+
+            def ensures(self, model: ir.Model) -> None:
+                # Record the number of nodes in the model passed to ensures
+                self.ensures_model_graph_nodes_count = len(list(model.graph))
+
+        pass_ = ModelModifyingPass()
+        input_model = ir.Model(graph=ir.Graph([], [], nodes=[]), ir_version=10)
+        result = pass_(input_model)
+        
+        # The ensures method should see the modified model with 1 node
+        self.assertEqual(pass_.ensures_model_graph_nodes_count, 1)
+
+
 if __name__ == "__main__":
     unittest.main()