Gr-diff retrofit: store comment state

Fixing UI data pipe line, re-rendering:
 - side, draft text and editing status in UI comment objects
 - update gr-diff UI model on comment save/update

Feature: Issue 3910
Change-Id: I96f714c7de9add6e316dcf64bb7d566690b9d3ae
diff --git a/polygerrit-ui/app/elements/diff/gr-diff-comment-thread/gr-diff-comment-thread.js b/polygerrit-ui/app/elements/diff/gr-diff-comment-thread/gr-diff-comment-thread.js
index aa5d8bd..7304ac7 100644
--- a/polygerrit-ui/app/elements/diff/gr-diff-comment-thread/gr-diff-comment-thread.js
+++ b/polygerrit-ui/app/elements/diff/gr-diff-comment-thread/gr-diff-comment-thread.js
@@ -61,10 +61,8 @@
       }
 
       var draft = this._newDraft(opt_lineNum);
+      draft.__editing = true;
       this.push('comments', draft);
-      this.async(function() {
-        this._commentElWithDraftID(draft.__draftID).editing = true;
-      }.bind(this), 1);
     },
 
     _getLoggedIn: function() {
@@ -121,13 +119,8 @@
             function(line) { return ' > ' + line; }).join('\n') + '\n\n';
       }
       var reply = this._newReply(comment.id, comment.line, quoteStr);
+      reply.__editing = true;
       this.push('comments', reply);
-
-      // Allow the reply to render in the dom-repeat.
-      this.async(function() {
-        var commentEl = this._commentElWithDraftID(reply.__draftID);
-        commentEl.editing = true;
-      }.bind(this), 1);
     },
 
     _handleCommentDone: function(e) {
diff --git a/polygerrit-ui/app/elements/diff/gr-diff-comment/gr-diff-comment.js b/polygerrit-ui/app/elements/diff/gr-diff-comment/gr-diff-comment.js
index 525bac0..7a15754 100644
--- a/polygerrit-ui/app/elements/diff/gr-diff-comment/gr-diff-comment.js
+++ b/polygerrit-ui/app/elements/diff/gr-diff-comment/gr-diff-comment.js
@@ -15,6 +15,7 @@
   'use strict';
 
   var STORAGE_DEBOUNCE_INTERVAL = 400;
+  var UPDATE_DEBOUNCE_INTERVAL = 500;
 
   Polymer({
     is: 'gr-diff-comment',
@@ -43,11 +44,18 @@
      * @event comment-save
      */
 
+    /**
+     * Fired when this comment is updated.
+     *
+     * @event comment-update
+     */
+
     properties: {
       changeNum: String,
       comment: {
         type: Object,
         notify: true,
+        observer: '_commentChanged',
       },
       disabled: {
         type: Boolean,
@@ -81,6 +89,10 @@
       '_loadLocalDraft(changeNum, patchNum, comment)',
     ],
 
+    detached: function() {
+      this.flushDebouncer('fire-update');
+    },
+
     save: function() {
       this.comment.message = this._messageText;
       this.disabled = true;
@@ -106,8 +118,7 @@
           }
           this.comment = comment;
           this.editing = false;
-          this.fire('comment-save');
-
+          this.fire('comment-save', {comment: this.comment});
           return obj;
         }.bind(this));
       }.bind(this)).catch(function(err) {
@@ -116,6 +127,16 @@
       }.bind(this));
     },
 
+    _commentChanged: function(comment) {
+      this.editing = !!comment.__editing;
+    },
+
+    _fireUpdate: function() {
+      this.debounce('fire-update', function() {
+        this.fire('comment-update', {comment: this.comment});
+      }, UPDATE_DEBOUNCE_INTERVAL);
+    },
+
     _draftChanged: function(draft) {
       this.$.container.classList.toggle('draft', draft);
     },
@@ -134,6 +155,10 @@
       if (this.comment && this.comment.id) {
         this.$$('.cancel').hidden = !editing;
       }
+      if (this.comment) {
+        this.comment.__editing = this.editing;
+      }
+      this._fireUpdate();
     },
 
     _computeLinkToComment: function(comment) {
@@ -174,6 +199,7 @@
         } else {
           this.$.storage.setDraftComment(commentLocation, message);
         }
+        this._fireUpdate();
       }, STORAGE_DEBOUNCE_INTERVAL);
     },
 
