Move onSnapshot handers to the queue so that a call to flush will call the onSnapshot callbacks no mater where the flush is called from.

Author: BrianChapman

src/firestore-document.js

@@ -23,14 +23,9 @@ function MockFirestoreDocument(path, data, parent, name, CollectionReference) {
   if (parent) parent.children[this.id] = this;
   this.data = null;
   this._dataChanged(_.cloneDeep(data) || null);
-  this.onSnapshotSubscribers = [];
 }
 
 MockFirestoreDocument.prototype.flush = function (delay) {
-  var self = this;
-  _.forEach(this.onSnapshotSubscribers, function (subscriber) {
-    self._defer('onSnapshot', _.toArray(arguments), subscriber);
-  });
   this.queue.flush(delay);
   return this;
 };
@@ -224,20 +219,18 @@ MockFirestoreDocument.prototype.onSnapshot = function (optionsOrObserverOrOnNext
     // compare the current state to the one from when this function was created
     // and send the data to the callback if different.
     if (err === null) {
-      if (JSON.stringify(this.data) !== JSON.stringify(context.data) || includeMetadataChanges) {
+      if (JSON.stringify(self.data) !== JSON.stringify(context.data) || includeMetadataChanges) {
         onNext(new DocumentSnapshot(self.id, self.ref, self._getData()));
-        context.data = this._getData();
+        context.data = self._getData();
       }
     } else {
       onError(err);
     }
   };
-  this.onSnapshotSubscribers.push(onSnapshot);
+  var unsubscribe = this.queue.onPostFlush(onSnapshot);
 
   // return the unsubscribe function
-  return function () {
-    self.onSnapshotSubscribers.pop(onSnapshot);
-  };
+  return unsubscribe;
 };
 
 

src/firestore-query.js

@@ -23,14 +23,9 @@ function MockFirestoreQuery(path, data, parent, name) {
   this.orderedDirections = [];
   this.limited = 0;
   this._setData(data);
-  this.onSnapshotSubscribers = [];
 }
 
 MockFirestoreQuery.prototype.flush = function (delay) {
-  var self = this;
-  _.forEach(this.onSnapshotSubscribers, function (subscriber) {
-    self._defer('onSnapshot', _.toArray(arguments), subscriber);
-  });
   this.queue.flush(delay);
   return this;
 };
@@ -163,7 +158,7 @@ MockFirestoreQuery.prototype.onSnapshot = function (optionsOrObserverOrOnNext, o
     // compare the current state to the one from when this function was created
     // and send the data to the callback if different.
     if (err === null) {
-      this.get().then(function (querySnapshot) {
+      self.get().then(function (querySnapshot) {
         var results = self._results();
         if (JSON.stringify(results) !== JSON.stringify(context.data) || includeMetadataChanges) {
           onNext(new QuerySnapshot(self.parent === null ? self : self.parent.collection(self.id), results));
@@ -175,12 +170,10 @@ MockFirestoreQuery.prototype.onSnapshot = function (optionsOrObserverOrOnNext, o
       onError(err);
     }
   };
-  this.onSnapshotSubscribers.push(onSnapshot);
+  var unsubscribe = this.queue.onPostFlush(onSnapshot);
 
   // return the unsubscribe function
-  return function () {
-    self.onSnapshotSubscribers.pop(onSnapshot);
-  };
+  return unsubscribe;
 };
 
 MockFirestoreQuery.prototype._results = function () {

src/queue.js

@@ -6,6 +6,7 @@ var EventEmitter = require('events').EventEmitter;
 
 function FlushQueue () {
   this.events = [];
+  this.postFlushListeners = [];
 }
 
 FlushQueue.prototype.push = function () {
@@ -23,6 +24,14 @@ FlushQueue.prototype.push = function () {
   }));
 };
 
+FlushQueue.prototype.onPostFlush = function(subscriber) {
+  this.postFlushListeners.push(subscriber);
+  var self = this;
+  return function() {
+    self.postFlushListeners.pop(subscriber);
+  };
+};
+
 FlushQueue.prototype.flushing = false;
 
 FlushQueue.prototype.flush = function (delay) {
@@ -33,6 +42,9 @@ FlushQueue.prototype.flush = function (delay) {
   }
   function process () {
     self.flushing = true;
+    _.forEach(self.postFlushListeners, function (subscriber) {
+      self.push(subscriber);
+    });
     while (self.events.length) {
       self.events[0].run();
     }

test/unit/firestore-collection.js

@@ -412,6 +412,7 @@ describe('MockFirestoreCollection', function () {
         expect(err).to.equal(error);
         done();
       });
+      collection.doc('a').update({name: 'A'}, {setMerge: true});
       collection.flush();
     });
 

test/unit/firestore-document.js

@@ -471,7 +471,7 @@ describe('MockFirestoreDocument', function () {
         done();
       });
       doc.update({newTitle: 'A new title'}, {setMerge: true});
-      doc.flush();
+      db.flush();
     });
 
     it('returns error if error occured', function (done) {
@@ -483,6 +483,7 @@ describe('MockFirestoreDocument', function () {
         expect(err).to.equal(error);
         done();
       });
+      doc.update({name: 'A'}, {setMerge: true});
       doc.flush();
     });