Fix aggregation failing if index does not exists, added test (#646)

Co-authored-by: Daniil Volynkin <dvvolynkin@avito.ru>

mongomock/aggregate.py

@@ -369,10 +369,14 @@ def _handle_project_operator(self, operator, values):
             values = self.parse(values) if isinstance(values, str) else self.parse_many(values)
             return _GROUPING_OPERATOR_MAP[operator](values)
         if operator == '$arrayElemAt':
-            key, index = values
+            key, value = values
             array = self._parse_basic_expression(key)
-            index = self.parse(index)
-            return array[index]
+            index = self.parse(value)
+            try:
+                return array[index]
+            except IndexError as error:
+                raise KeyError('Array have length less than index value') from error
+
         raise NotImplementedError("Although '%s' is a valid project operator for the "
                                   'aggregation pipeline, it is currently not implemented '
                                   'in Mongomock.' % operator)

tests/test__mongomock.py

@@ -3189,6 +3189,20 @@ def test__aggregate_bug_473(self):
             }},
         ])
 
+    def test__aggregate_array_eleme_at(self):
+        self.cmp.do.drop()
+        self.cmp.do.insert_many([
+            {'values_list': [1, 2]},
+            {'values_list': [1, 2, 3]},
+        ])
+
+        self.cmp.compare.aggregate([{
+            '$project': {
+                'first_user_id': {'$arrayElemAt': ['$values_list', 2]},
+                'other_user_id': {'$arrayElemAt': ['$values_list', -1]},
+            },
+        }])
+
     def test_aggregate_bug_607(self):
         """Regression test for bug https://github.com/mongomock/mongomock/issues/607."""
         self.cmp.do.drop()