@@ -218,7 +244,7 @@
     _handleCancel: function(e) {
       this._preventDefaultAndBlur(e);
       if (this.comment.message == null || this.comment.message.length == 0) {
-        this.fire('comment-discard');
+        this.fire('comment-discard', {comment: this.comment});
         return;
       }
       this._messageText = this.comment.message;
@@ -234,20 +260,20 @@
       this.disabled = true;
       if (!this.comment.id) {
         this.disabled = false;
-        this.fire('comment-discard');
+        this.fire('comment-discard', {comment: this.comment});
         return;
       }
 
-      this._xhrPromise =
-          this._deleteDraft(this.comment).then(function(response) {
-        this.disabled = false;
-        if (!response.ok) { return response; }
+      this._xhrPromise = this._deleteDraft(this.comment).then(
+          function(response) {
+            this.disabled = false;
+            if (!response.ok) { return response; }
 
-        this.fire('comment-discard');
-      }.bind(this)).catch(function(err) {
-        this.disabled = false;
-        throw err;
-      }.bind(this));;
+            this.fire('comment-discard', {comment: this.comment});
+          }.bind(this)).catch(function(err) {
+            this.disabled = false;
+            throw err;
+          }.bind(this));
     },
 
     _preventDefaultAndBlur: function(e) {
diff --git a/polygerrit-ui/app/elements/diff/gr-diff-comment/gr-diff-comment_test.html b/polygerrit-ui/app/elements/diff/gr-diff-comment/gr-diff-comment_test.html
index 8ef222e..6c27a36 100644
--- a/polygerrit-ui/app/elements/diff/gr-diff-comment/gr-diff-comment_test.html
+++ b/polygerrit-ui/app/elements/diff/gr-diff-comment/gr-diff-comment_test.html
@@ -204,16 +204,50 @@
     });
 
     test('draft saving/editing', function(done) {
+      var fireStub = sinon.stub(element, 'fire');
+
       element.draft = true;
       MockInteractions.tap(element.$$('.edit'));
       element._messageText = 'good news, everyone!';
+      element.flushDebouncer('fire-update');
+      element.flushDebouncer('store');
+      assert(fireStub.calledWith('comment-update'),
+             'comment-update should be sent');
+      assert.deepEqual(fireStub.lastCall.args, [
+        'comment-update', {
+          comment: {
+            __draft: true,
+            __draftID: 'temp_draft_id',
+            __editing: true,
+            line: 5,
+            path: '/path/to/file',
+          },
+        },
+      ]);
       MockInteractions.tap(element.$$('.save'));
+
       assert.isTrue(element.disabled,
           'Element should be disabled when creating draft.');
 
       element._xhrPromise.then(function(draft) {
+        assert(fireStub.calledWith('comment-save'),
+               'comment-save should be sent');
+        assert.deepEqual(fireStub.lastCall.args, [
+          'comment-save', {
+            comment: {
+              __draft: true,
+              __draftID: 'temp_draft_id',
+              __editing: false,
+              id: 'baf0414d_40572e03',
+              line: 5,
+              message: 'saved!',
+              path: '/path/to/file',
+              updated: '2015-12-08 21:52:36.177000000',
+            },
+          },
+        ]);
         assert.isFalse(element.disabled,
-            'Element should be enabled when done creating draft.');
+                       'Element should be enabled when done creating draft.');
         assert.equal(draft.message, 'saved!');
         assert.isFalse(element.editing);
       }).then(function() {
diff --git a/polygerrit-ui/app/elements/diff/gr-diff/gr-diff.js b/polygerrit-ui/app/elements/diff/gr-diff/gr-diff.js
index df91a38..ea7111a 100644
--- a/polygerrit-ui/app/elements/diff/gr-diff/gr-diff.js
+++ b/polygerrit-ui/app/elements/diff/gr-diff/gr-diff.js
@@ -71,15 +71,17 @@
       '_prefsChanged(prefs.*, viewMode)',
     ],
 
+    listeners: {
+      'thread-discard': '_handleThreadDiscard',
+      'comment-discard': '_handleCommentDiscard',
+      'comment-update': '_handleCommentUpdate',
+      'comment-save': '_handleCommentSave',
+    },
+
     attached: function() {
       this._getLoggedIn().then(function(loggedIn) {
         this._loggedIn = loggedIn;
       }.bind(this));
-
-      this.addEventListener('thread-discard',
-          this._handleThreadDiscard.bind(this));
-      this.addEventListener('comment-discard',
-          this._handleCommentDiscard.bind(this));
     },
 
     reload: function() {
@@ -234,29 +236,71 @@
     },
 
     _handleCommentDiscard: function(e) {
-      var comment = Polymer.dom(e).rootTarget.comment;
-      this._removeComment(comment);
+      var comment = e.detail.comment;
+      this._removeComment(comment, e.target.patchNum);
     },
 
-    _removeComment: function(comment) {
-      if (!comment.id) { return; }
-      this._removeCommentFromSide(comment, DiffSide.LEFT) ||
-          this._removeCommentFromSide(comment, DiffSide.RIGHT);
+    _removeComment: function(comment, opt_patchNum) {
+      var side = this._findCommentSide(comment, opt_patchNum);
+      this._removeCommentFromSide(comment, side);
+    },
+
+    _findCommentSide: function(comment, opt_patchNum) {
+      if (comment.side === 'PARENT') {
+        return DiffSide.LEFT;
+      } else {
+        return this._comments.meta.patchRange.basePatchNum === opt_patchNum ?
+            DiffSide.LEFT : DiffSide.RIGHT;
+      }
+    },
+
+    _handleCommentSave: function(e) {
+      var comment = e.detail.comment;
+      var side = this._findCommentSide(comment, e.target.patchNum);
+      var idx = this._findDraftIndex(comment, side);
+      this.set(['_comments', side, idx], comment);
+    },
+
+    _handleCommentUpdate: function(e) {
+      var comment = e.detail.comment;
+      var side = this._findCommentSide(comment, e.target.patchNum);
+      var idx = this._findCommentIndex(comment, side);
+      if (idx === -1) {
+        idx = this._findDraftIndex(comment, side);
+      }
+      if (idx !== -1) { // Update draft or comment.
+        this.set(['_comments', side, idx], comment);
+      } else { // Create new draft.
+        this.push(['_comments', side], comment);
+      }
     },
 
     _removeCommentFromSide: function(comment, side) {
-      var idx = -1;
-      for (var i = 0; i < this._comments[side].length; i++) {
-        if (this._comments[side][i].id === comment.id) {
-          idx = i;
-          break;
-        }
+      var idx = this._findCommentIndex(comment, side);
+      if (idx === -1) {
+        idx = this._findDraftIndex(comment, side);
       }
       if (idx !== -1) {
         this.splice('_comments.' + side, idx, 1);
-        return true;
       }
-      return false;
+    },
+
+    _findCommentIndex: function(comment, side) {
+      if (!comment.id || !this._comments[side]) {
+        return -1;
+      }
+      return this._comments[side].findIndex(function(item) {
+        return item.id === comment.id;
+      });
+    },
+
+    _findDraftIndex: function(comment, side) {
+      if (!comment.__draftID || !this._comments[side]) {
+        return -1;
+      }
+      return this._comments[side].findIndex(function(item) {
+        return item.__draftID === comment.__draftID;
+      });
     },
 
     _handleMouseDown: function(e) {
diff --git a/polygerrit-ui/app/elements/diff/gr-diff/gr-diff_test.html b/polygerrit-ui/app/elements/diff/gr-diff/gr-diff_test.html
index d03831f..aec32b6 100644
--- a/polygerrit-ui/app/elements/diff/gr-diff/gr-diff_test.html
+++ b/polygerrit-ui/app/elements/diff/gr-diff/gr-diff_test.html
@@ -35,85 +35,29 @@
   suite('gr-diff tests', function() {
     var element;
 
-    setup(function() {
-      stub('gr-rest-api-interface', {
-        getLoggedIn: function() { return Promise.resolve(false); },
+    suite('not logged in', function() {
+
+      setup(function() {
+        stub('gr-rest-api-interface', {
+          getLoggedIn: function() { return Promise.resolve(false); },
+        });
+        element = fixture('basic');
       });
-      element = fixture('basic');
-    });
 
-    test('get drafts logged out', function(done) {
-      element.patchRange = {basePatchNum: 0, patchNum: 0};
+      test('get drafts', function(done) {
+        element.patchRange = {basePatchNum: 0, patchNum: 0};
 
-      var getDraftsStub = sinon.stub(element.$.restAPI, 'getDiffDrafts');
-      var loggedInStub = sinon.stub(element, '_getLoggedIn',
-          function() { return Promise.resolve(false); });
-      element._getDiffDrafts().then(function(result) {
-        assert.deepEqual(result, {baseComments: [], comments: []});
-        sinon.assert.notCalled(getDraftsStub);
-        loggedInStub.restore();
-        getDraftsStub.restore();
-        done();
+        var getDraftsStub = sinon.stub(element.$.restAPI, 'getDiffDrafts');
+        element._getDiffDrafts().then(function(result) {
+          assert.deepEqual(result, {baseComments: [], comments: []});
+          sinon.assert.notCalled(getDraftsStub);
+          getDraftsStub.restore();
+          done();
+        });
       });
-    });
 
-    test('get drafts logged in', function(done) {
-      element.patchRange = {basePatchNum: 0, patchNum: 0};
-      var draftsResponse = {
-        baseComments: [{id: 'foo'}],
-        comments: [{id: 'bar'}],
-      };
-      var getDraftsStub = sinon.stub(element.$.restAPI, 'getDiffDrafts',
-          function() { return Promise.resolve(draftsResponse); });
-      var loggedInStub = sinon.stub(element, '_getLoggedIn',
-          function() { return Promise.resolve(true); });
-      element._getDiffDrafts().then(function(result) {
-        assert.deepEqual(result, draftsResponse);
-        loggedInStub.restore();
-        getDraftsStub.restore();
-        done();
-      });
-    });
-
-    test('get comments and drafts', function(done) {
-      var loggedInStub = sinon.stub(element, '_getLoggedIn',
-          function() { return Promise.resolve(true); });
-      var comments = {
-        baseComments: [
-          {id: 'bc1'},
-          {id: 'bc2'},
-        ],
-        comments: [
-          {id: 'c1'},
-          {id: 'c2'},
-        ],
-      };
-      var diffCommentsStub = sinon.stub(element, '_getDiffComments',
-          function() { return Promise.resolve(comments); });
-
-      var drafts = {
-        baseComments: [
-          {id: 'bd1'},
-          {id: 'bd2'},
-        ],
-        comments: [
-          {id: 'd1'},
-          {id: 'd2'},
-        ],
-      };
-      var diffDraftsStub = sinon.stub(element, '_getDiffDrafts',
-          function() { return Promise.resolve(drafts); });
-
-      element.changeNum = '42';
-      element.patchRange = {
-        basePatchNum: 'PARENT',
-        patchNum: 3,
-      };
-      element.path = '/path/to/foo';
-      element.projectConfig = {foo: 'bar'};
-
-      element._getDiffCommentsAndDrafts().then(function(result) {
-        assert.deepEqual(result, {
+      test('remove comment', function() {
+        element._comments = {
           meta: {
             changeNum: '42',
             patchRange: {
@@ -124,10 +68,10 @@
             projectConfig: {foo: 'bar'},
           },
           left: [
-            {id: 'bc1'},
-            {id: 'bc2'},
-            {id: 'bd1', __draft: true},
-            {id: 'bd2', __draft: true},
+            {id: 'bc1', side: 'PARENT'},
+            {id: 'bc2', side: 'PARENT'},
+            {id: 'bd1', __draft: true, side: 'PARENT'},
+            {id: 'bd2', __draft: true, side: 'PARENT'},
           ],
           right: [
             {id: 'c1'},
@@ -135,212 +79,322 @@
             {id: 'd1', __draft: true},
             {id: 'd2', __draft: true},
           ],
-        });
+        };
 
-        diffCommentsStub.restore();
-        diffDraftsStub.restore();
-        loggedInStub.restore();
-        done();
+        element._removeComment({});
+        // Using JSON.stringify because Safari 9.1 (11601.5.17.1) doesn’t seem
+        // to believe that one object deepEquals another even when they do :-/.
+        assert.equal(JSON.stringify(element._comments), JSON.stringify({
+          meta: {
+            changeNum: '42',
+            patchRange: {
+              basePatchNum: 'PARENT',
+              patchNum: 3,
+            },
+            path: '/path/to/foo',
+            projectConfig: {foo: 'bar'},
+          },
+          left: [
+            {id: 'bc1', side: 'PARENT'},
+            {id: 'bc2', side: 'PARENT'},
+            {id: 'bd1', __draft: true, side: 'PARENT'},
+            {id: 'bd2', __draft: true, side: 'PARENT'},
+          ],
+          right: [
+            {id: 'c1'},
+            {id: 'c2'},
+            {id: 'd1', __draft: true},
+            {id: 'd2', __draft: true},
+          ],
+        }));
+
+        element._removeComment({id: 'bc2', side: 'PARENT'});
+        assert.equal(JSON.stringify(element._comments), JSON.stringify({
+          meta: {
+            changeNum: '42',
+            patchRange: {
+              basePatchNum: 'PARENT',
+              patchNum: 3,
+            },
+            path: '/path/to/foo',
+            projectConfig: {foo: 'bar'},
+          },
+          left: [
+            {id: 'bc1', side: 'PARENT'},
+            {id: 'bd1', __draft: true, side: 'PARENT'},
+            {id: 'bd2', __draft: true, side: 'PARENT'},
+          ],
+          right: [
+            {id: 'c1'},
+            {id: 'c2'},
+            {id: 'd1', __draft: true},
+            {id: 'd2', __draft: true},
+          ],
+        }));
+
+        element._removeComment({id: 'd2'});
+        assert.deepEqual(JSON.stringify(element._comments), JSON.stringify({
+          meta: {
+            changeNum: '42',
+            patchRange: {
+              basePatchNum: 'PARENT',
+              patchNum: 3,
+            },
+            path: '/path/to/foo',
+            projectConfig: {foo: 'bar'},
+          },
+          left: [
+            {id: 'bc1', side: 'PARENT'},
+            {id: 'bd1', __draft: true, side: 'PARENT'},
+            {id: 'bd2', __draft: true, side: 'PARENT'},
+          ],
+          right: [
+            {id: 'c1'},
+            {id: 'c2'},
+            {id: 'd1', __draft: true},
+          ],
+        }));
+      });
+
+      test('renders image diffs', function(done) {
+        var mockDiff = {
+          meta_a: {name: 'carrot.jpg', content_type: 'image/jpeg', lines: 66},
+          meta_b: {name: 'carrot.jpg', content_type: 'image/jpeg', lines: 560},
+          intraline_status: 'OK',
+          change_type: 'MODIFIED',
+          diff_header: [
+            'diff --git a/carrot.jpg b/carrot.jpg',
+            'index 2adc47d..f9c2f2c 100644',
+            '--- a/carrot.jpg',
+            '+++ b/carrot.jpg',
+            'Binary files differ',
+          ],
+          content: [{skip: 66}],
+          binary: true,
+        };
+        var mockFile1 = {
+          body: 'Qk06AAAAAAAAADYAAAAoAAAAAQAAAP////8BACAAAAAAAAAAAAATCwAAEwsA' +
+              'AAAAAAAAAAAAAAAA/w==',
+          type: 'image/bmp',
+        };
+        var mockFile2 = {
+          body: 'Qk06AAAAAAAAADYAAAAoAAAAAQAAAP////8BACAAAAAAAAAAAAATCwAAEwsA' +
+              'AAAAAAAAAAAA/////w==',
+          type: 'image/bmp'
+        };
+        var mockCommit = {
+          commit: '9a1a1d10baece5efbba10bc4ccf808a67a50ac0a',
+          parents: [{
+            commit: '7338aa9adfe57909f1fdaf88975cdea467d3382f',
+            subject: 'Added a carrot',
+          }],
+          author: {
+            name: 'Wyatt Allen',
+            email: 'wyatta@google.com',
+            date: '2016-05-23 21:44:51.000000000',
+            tz: -420,
+          },
+          committer: {
+            name: 'Wyatt Allen',
+            email: 'wyatta@google.com',
+            date: '2016-05-25 00:25:41.000000000',
+            tz: -420,
+          },
+          subject: 'Updated the carrot',
+          message: 'Updated the carrot\n\nChange-Id: Iabcd123\n',
+        };
+        var mockComments = {baseComments: [], comments: []};
+
+        var stubs = [];
+        stubs.push(sinon.stub(element, '_getDiff',
+            function() { return Promise.resolve(mockDiff); }));
+        stubs.push(sinon.stub(element.$.restAPI, 'getCommitInfo',
+            function() { return Promise.resolve(mockCommit); }));
+        stubs.push(sinon.stub(element.$.restAPI,
+            'getCommitFileContents',
+            function() { return Promise.resolve(mockFile1); }));
+        stubs.push(sinon.stub(element.$.restAPI,
+            'getChangeFileContents',
+            function() { return Promise.resolve(mockFile2); }));
+        stubs.push(sinon.stub(element.$.restAPI, '_getDiffComments',
+            function() { return Promise.resolve(mockComments); }));
+        stubs.push(sinon.stub(element.$.restAPI, 'getDiffDrafts',
+            function() { return Promise.resolve(mockComments); }));
+
+        element.patchRange = {basePatchNum: 'PARENT', patchNum: 1};
+
+        var rendered = function() {
+          // Recognizes that it should be an image diff.
+          assert.isTrue(element.isImageDiff);
+          assert.instanceOf(element._getDiffBuilder(element._diff,
+              element._comments, element.prefs), GrDiffBuilderImage);
+
+          // Left image rendered with the parent commit's version of the file.
+          var leftInmage = element.$.diffTable.querySelector('td.left img');
+          assert.isOk(leftInmage);
+          assert.equal(leftInmage.getAttribute('src'),
+              'data:image/bmp;base64, ' + mockFile1.body);
+
+          // Right image rendered with this change's revision of the image.
+          var rightInmage = element.$.diffTable.querySelector('td.right img');
+          assert.isOk(rightInmage);
+          assert.equal(rightInmage.getAttribute('src'),
+              'data:image/bmp;base64, ' + mockFile2.body);
+
+          // Cleanup.
+          element.removeEventListener('render', rendered);
+          stubs.forEach(function(stub) { stub.restore(); });
+
+          done();
+        };
+
+        element.addEventListener('render', rendered);
+
+        element.$.restAPI.getDiffPreferences().then(function(prefs) {
+          element.prefs = prefs;
+          element.reload();
+        });
       });
     });
 
-    test('remove comment', function() {
-      element._comments = {
-        meta: {
-          changeNum: '42',
-          patchRange: {
-            basePatchNum: 'PARENT',
-            patchNum: 3,
-          },
-          path: '/path/to/foo',
-          projectConfig: {foo: 'bar'},
-        },
-        left: [
-          {id: 'bc1'},
-          {id: 'bc2'},
-          {id: 'bd1', __draft: true},
-          {id: 'bd2', __draft: true},
-        ],
-        right: [
-          {id: 'c1'},
-          {id: 'c2'},
-          {id: 'd1', __draft: true},
-          {id: 'd2', __draft: true},
-        ],
-      };
+    suite('logged in', function() {
 
-      element._removeComment({});
-      // Using JSON.stringify because Safari 9.1 (11601.5.17.1) doesn’t seem to
-      // believe that one object deepEquals another even when they do :-/.
-      assert.equal(JSON.stringify(element._comments), JSON.stringify({
-        meta: {
-          changeNum: '42',
-          patchRange: {
-            basePatchNum: 'PARENT',
-            patchNum: 3,
-          },
-          path: '/path/to/foo',
-          projectConfig: {foo: 'bar'},
-        },
-        left: [
-          {id: 'bc1'},
-          {id: 'bc2'},
-          {id: 'bd1', __draft: true},
-          {id: 'bd2', __draft: true},
-        ],
-        right: [
-          {id: 'c1'},
-          {id: 'c2'},
-          {id: 'd1', __draft: true},
-          {id: 'd2', __draft: true},
-        ],
-      }));
+      setup(function() {
+        stub('gr-rest-api-interface', {
+          getLoggedIn: function() { return Promise.resolve(true); },
+        });
+        element = fixture('basic');
+      });
 
-      element._removeComment({id: 'bc2'});
-      assert.equal(JSON.stringify(element._comments), JSON.stringify({
-        meta: {
-          changeNum: '42',
-          patchRange: {
-            basePatchNum: 'PARENT',
-            patchNum: 3,
-          },
-          path: '/path/to/foo',
-          projectConfig: {foo: 'bar'},
-        },
-        left: [
-          {id: 'bc1'},
-          {id: 'bd1', __draft: true},
-          {id: 'bd2', __draft: true},
-        ],
-        right: [
-          {id: 'c1'},
-          {id: 'c2'},
-          {id: 'd1', __draft: true},
-          {id: 'd2', __draft: true},
-        ],
-      }));
+      test('get drafts', function(done) {
+        element.patchRange = {basePatchNum: 0, patchNum: 0};
+        var draftsResponse = {
+          baseComments: [{id: 'foo'}],
+          comments: [{id: 'bar'}],
+        };
+        var getDraftsStub = sinon.stub(element.$.restAPI, 'getDiffDrafts',
+            function() { return Promise.resolve(draftsResponse); });
+        element._getDiffDrafts().then(function(result) {
+          assert.deepEqual(result, draftsResponse);
+          getDraftsStub.restore();
+          done();
+        });
+      });
 
-      element._removeComment({id: 'd2'});
-      assert.deepEqual(JSON.stringify(element._comments), JSON.stringify({
-        meta: {
-          changeNum: '42',
-          patchRange: {
-            basePatchNum: 'PARENT',
-            patchNum: 3,
-          },
-          path: '/path/to/foo',
-          projectConfig: {foo: 'bar'},
-        },
-        left: [
-          {id: 'bc1'},
-          {id: 'bd1', __draft: true},
-          {id: 'bd2', __draft: true},
-        ],
-        right: [
-          {id: 'c1'},
-          {id: 'c2'},
-          {id: 'd1', __draft: true},
-        ],
-      }));
-    });
+      test('get comments and drafts', function(done) {
+        var comments = {
+          baseComments: [
+            {id: 'bc1'},
+            {id: 'bc2'},
+          ],
+          comments: [
+            {id: 'c1'},
+            {id: 'c2'},
+          ],
+        };
+        var diffCommentsStub = sinon.stub(element, '_getDiffComments',
+            function() { return Promise.resolve(comments); });
 
-    test('renders image diffs', function(done) {
-      var mockDiff = {
-        meta_a: {name: 'carrot.jpg', content_type: 'image/jpeg', lines: 66},
-        meta_b: {name: 'carrot.jpg', content_type: 'image/jpeg', lines: 560},
-        intraline_status: 'OK',
-        change_type: 'MODIFIED',
-        diff_header: [
-          'diff --git a/carrot.jpg b/carrot.jpg',
-          'index 2adc47d..f9c2f2c 100644',
-          '--- a/carrot.jpg',
-          '+++ b/carrot.jpg',
-          'Binary files differ',
-        ],
-        content: [{skip: 66}],
-        binary: true,
-      };
-      var mockFile1 = {
-        body: 'Qk06AAAAAAAAADYAAAAoAAAAAQAAAP////8BACAAAAAAAAAAAAATCwAAEwsAAA' +
-            'AAAAAAAAAAAAAA/w==',
-        type: 'image/bmp',
-      };
-      var mockFile2 = {
-        body: 'Qk06AAAAAAAAADYAAAAoAAAAAQAAAP////8BACAAAAAAAAAAAAATCwAAEwsAAA' +
-            'AAAAAAAAAA/////w==',
-        type: 'image/bmp'
-      };
-      var mockCommit = {
-        commit: '9a1a1d10baece5efbba10bc4ccf808a67a50ac0a',
-        parents: [{
-          commit: '7338aa9adfe57909f1fdaf88975cdea467d3382f',
-          subject: 'Added a carrot',
-        }],
-        author: {
-          name: 'Wyatt Allen',
-          email: 'wyatta@google.com',
-          date: '2016-05-23 21:44:51.000000000',
-          tz: -420,
-        },
-        committer: {
-          name: 'Wyatt Allen',
-          email: 'wyatta@google.com',
-          date: '2016-05-25 00:25:41.000000000',
-          tz: -420,
-        },
-        subject: 'Updated the carrot',
-        message: 'Updated the carrot\n\nChange-Id: Iabcd123\n',
-      };
-      var mockComments = {baseComments: [], comments: []};
+        var drafts = {
+          baseComments: [
+            {id: 'bd1'},
+            {id: 'bd2'},
+          ],
+          comments: [
+            {id: 'd1'},
+            {id: 'd2'},
+          ],
+        };
+        var diffDraftsStub = sinon.stub(element, '_getDiffDrafts',
+            function() { return Promise.resolve(drafts); });
 
-      var stubs = [];
-      stubs.push(sinon.stub(element, '_getDiff',
-          function() { return Promise.resolve(mockDiff); }));
-      stubs.push(sinon.stub(element.$.restAPI, 'getCommitInfo',
-          function() { return Promise.resolve(mockCommit); }));
-      stubs.push(sinon.stub(element.$.restAPI,
-          'getCommitFileContents',
-          function() { return Promise.resolve(mockFile1); }));
-      stubs.push(sinon.stub(element.$.restAPI,
-          'getChangeFileContents',
-          function() { return Promise.resolve(mockFile2); }));
-      stubs.push(sinon.stub(element.$.restAPI, '_getDiffComments',
-          function() { return Promise.resolve(mockComments); }));
-      stubs.push(sinon.stub(element.$.restAPI, 'getDiffDrafts',
-          function() { return Promise.resolve(mockComments); }));
+        element.changeNum = '42';
+        element.patchRange = {
+          basePatchNum: 'PARENT',
+          patchNum: 3,
+        };
+        element.path = '/path/to/foo';
+        element.projectConfig = {foo: 'bar'};
 
-      element.patchRange = {basePatchNum: 'PARENT', patchNum: 1};
+        element._getDiffCommentsAndDrafts().then(function(result) {
+          assert.deepEqual(result, {
+            meta: {
+              changeNum: '42',
+              patchRange: {
+                basePatchNum: 'PARENT',
+                patchNum: 3,
+              },
+              path: '/path/to/foo',
+              projectConfig: {foo: 'bar'},
+            },
+            left: [
+              {id: 'bc1'},
+              {id: 'bc2'},
+              {id: 'bd1', __draft: true},
+              {id: 'bd2', __draft: true},
+            ],
+            right: [
+              {id: 'c1'},
+              {id: 'c2'},
+              {id: 'd1', __draft: true},
+              {id: 'd2', __draft: true},
+            ],
+          });
 
-      var rendered = function() {
-        // Recognizes that it should be an image diff.
-        assert.isTrue(element.isImageDiff);
-        assert.instanceOf(element._getDiffBuilder(element._diff,
-            element._comments, element.prefs), GrDiffBuilderImage);
+          diffCommentsStub.restore();
+          diffDraftsStub.restore();
+          done();
+        });
+      });
 
-        // The left image rendered with the parent commit's version of the file.
-        var leftInmage = element.$.diffTable.querySelector('td.left img');
-        assert.isOk(leftInmage);
-        assert.equal(leftInmage.getAttribute('src'),
-            'data:image/bmp;base64, ' + mockFile1.body);
+      suite('handle comment-update', function() {
 
-        // The right image rendered with this change's revision of the image.
-        var rightInmage = element.$.diffTable.querySelector('td.right img');
-        assert.isOk(rightInmage);
-        assert.equal(rightInmage.getAttribute('src'),
-            'data:image/bmp;base64, ' + mockFile2.body);
+        setup(function() {
+          element._comments = {
+            meta: {
+              changeNum: '42',
+              patchRange: {
+                basePatchNum: 'PARENT',
+                patchNum: 3,
+              },
+              path: '/path/to/foo',
+              projectConfig: {foo: 'bar'},
+            },
+            left: [
+              {id: 'bc1', side: 'PARENT'},
+              {id: 'bc2', side: 'PARENT'},
+              {id: 'bd1', __draft: true, side: 'PARENT'},
+              {id: 'bd2', __draft: true, side: 'PARENT'},
+            ],
+            right: [
+              {id: 'c1'},
+              {id: 'c2'},
+              {id: 'd1', __draft: true},
+              {id: 'd2', __draft: true},
+            ],
+          };
+        });
 
-        // Cleanup.
-        element.removeEventListener('render', rendered);
-        stubs.forEach(function(stub) { stub.restore(); });
+        test('creating a draft', function() {
+          var comment = {__draft: true, __draftID: 'tempID', side: 'PARENT'};
+          element.fire('comment-update', {comment: comment});
+          assert.include(element._comments.left, comment);
+        });
 
-        done();
-      };
-
-      element.addEventListener('render', rendered);
-
-      element.$.restAPI.getDiffPreferences().then(function(prefs) {
-        element.prefs = prefs;
-        element.reload();
+        test('saving a draft', function() {
+          var draftID = 'tempID';
+          var id = 'savedID';
+          element._comments.left.push(
+              {__draft: true, __draftID: draftID, side: 'PARENT'});
+          element.fire('comment-update', {comment:
+              {id: id, __draft: true, __draftID: draftID, side: 'PARENT'},
+          });
+          var drafts = element._comments.left.filter(function(item) {
+            return item.__draftID === draftID;
+          });
+          assert.equal(drafts.length, 1);
+          assert.equal(drafts[0].id, id);
+        });
       });
     });