diff --git a/polygerrit-ui/app/elements/diff/gr-diff-host/gr-diff-host_test.html b/polygerrit-ui/app/elements/diff/gr-diff-host/gr-diff-host_test.html
index 88d7ab6..5344256 100644
--- a/polygerrit-ui/app/elements/diff/gr-diff-host/gr-diff-host_test.html
+++ b/polygerrit-ui/app/elements/diff/gr-diff-host/gr-diff-host_test.html
@@ -52,6 +52,189 @@
       sandbox.restore();
     });
 
+    suite('handle comment-update', () => {
+      setup(() => {
+        sandbox.stub(element, '_commentsChanged');
+        element.comments = {
+          meta: {
+            changeNum: '42',
+            patchRange: {
+              basePatchNum: 'PARENT',
+              patchNum: 3,
+            },
+            path: '/path/to/foo',
+            projectConfig: {foo: 'bar'},
+          },
+          left: [
+            {id: 'bc1', side: 'PARENT', __commentSide: 'left'},
+            {id: 'bc2', side: 'PARENT', __commentSide: 'left'},
+            {id: 'bd1', __draft: true, side: 'PARENT', __commentSide: 'left'},
+            {id: 'bd2', __draft: true, side: 'PARENT', __commentSide: 'left'},
+          ],
+          right: [
+            {id: 'c1', __commentSide: 'right'},
+            {id: 'c2', __commentSide: 'right'},
+            {id: 'd1', __draft: true, __commentSide: 'right'},
+            {id: 'd2', __draft: true, __commentSide: 'right'},
+          ],
+        };
+      });
+
+      test('creating a draft', () => {
+        const comment = {__draft: true, __draftID: 'tempID', side: 'PARENT',
+          __commentSide: 'left'};
+        element.fire('comment-update', {comment});
+        assert.include(element.comments.left, comment);
+      });
+
+      test('discarding a draft', () => {
+        const draftID = 'tempID';
+        const id = 'savedID';
+        const comment = {
+          __draft: true,
+          __draftID: draftID,
+          side: 'PARENT',
+          __commentSide: 'left',
+        };
+        const diffCommentsModifiedStub = sandbox.stub();
+        element.addEventListener('diff-comments-modified',
+            diffCommentsModifiedStub);
+        element.comments.left.push(comment);
+        comment.id = id;
+        element.fire('comment-discard', {comment});
+        const drafts = element.comments.left.filter(item => {
+          return item.__draftID === draftID;
+        });
+        assert.equal(drafts.length, 0);
+        assert.isTrue(diffCommentsModifiedStub.called);
+      });
+
+      test('saving a draft', () => {
+        const draftID = 'tempID';
+        const id = 'savedID';
+        const comment = {
+          __draft: true,
+          __draftID: draftID,
+          side: 'PARENT',
+          __commentSide: 'left',
+        };
+        const diffCommentsModifiedStub = sandbox.stub();
+        element.addEventListener('diff-comments-modified',
+            diffCommentsModifiedStub);
+        element.comments.left.push(comment);
+        comment.id = id;
+        element.fire('comment-save', {comment});
+        const drafts = element.comments.left.filter(item => {
+          return item.__draftID === draftID;
+        });
+        assert.equal(drafts.length, 1);
+        assert.equal(drafts[0].id, id);
+        assert.isTrue(diffCommentsModifiedStub.called);
+      });
+    });
+
+    test('remove comment', () => {
+      sandbox.stub(element, '_commentsChanged');
+      element.comments = {
+        meta: {
+          changeNum: '42',
+          patchRange: {
+            basePatchNum: 'PARENT',
+            patchNum: 3,
+          },
+          path: '/path/to/foo',
+          projectConfig: {foo: 'bar'},
+        },
+        left: [
+          {id: 'bc1', side: 'PARENT', __commentSide: 'left'},
+          {id: 'bc2', side: 'PARENT', __commentSide: 'left'},
+          {id: 'bd1', __draft: true, side: 'PARENT', __commentSide: 'left'},
+          {id: 'bd2', __draft: true, side: 'PARENT', __commentSide: 'left'},
+        ],
+        right: [
+          {id: 'c1', __commentSide: 'right'},
+          {id: 'c2', __commentSide: 'right'},
+          {id: 'd1', __draft: true, __commentSide: 'right'},
+          {id: 'd2', __draft: true, __commentSide: 'right'},
+        ],
+      };
+
+      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', __commentSide: 'left'},
+          {id: 'bc2', side: 'PARENT', __commentSide: 'left'},
+          {id: 'bd1', __draft: true, side: 'PARENT', __commentSide: 'left'},
+          {id: 'bd2', __draft: true, side: 'PARENT', __commentSide: 'left'},
+        ],
+        right: [
+          {id: 'c1', __commentSide: 'right'},
+          {id: 'c2', __commentSide: 'right'},
+          {id: 'd1', __draft: true, __commentSide: 'right'},
+          {id: 'd2', __draft: true, __commentSide: 'right'},
+        ],
+      }));
+
+      element._removeComment({id: 'bc2', side: 'PARENT',
+        __commentSide: 'left'});
+      assert.deepEqual(element.comments, {
+        meta: {
+          changeNum: '42',
+          patchRange: {
+            basePatchNum: 'PARENT',
+            patchNum: 3,
+          },
+          path: '/path/to/foo',
+          projectConfig: {foo: 'bar'},
+        },
+        left: [
+          {id: 'bc1', side: 'PARENT', __commentSide: 'left'},
+          {id: 'bd1', __draft: true, side: 'PARENT', __commentSide: 'left'},
+          {id: 'bd2', __draft: true, side: 'PARENT', __commentSide: 'left'},
+        ],
+        right: [
+          {id: 'c1', __commentSide: 'right'},
+          {id: 'c2', __commentSide: 'right'},
+          {id: 'd1', __draft: true, __commentSide: 'right'},
+          {id: 'd2', __draft: true, __commentSide: 'right'},
+        ],
+      });
+
+      element._removeComment({id: 'd2', __commentSide: 'right'});
+      assert.deepEqual(element.comments, {
+        meta: {
+          changeNum: '42',
+          patchRange: {
+            basePatchNum: 'PARENT',
+            patchNum: 3,
+          },
+          path: '/path/to/foo',
+          projectConfig: {foo: 'bar'},
+        },
+        left: [
+          {id: 'bc1', side: 'PARENT', __commentSide: 'left'},
+          {id: 'bd1', __draft: true, side: 'PARENT', __commentSide: 'left'},
+          {id: 'bd2', __draft: true, side: 'PARENT', __commentSide: 'left'},
+        ],
+        right: [
+          {id: 'c1', __commentSide: 'right'},
+          {id: 'c2', __commentSide: 'right'},
+          {id: 'd1', __draft: true, __commentSide: 'right'},
+        ],
+      });
+    });
+
     test('thread-discard handling', () => {
       const threads = [
         {comments: [{id: 4711}]},
@@ -682,12 +865,6 @@
       assert.equal(element.$.diff.noRenderOnPrefsChange, value);
     });
 
-    test('passes in comments', () => {
-      const value = {left: [], right: []};
-      element.comments = value;
-      assert.equal(element.$.diff.comments, value);
-    });
-
     test('passes in lineWrapping', () => {
       const value = true;
       element.lineWrapping = value;