Include tests for handling list index in ReturnValue

nucypher/policy/conditions/lingo.py

@@ -126,11 +126,18 @@ class ReturnValueTest:
                 f"'{self.key}' is an unprocessed context variable and is not valid "
                 f"for condition evaluation."
             )
-        if self.key:
+        if self.key and type(data) is dict:
             try:
                 data = data[self.key]
             except KeyError:
                 raise KeyError(f"Key '{self.key}' not found in return data.")
+
+        if type(self.key) is int and type(data) is list:
+            try:
+                data = data[self.key]
+            except IndexError:
+                raise IndexError(f"Index '{self.key}' not found in return data.")
+
         left_operand = self._sanitize_value(data)
         right_operand = self._sanitize_value(self.value)
         result = self._COMPARATOR_FUNCTIONS[self.comparator](left_operand, right_operand)

tests/unit/conditions/test_return_value.py

@@ -31,6 +31,25 @@ def test_return_value_key():
     with pytest.raises(KeyError):
         test.eval({"bond": 1})
 
+    test = ReturnValueTest(comparator=">", value="0", key=4)
+    assert test.eval({4: 1})
+    assert not test.eval({4: -1})
+
+    with pytest.raises(KeyError):
+        test.eval({5: 1})
+
+
+def test_return_value_index():
+    test = ReturnValueTest(comparator=">", value="0", key=0)
+    assert test.eval([1])
+    assert not test.eval([-1])
+
+    test = ReturnValueTest(comparator="==", value='"james"', key=3)
+    assert test.eval([0, 1, 2, '"james"'])
+
+    with pytest.raises(IndexError):
+        test.eval([0, 1, 2])
+
 
 def test_return_value_test_invalid_comparators():
     with pytest.raises(ReturnValueTest.InvalidExpression):