Update tests

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

src/onnx_ir/passes/_pass_infra_test.py

@@ -97,6 +97,32 @@ def ensures(self, model: ir.Model) -> None:
         self.assertIs(pass_.ensures_called_with, result.model)
         self.assertIs(pass_.ensures_called_with, input_model)
 
+    def test_ensures_called_with_result_model_functional_pass(self):
+        """Test that ensures() is called with result.model for functional passes."""
+
+        class TestPass(_pass_infra.FunctionalPass):
+            def __init__(self):
+                self.ensures_called_with = None
+
+            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_postcondition_error_raised_when_ensures_fails(self):
         """Test that PostconditionError is raised when ensures() raises an exception."""
 
@@ -153,36 +179,6 @@ def ensures(self, model: ir.Model) -> None:
         ):
             pass_(model)
 
-    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_model = ir.Model(
-                    graph=ir.Graph(
-                        [],
-                        [output := ir.Value(name="output")],
-                        nodes=[ir.node("TestOp", [], outputs=[output])],
-                    ),
-                    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)
-        _ = 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()