From 46136e6c4d19fc772fa34da2c091100366d2ff49 Mon Sep 17 00:00:00 2001
From: Chris Bianca <me@chrisbianca.com>
Date: Fri, 6 Oct 2017 12:36:41 +0100
Subject: [PATCH] [firestore] Support `update` variations on DocumentReference
 and WriteBatch

---
 lib/modules/firestore/DocumentReference.js    | 22 ++++++++++++++++--
 lib/modules/firestore/WriteBatch.js           | 23 ++++++++++++++++---
 .../tests/firestore/documentReferenceTests.js | 16 +++++++++++--
 tests/src/tests/firestore/firestoreTests.js   |  2 +-
 4 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/lib/modules/firestore/DocumentReference.js b/lib/modules/firestore/DocumentReference.js
index 2210bca1..986f6165 100644
--- a/lib/modules/firestore/DocumentReference.js
+++ b/lib/modules/firestore/DocumentReference.js
@@ -5,7 +5,7 @@
 import CollectionReference from './CollectionReference';
 import DocumentSnapshot from './DocumentSnapshot';
 import Path from './Path';
-import { firestoreAutoId, isFunction, isObject } from '../../utils';
+import { firestoreAutoId, isFunction, isObject, isString } from '../../utils';
 
 export type WriteOptions = {
   merge?: boolean,
@@ -156,7 +156,25 @@ export default class DocumentReference {
       .documentSet(this.path, data, writeOptions);
   }
 
-  update(data: Object): Promise<void> {
+  update(...args: Object | string[]): Promise<void> {
+    let data = {};
+    if (args.length === 1) {
+      if (!isObject(args[0])) {
+        throw new Error('DocumentReference.update failed: If using a single argument, it must be an object.');
+      }
+      data = args[0];
+    } else if (args.length % 2 === 1) {
+      throw new Error('DocumentReference.update failed: Must have either a single object argument, or equal numbers of key/value pairs.');
+    } else {
+      for (let i = 0; i < args.length; i += 2) {
+        const key = args[i];
+        const value = args[i + 1];
+        if (!isString(key)) {
+          throw new Error(`DocumentReference.update failed: Argument at index ${i} must be a string`);
+        }
+        data[key] = value;
+      }
+    }
     return this._firestore._native
       .documentUpdate(this.path, data);
   }
diff --git a/lib/modules/firestore/WriteBatch.js b/lib/modules/firestore/WriteBatch.js
index 7ada92d9..0ed02094 100644
--- a/lib/modules/firestore/WriteBatch.js
+++ b/lib/modules/firestore/WriteBatch.js
@@ -3,6 +3,7 @@
  * WriteBatch representation wrapper
  */
 import DocumentReference from './DocumentReference';
+import { isObject, isString } from '../../utils';
 
 import type { WriteOptions } from './DocumentReference';
 
@@ -58,11 +59,27 @@ export default class WriteBatch {
     return this;
   }
 
-  // TODO: Update to new method signature
-  update(docRef: DocumentReference, data: Object): WriteBatch {
+  update(docRef: DocumentReference, ...args: Object | string[]): WriteBatch {
     // TODO: Validation
     // validate.isDocumentReference('docRef', docRef);
-    // validate.isDocument('data', data, true);
+    let data = {};
+    if (args.length === 1) {
+      if (!isObject(args[0])) {
+        throw new Error('DocumentReference.update failed: If using two arguments, the second must be an object.');
+      }
+      data = args[0];
+    } else if (args.length % 2 === 1) {
+      throw new Error('DocumentReference.update failed: Must have a document reference, followed by either a single object argument, or equal numbers of key/value pairs.');
+    } else {
+      for (let i = 0; i < args.length; i += 2) {
+        const key = args[i];
+        const value = args[i + 1];
+        if (!isString(key)) {
+          throw new Error(`DocumentReference.update failed: Argument at index ${i + 1} must be a string`);
+        }
+        data[key] = value;
+      }
+    }
 
     this._writes.push({
       data,
diff --git a/tests/src/tests/firestore/documentReferenceTests.js b/tests/src/tests/firestore/documentReferenceTests.js
index 846e86a7..4fbd4ca4 100644
--- a/tests/src/tests/firestore/documentReferenceTests.js
+++ b/tests/src/tests/firestore/documentReferenceTests.js
@@ -388,10 +388,22 @@ function documentReferenceTests({ describe, it, context, firebase }) {
     });
 
     context('update()', () => {
-      it('should update Document', () => {
+      it('should update Document using object', () => {
         return firebase.native.firestore()
           .doc('document-tests/doc1')
-          .set({ name: 'updated' })
+          .update({ name: 'updated' })
+          .then(async () => {
+            const doc = await firebase.native.firestore().doc('document-tests/doc1').get();
+            doc.data().name.should.equal('updated');
+          });
+      });
+    });
+
+    context('update()', () => {
+      it('should update Document using key/value pairs', () => {
+        return firebase.native.firestore()
+          .doc('document-tests/doc1')
+          .update('name', 'updated')
           .then(async () => {
             const doc = await firebase.native.firestore().doc('document-tests/doc1').get();
             doc.data().name.should.equal('updated');
diff --git a/tests/src/tests/firestore/firestoreTests.js b/tests/src/tests/firestore/firestoreTests.js
index 26a29893..29fc9703 100644
--- a/tests/src/tests/firestore/firestoreTests.js
+++ b/tests/src/tests/firestore/firestoreTests.js
@@ -36,7 +36,7 @@ function firestoreTests({ describe, it, context, firebase }) {
             .set(nycRef, { name: 'New York City' })
             .set(sfRef, { name: 'San Francisco' })
             .update(nycRef, { population: 1000000 })
-            .update(sfRef, { name: 'San Fran' })
+            .update(sfRef, 'name', 'San Fran')
             .set(lRef, { population: 3000000 }, { merge: true })
             .delete(ayRef)
             .commit()