diff --git a/nucypher/policy/conditions/lingo.py b/nucypher/policy/conditions/lingo.py index 4e220c080..666888616 100644 --- a/nucypher/policy/conditions/lingo.py +++ b/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) diff --git a/tests/unit/conditions/test_return_value.py b/tests/unit/conditions/test_return_value.py index 1964f8da1..8e2100e4a 100644 --- a/tests/unit/conditions/test_return_value.py +++ b/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):