Add tests for templatizing rules returning tuples

Pyomo · Aug 31, 2023 · ecbd709 · ecbd709
1 parent 677289a
commit ecbd709
Showing 1 changed file with 47 additions and 0 deletions.
diff --git a/pyomo/core/tests/unit/test_template_expr.py b/pyomo/core/tests/unit/test_template_expr.py
@@ -353,6 +353,53 @@ def c(m, i):
         indices[0].set_value(2)
         self.assertEqual(str(resolve_template(template)), 'x[2]  <=  0')
 
+    def test_tuple_rules(self):
+        m = ConcreteModel()
+        m.I = RangeSet(3)
+        m.x = Var(m.I)
+
+        @m.Constraint(m.I)
+        def c(m, i):
+            return (None, m.x[i], 0)
+
+        template, indices = templatize_constraint(m.c)
+        self.assertEqual(len(indices), 1)
+        self.assertIs(indices[0]._set, m.I)
+        self.assertEqual(str(template), "x[_1]  <=  0")
+        # Test that the RangeSet iterator was put back
+        self.assertEqual(list(m.I), list(range(1, 4)))
+        # Evaluate the template
+        indices[0].set_value(2)
+        self.assertEqual(str(resolve_template(template)), 'x[2]  <=  0')
+
+        @m.Constraint(m.I)
+        def d(m, i):
+            return (0, m.x[i], 10)
+
+        template, indices = templatize_constraint(m.d)
+        self.assertEqual(len(indices), 1)
+        self.assertIs(indices[0]._set, m.I)
+        self.assertEqual(str(template), "0  <=  x[_1]  <=  10")
+        # Test that the RangeSet iterator was put back
+        self.assertEqual(list(m.I), list(range(1, 4)))
+        # Evaluate the template
+        indices[0].set_value(2)
+        self.assertEqual(str(resolve_template(template)), '0  <=  x[2]  <=  10')
+
+        @m.Constraint(m.I)
+        def e(m, i):
+            return (m.x[i], 0)
+
+        template, indices = templatize_constraint(m.e)
+        self.assertEqual(len(indices), 1)
+        self.assertIs(indices[0]._set, m.I)
+        self.assertEqual(str(template), "x[_1]  ==  0")
+        # Test that the RangeSet iterator was put back
+        self.assertEqual(list(m.I), list(range(1, 4)))
+        # Evaluate the template
+        indices[0].set_value(2)
+        self.assertEqual(str(resolve_template(template)), 'x[2]  ==  0')
+
     def test_simple_rule_nonfinite_set(self):
         m = ConcreteModel()
         m.x = Var(Integers, dense=False)