1 files changed, 13 insertions, 3 deletions
diff --git a/Firestore/Example/Tests/Util/FSTEventAccumulator.m b/Firestore/Example/Tests/Util/FSTEventAccumulator.m
index b44ec67..1b1e276 100644
--- a/Firestore/Example/Tests/Util/FSTEventAccumulator.m
+++ b/Firestore/Example/Tests/Util/FSTEventAccumulator.m
@@ -68,9 +68,9 @@ NS_ASSUME_NONNULL_BEGIN
   return events[0];
 }
 
-// Overrides the handler property
-- (void (^)(id _Nullable, NSError *))handler {
-  return ^void(id _Nullable value, NSError *error) {
+// Override the valueEventHandler property
+- (void (^)(id _Nullable, NSError *_Nullable))valueEventHandler {
+  return ^void(id _Nullable value, NSError *_Nullable error) {
     // We can't store nil in the _events array, but these are still interesting to tests so store
     // NSNull instead.
     id event = value ? value : [NSNull null];
@@ -82,6 +82,16 @@ NS_ASSUME_NONNULL_BEGIN
   };
 }
 
+// Override the errorEventHandler property
+- (void (^)(NSError *_Nullable))errorEventHandler {
+  return ^void(NSError *_Nullable error) {
+    @synchronized(self) {
+      [_events addObject:[NSNull null]];
+      [self checkFulfilled];
+    }
+  };
+}
+
 - (void)checkFulfilled {
   if (_events.count >= self.maxEvents) {
     [self.expectation fulfill];