3 files changed, 189 insertions, 1 deletions
diff --git a/spec/frontend/notebook/cells/output/dataframe_spec.js b/spec/frontend/notebook/cells/output/dataframe_spec.js
new file mode 100644
index 00000000000..abf6631353c
--- /dev/null
+++ b/spec/frontend/notebook/cells/output/dataframe_spec.js
@@ -0,0 +1,59 @@
+import { shallowMount } from '@vue/test-utils';
+import DataframeOutput from '~/notebook/cells/output/dataframe.vue';
+import JSONTable from '~/behaviors/components/json_table.vue';
+import { outputWithDataframe } from '../../mock_data';
+
+describe('~/notebook/cells/output/DataframeOutput', () => {
+  let wrapper;
+
+  function createComponent(rawCode) {
+    wrapper = shallowMount(DataframeOutput, {
+      propsData: {
+        rawCode,
+        count: 0,
+        index: 0,
+      },
+    });
+  }
+
+  const findTable = () => wrapper.findComponent(JSONTable);
+
+  describe('with valid dataframe', () => {
+    beforeEach(() => createComponent(outputWithDataframe.data['text/html'].join('')));
+
+    it('mounts the table', () => {
+      expect(findTable().exists()).toBe(true);
+    });
+
+    it('table caption is empty', () => {
+      expect(findTable().props().caption).toEqual('');
+    });
+
+    it('allows filtering', () => {
+      expect(findTable().props().hasFilter).toBe(true);
+    });
+
+    it('sets the correct fields', () => {
+      expect(findTable().props().fields).toEqual([
+        { key: 'index', label: '', sortable: true },
+        { key: 'column_1', label: 'column_1', sortable: true },
+        { key: 'column_2', label: 'column_2', sortable: true },
+      ]);
+    });
+
+    it('sets the correct items', () => {
+      expect(findTable().props().items).toEqual([
+        { index: 0, column_1: 'abc de f', column_2: 'a' },
+        { index: 1, column_1: 'True', column_2: '0.1' },
+      ]);
+    });
+  });
+
+  describe('invalid dataframe', () => {
+    it('still displays the table', () => {
+      createComponent('dataframe');
+
+      expect(findTable().exists()).toBe(true);
+    });
+  });
+});
diff --git a/spec/frontend/notebook/cells/output/dataframe_util_spec.js b/spec/frontend/notebook/cells/output/dataframe_util_spec.js
new file mode 100644
index 00000000000..ddc1b3cfe26
--- /dev/null
+++ b/spec/frontend/notebook/cells/output/dataframe_util_spec.js
@@ -0,0 +1,113 @@
+import { isDataframe, convertHtmlTableToJson } from '~/notebook/cells/output/dataframe_util';
+import { outputWithDataframeContent } from '../../mock_data';
+import sanitizeTests from './html_sanitize_fixtures';
+
+describe('notebook/cells/output/dataframe_utils', () => {
+  describe('isDataframe', () => {
+    describe('when output data has no text/html', () => {
+      it('is is not a dataframe', () => {
+        const input = { data: { 'image/png': ['blah'] } };
+
+        expect(isDataframe(input)).toBe(false);
+      });
+    });
+
+    describe('when output data has no text/html, but no mention of dataframe', () => {
+      it('is is not a dataframe', () => {
+        const input = { data: { 'text/html': ['blah'] } };
+
+        expect(isDataframe(input)).toBe(false);
+      });
+    });
+
+    describe('when output data has text/html, but no mention of dataframe in the first 20 lines', () => {
+      it('is is not a dataframe', () => {
+        const input = { data: { 'text/html': [...new Array(20).fill('a'), 'dataframe'] } };
+
+        expect(isDataframe(input)).toBe(false);
+      });
+    });
+
+    describe('when output data has text/html, and includes "dataframe" within the first 20 lines', () => {
+      it('is is not a dataframe', () => {
+        const input = { data: { 'text/html': ['dataframe'] } };
+
+        expect(isDataframe(input)).toBe(true);
+      });
+    });
+  });
+
+  describe('convertHtmlTableToJson', () => {
+    it('converts table correctly', () => {
+      const input = outputWithDataframeContent;
+
+      const output = {
+        fields: [
+          { key: 'index', label: '', sortable: true },
+          { key: 'column_1', label: 'column_1', sortable: true },
+          { key: 'column_2', label: 'column_2', sortable: true },
+        ],
+        items: [
+          { index: 0, column_1: 'abc de f', column_2: 'a' },
+          { index: 1, column_1: 'True', column_2: '0.1' },
+        ],
+      };
+
+      expect(convertHtmlTableToJson(input)).toEqual(output);
+    });
+
+    describe('sanitizes input before parsing table', () => {
+      it('sanitizes input html', () => {
+        const parser = new DOMParser();
+        const spy = jest.spyOn(parser, 'parseFromString');
+        const input = 'hello<style>p {width:50%;}</style><script>alert(1)</script>';
+
+        convertHtmlTableToJson(input, parser);
+
+        expect(spy).toHaveBeenCalledWith('hello', 'text/html');
+      });
+    });
+
+    describe('does not include harmful html', () => {
+      const makeDataframeWithHtml = (html) => {
+        return [
+          '<table border="1" class="dataframe">\n',
+          '  <thead>\n',
+          '    <tr style="text-align: right;">\n',
+          '      <th></th>\n',
+          '      <th>column_1</th>\n',
+          '    </tr>\n',
+          '  </thead>\n',
+          '  <tbody>\n',
+          '    <tr>\n',
+          '      <th>0</th>\n',
+          `      <td>${html}</td>\n`,
+          '    </tr>\n',
+          '  </tbody>\n',
+          '</table>\n',
+          '</div>',
+        ];
+      };
+
+      it.each([
+        ['table', 0],
+        ['style', 1],
+        ['iframe', 2],
+        ['svg', 3],
+      ])('sanitizes output for: %p', (tag, index) => {
+        const inputHtml = makeDataframeWithHtml(sanitizeTests[index][1].input);
+        const convertedHtml = convertHtmlTableToJson(inputHtml).items[0].column_1;
+
+        expect(convertedHtml).not.toContain(tag);
+      });
+    });
+
+    describe('when dataframe is invalid', () => {
+      it('returns empty', () => {
+        const input = [' dataframe', ' blah'];
+
+        expect(convertHtmlTableToJson(input)).toEqual({ fields: [], items: [] });
+      });
+    });
+  });
+});
diff --git a/spec/frontend/notebook/cells/output/index_spec.js b/spec/frontend/notebook/cells/output/index_spec.js
index 1241c133b89..efbdfca8d8c 100644
--- a/spec/frontend/notebook/cells/output/index_spec.js
+++ b/spec/frontend/notebook/cells/output/index_spec.js
@@ -2,7 +2,13 @@ import { mount } from '@vue/test-utils';
 import json from 'test_fixtures/blob/notebook/basic.json';
 import Output from '~/notebook/cells/output/index.vue';
 import MarkdownOutput from '~/notebook/cells/output/markdown.vue';
-import { relativeRawPath, markdownCellContent } from '../../mock_data';
+import DataframeOutput from '~/notebook/cells/output/dataframe.vue';
+import {
+  relativeRawPath,
+  markdownCellContent,
+  outputWithDataframe,
+  outputWithDataframeContent,
+} from '../../mock_data';
 
 describe('Output component', () => {
   let wrapper;
@@ -105,6 +111,16 @@ describe('Output component', () => {
     });
   });
 
+  describe('Dataframe output', () => {
+    it('renders DataframeOutput component', () => {
+      createComponent(outputWithDataframe);
+
+      expect(wrapper.findComponent(DataframeOutput).props('rawCode')).toBe(
+        outputWithDataframeContent.join(''),
+      );
+    });
+  });
+
   describe('default to plain text', () => {
     beforeEach(() => {
       const unknownType = json.cells[6];