From 6d6f304d68bbab0859b54d58dfa846a9e85967e8 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 7 May 2019 11:57:37 +0200
Subject: [PATCH 01/41] add nativerenderer component

---
 x-pack/plugins/lens/public/app_plugin/app.tsx |   9 +-
 .../utils/native_renderer/__mocks__/react.ts  |  12 +++
 .../native_renderer/native_renderer.test.tsx  | 100 ++++++++++++++++++
 .../utils/native_renderer/native_renderer.tsx |  65 ++++++++++++
 4 files changed, 179 insertions(+), 7 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts
 create mode 100644 x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
 create mode 100644 x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx

diff --git a/x-pack/plugins/lens/public/app_plugin/app.tsx b/x-pack/plugins/lens/public/app_plugin/app.tsx
index 3e16a78083092..819109ea6d276 100644
--- a/x-pack/plugins/lens/public/app_plugin/app.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/app.tsx
@@ -8,20 +8,15 @@ import { I18nProvider } from '@kbn/i18n/react';
 import React, { useCallback } from 'react';
 
 import { EditorFrameSetup } from '../types';
+import { NativeRenderer } from '../utils/native_renderer/native_renderer';
 
 export function App({ editorFrame }: { editorFrame: EditorFrameSetup }) {
-  const renderFrame = useCallback(node => {
-    if (node !== null) {
-      editorFrame.render(node);
-    }
-  }, []);
-
   return (
     <I18nProvider>
       <div>
         <h1>Lens</h1>
 
-        <div ref={renderFrame} />
+        <NativeRenderer render={editorFrame.render} actualProps={undefined} />
       </div>
     </I18nProvider>
   );
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts b/x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts
new file mode 100644
index 0000000000000..b09289b0516ce
--- /dev/null
+++ b/x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts
@@ -0,0 +1,12 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+
+// This mock overwrites react to use the layout effect hook instead of the normal one.
+// This is done to have effects executed synchronously to be able to test them without
+// setTimeout hacks.
+export default { ...React, useEffect: React.useLayoutEffect };
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
new file mode 100644
index 0000000000000..1340a477d570c
--- /dev/null
+++ b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
@@ -0,0 +1,100 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+import { render } from 'react-dom';
+import { NativeRenderer } from './native_renderer';
+
+describe('agnostic_renderer', () => {
+  let mountpoint: Element;
+
+  beforeEach(() => {
+    mountpoint = document.createElement('div');
+  });
+
+  afterEach(() => {
+    mountpoint.remove();
+  });
+
+  it('should render element in container', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).toHaveBeenCalledWith(containerElement, testProps);
+  });
+
+  it('should not render again if props do not change', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('should not render again if props do not change shallowly', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} actualProps={{ ...testProps }} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('should render again once if props change', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} actualProps={{ a: 'def' }} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} actualProps={{ a: 'def' }} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(2);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).lastCalledWith(containerElement, { a: 'def' });
+  });
+
+  it('should render again if props are extended', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} actualProps={{ a: 'abc', b: 'def' }} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(2);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).lastCalledWith(containerElement, { a: 'abc', b: 'def' });
+  });
+
+  it('should render again if props are limited', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc', b: 'def' };
+
+    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} actualProps={{ a: 'abc' }} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(2);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).lastCalledWith(containerElement, { a: 'abc' });
+  });
+
+  it('should render a div as container', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    const containerElement: Element = mountpoint.firstElementChild!;
+    expect(containerElement.nodeName).toBe('DIV');
+  });
+
+  it('should render a specified element as container', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} tag="span" actualProps={testProps} />, mountpoint);
+    const containerElement: Element = mountpoint.firstElementChild!;
+    expect(containerElement.nodeName).toBe('SPAN');
+  });
+});
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
new file mode 100644
index 0000000000000..46b0f78b0539a
--- /dev/null
+++ b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
@@ -0,0 +1,65 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import PropTypes from 'prop-types';
+import React, { useEffect, useRef } from 'react';
+
+export interface NativeRendererProps<T> {
+  render: (domElement: Element, props: T) => void;
+  actualProps: T;
+  tag?: string;
+  children?: never;
+}
+
+function isShallowDifferent<T>(a: T, b: T): boolean {
+  for (const key in a) {
+    if (!(key in b) || a[key] !== b[key]) {
+      return true;
+    }
+  }
+  for (const key in b) {
+    if (!(key in a) || a[key] !== b[key]) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+export function NativeRenderer<T>({ render, actualProps, tag }: NativeRendererProps<T>) {
+  const elementRef = useRef<Element | null>(null);
+  const propsRef = useRef<T | null>(null);
+
+  function renderAndUpdate(element: Element) {
+        elementRef.current = element;
+        propsRef.current = actualProps;
+        render(element, actualProps);
+  }
+
+  useEffect(
+    () => {
+      if (elementRef.current && isShallowDifferent(propsRef.current, actualProps)) {
+        renderAndUpdate(elementRef.current);
+      }
+    },
+    [actualProps]
+  );
+
+  return React.createElement(tag || 'div', {
+    ref: element => {
+      if (element && element !== elementRef.current) {
+        renderAndUpdate(element);
+      }
+    },
+  });
+}
+
+NativeRenderer.propTypes = {
+  render: PropTypes.func.isRequired,
+  actualProps: PropTypes.object,
+  
+  tag: PropTypes.string
+};
\ No newline at end of file

From 44894dfb69d79a99d4fef5c66865957ca096e891 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 7 May 2019 12:42:27 +0200
Subject: [PATCH 02/41] use native renderer in app and editor frame

---
 x-pack/plugins/lens/public/app_plugin/app.tsx      |  2 +-
 .../public/editor_frame_plugin/editor_frame.tsx    | 14 ++++++--------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/x-pack/plugins/lens/public/app_plugin/app.tsx b/x-pack/plugins/lens/public/app_plugin/app.tsx
index 819109ea6d276..d48badab71016 100644
--- a/x-pack/plugins/lens/public/app_plugin/app.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/app.tsx
@@ -5,7 +5,7 @@
  */
 
 import { I18nProvider } from '@kbn/i18n/react';
-import React, { useCallback } from 'react';
+import React from 'react';
 
 import { EditorFrameSetup } from '../types';
 import { NativeRenderer } from '../utils/native_renderer/native_renderer';
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
index 091862e635d99..6dcc14542e21c 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
@@ -6,6 +6,7 @@
 
 import React from 'react';
 import { Datasource, Visualization } from '../types';
+import { NativeRenderer } from '../utils/native_renderer/native_renderer';
 
 interface EditorFrameProps {
   datasources: { [key: string]: Datasource };
@@ -20,15 +21,12 @@ export function EditorFrame(props: EditorFrameProps) {
       <h2>Editor Frame</h2>
 
       {keys.map(key => (
-        <div
+        <NativeRenderer
           key={key}
-          ref={domElement => {
-            if (domElement) {
-              props.datasources[key].renderDataPanel(domElement, {
-                state: {},
-                setState: () => {},
-              });
-            }
+          render={props.datasources[key].renderDataPanel}
+          actualProps={{
+            state: {},
+            setState: () => {},
           }}
         />
       ))}

From 943fca82e2cc9132199a1d83c7c13c1860c0a537 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 7 May 2019 13:44:40 +0200
Subject: [PATCH 03/41] remove prop types

---
 .../utils/native_renderer/native_renderer.tsx      | 14 +++-----------
 1 file changed, 3 insertions(+), 11 deletions(-)

diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
index 46b0f78b0539a..936433e4b2c78 100644
--- a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
+++ b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
@@ -4,7 +4,6 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import PropTypes from 'prop-types';
 import React, { useEffect, useRef } from 'react';
 
 export interface NativeRendererProps<T> {
@@ -34,9 +33,9 @@ export function NativeRenderer<T>({ render, actualProps, tag }: NativeRendererPr
   const propsRef = useRef<T | null>(null);
 
   function renderAndUpdate(element: Element) {
-        elementRef.current = element;
-        propsRef.current = actualProps;
-        render(element, actualProps);
+    elementRef.current = element;
+    propsRef.current = actualProps;
+    render(element, actualProps);
   }
 
   useEffect(
@@ -56,10 +55,3 @@ export function NativeRenderer<T>({ render, actualProps, tag }: NativeRendererPr
     },
   });
 }
-
-NativeRenderer.propTypes = {
-  render: PropTypes.func.isRequired,
-  actualProps: PropTypes.object,
-  
-  tag: PropTypes.string
-};
\ No newline at end of file

From 731e2647996ba5059b2928711a4adedfe58527a6 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Wed, 8 May 2019 08:01:38 +0200
Subject: [PATCH 04/41] add unit test and fix linting issues

---
 x-pack/plugins/lens/public/index.ts           |  2 +-
 .../utils/native_renderer/__mocks__/react.ts  |  2 ++
 .../native_renderer/native_renderer.test.tsx  | 22 +++++++++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/x-pack/plugins/lens/public/index.ts b/x-pack/plugins/lens/public/index.ts
index 874249ff75d85..f8f074cdb99bf 100644
--- a/x-pack/plugins/lens/public/index.ts
+++ b/x-pack/plugins/lens/public/index.ts
@@ -6,8 +6,8 @@
 
 export * from './types';
 
-import { IScope } from 'angular';
 import { render, unmountComponentAtNode } from 'react-dom';
+import { IScope } from 'angular';
 import chrome from 'ui/chrome';
 import { appSetup, appStop } from './app_plugin';
 
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts b/x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts
index b09289b0516ce..6039f0888f596 100644
--- a/x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts
+++ b/x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts
@@ -9,4 +9,6 @@ import React from 'react';
 // This mock overwrites react to use the layout effect hook instead of the normal one.
 // This is done to have effects executed synchronously to be able to test them without
 // setTimeout hacks.
+
+// eslint-disable-next-line import/no-default-export
 export default { ...React, useEffect: React.useLayoutEffect };
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
index 1340a477d570c..9ff986a4707aa 100644
--- a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
+++ b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
@@ -46,6 +46,28 @@ describe('agnostic_renderer', () => {
     expect(renderSpy).toHaveBeenCalledTimes(1);
   });
 
+  it('should not render again for unchanged callback functions', () => {
+    const renderSpy = jest.fn();
+    const testCallback = () => {};
+    const testState = { a: 'abc' };
+
+    render(
+      <NativeRenderer
+        render={renderSpy}
+        actualProps={{ state: testState, setState: testCallback }}
+      />,
+      mountpoint
+    );
+    render(
+      <NativeRenderer
+        render={renderSpy}
+        actualProps={{ state: testState, setState: testCallback }}
+      />,
+      mountpoint
+    );
+    expect(renderSpy).toHaveBeenCalledTimes(1);
+  });
+
   it('should render again once if props change', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };

From cc1a0dbaeac72b93140e33aa66fc73943e72256a Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Wed, 8 May 2019 15:17:26 +0200
Subject: [PATCH 05/41] rename test suite

---
 .../lens/public/utils/native_renderer/native_renderer.test.tsx  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
index 9ff986a4707aa..3af468e391668 100644
--- a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
+++ b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
@@ -8,7 +8,7 @@ import React from 'react';
 import { render } from 'react-dom';
 import { NativeRenderer } from './native_renderer';
 
-describe('agnostic_renderer', () => {
+describe('native_renderer', () => {
   let mountpoint: Element;
 
   beforeEach(() => {

From 18f9d838e5a358acad1ca2baac2a5681b1582ba7 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Thu, 9 May 2019 17:01:04 +0200
Subject: [PATCH 06/41] rename prop and adjust shallow comparison function

---
 x-pack/plugins/lens/public/app_plugin/app.tsx |  2 +-
 .../editor_frame_plugin/editor_frame.tsx      |  2 +-
 .../native_renderer/native_renderer.test.tsx  | 32 +++++++++----------
 .../utils/native_renderer/native_renderer.tsx | 25 ++++++++-------
 4 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/x-pack/plugins/lens/public/app_plugin/app.tsx b/x-pack/plugins/lens/public/app_plugin/app.tsx
index d48badab71016..ba426fa9685a9 100644
--- a/x-pack/plugins/lens/public/app_plugin/app.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/app.tsx
@@ -16,7 +16,7 @@ export function App({ editorFrame }: { editorFrame: EditorFrameSetup }) {
       <div>
         <h1>Lens</h1>
 
-        <NativeRenderer render={editorFrame.render} actualProps={undefined} />
+        <NativeRenderer render={editorFrame.render} nativeProps={undefined} />
       </div>
     </I18nProvider>
   );
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
index 6dcc14542e21c..bbfd2e6933e3c 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
@@ -24,7 +24,7 @@ export function EditorFrame(props: EditorFrameProps) {
         <NativeRenderer
           key={key}
           render={props.datasources[key].renderDataPanel}
-          actualProps={{
+          nativeProps={{
             state: {},
             setState: () => {},
           }}
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
index 3af468e391668..22c0a8d394043 100644
--- a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
+++ b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
@@ -23,7 +23,7 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).toHaveBeenCalledWith(containerElement, testProps);
   });
@@ -32,8 +32,8 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
     expect(renderSpy).toHaveBeenCalledTimes(1);
   });
 
@@ -41,8 +41,8 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} actualProps={{ ...testProps }} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ ...testProps }} />, mountpoint);
     expect(renderSpy).toHaveBeenCalledTimes(1);
   });
 
@@ -54,14 +54,14 @@ describe('native_renderer', () => {
     render(
       <NativeRenderer
         render={renderSpy}
-        actualProps={{ state: testState, setState: testCallback }}
+        nativeProps={{ state: testState, setState: testCallback }}
       />,
       mountpoint
     );
     render(
       <NativeRenderer
         render={renderSpy}
-        actualProps={{ state: testState, setState: testCallback }}
+        nativeProps={{ state: testState, setState: testCallback }}
       />,
       mountpoint
     );
@@ -72,9 +72,9 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} actualProps={{ a: 'def' }} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} actualProps={{ a: 'def' }} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'def' }} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'def' }} />, mountpoint);
     expect(renderSpy).toHaveBeenCalledTimes(2);
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).lastCalledWith(containerElement, { a: 'def' });
@@ -84,8 +84,8 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} actualProps={{ a: 'abc', b: 'def' }} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'abc', b: 'def' }} />, mountpoint);
     expect(renderSpy).toHaveBeenCalledTimes(2);
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).lastCalledWith(containerElement, { a: 'abc', b: 'def' });
@@ -95,8 +95,8 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc', b: 'def' };
 
-    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} actualProps={{ a: 'abc' }} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'abc' }} />, mountpoint);
     expect(renderSpy).toHaveBeenCalledTimes(2);
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).lastCalledWith(containerElement, { a: 'abc' });
@@ -106,7 +106,7 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
     const containerElement: Element = mountpoint.firstElementChild!;
     expect(containerElement.nodeName).toBe('DIV');
   });
@@ -115,7 +115,7 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} tag="span" actualProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} tag="span" nativeProps={testProps} />, mountpoint);
     const containerElement: Element = mountpoint.firstElementChild!;
     expect(containerElement.nodeName).toBe('SPAN');
   });
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
index 936433e4b2c78..c1af413c8ed29 100644
--- a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
+++ b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
@@ -8,43 +8,46 @@ import React, { useEffect, useRef } from 'react';
 
 export interface NativeRendererProps<T> {
   render: (domElement: Element, props: T) => void;
-  actualProps: T;
+  nativeProps: T;
   tag?: string;
   children?: never;
 }
 
 function isShallowDifferent<T>(a: T, b: T): boolean {
+  if (a === b) {
+    return false;
+  }
+
+  if (Object.keys(a).length !== Object.keys(b).length) {
+    return true;
+  }
+
   for (const key in a) {
     if (!(key in b) || a[key] !== b[key]) {
       return true;
     }
   }
-  for (const key in b) {
-    if (!(key in a) || a[key] !== b[key]) {
-      return true;
-    }
-  }
 
   return false;
 }
 
-export function NativeRenderer<T>({ render, actualProps, tag }: NativeRendererProps<T>) {
+export function NativeRenderer<T>({ render, nativeProps, tag }: NativeRendererProps<T>) {
   const elementRef = useRef<Element | null>(null);
   const propsRef = useRef<T | null>(null);
 
   function renderAndUpdate(element: Element) {
     elementRef.current = element;
-    propsRef.current = actualProps;
-    render(element, actualProps);
+    propsRef.current = nativeProps;
+    render(element, nativeProps);
   }
 
   useEffect(
     () => {
-      if (elementRef.current && isShallowDifferent(propsRef.current, actualProps)) {
+      if (elementRef.current && isShallowDifferent(propsRef.current, nativeProps)) {
         renderAndUpdate(elementRef.current);
       }
     },
-    [actualProps]
+    [nativeProps]
   );
 
   return React.createElement(tag || 'div', {

From 60cd8911a98b7be1a7b47b7a454b0a9201747c0c Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Thu, 9 May 2019 17:17:23 +0200
Subject: [PATCH 07/41] add documentation

---
 .../public/utils/native_renderer/native_renderer.tsx     | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
index c1af413c8ed29..f2693ae23c697 100644
--- a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
+++ b/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
@@ -31,6 +31,15 @@ function isShallowDifferent<T>(a: T, b: T): boolean {
   return false;
 }
 
+/**
+ * A component which takes care of providing a mountpoint for a generic render
+ * function which takes an html element and an optional props object.
+ * It also takes care of calling render again if the props object changes.
+ * By default the mountpoint element will be a div, this can be changed with the
+ * `tag` prop.
+ *
+ * @param props
+ */
 export function NativeRenderer<T>({ render, nativeProps, tag }: NativeRendererProps<T>) {
   const elementRef = useRef<Element | null>(null);
   const propsRef = useRef<T | null>(null);

From 0eb1f317c8247ef23d473a2fceda8668b0a2b2eb Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Thu, 9 May 2019 17:23:25 +0200
Subject: [PATCH 08/41] move native renderer to top level directory

---
 x-pack/plugins/lens/public/app_plugin/app.tsx              | 2 +-
 .../lens/public/editor_frame_plugin/editor_frame.tsx       | 2 +-
 .../public/{utils => }/native_renderer/__mocks__/react.ts  | 0
 x-pack/plugins/lens/public/native_renderer/index.ts        | 7 +++++++
 .../{utils => }/native_renderer/native_renderer.test.tsx   | 0
 .../public/{utils => }/native_renderer/native_renderer.tsx | 0
 6 files changed, 9 insertions(+), 2 deletions(-)
 rename x-pack/plugins/lens/public/{utils => }/native_renderer/__mocks__/react.ts (100%)
 create mode 100644 x-pack/plugins/lens/public/native_renderer/index.ts
 rename x-pack/plugins/lens/public/{utils => }/native_renderer/native_renderer.test.tsx (100%)
 rename x-pack/plugins/lens/public/{utils => }/native_renderer/native_renderer.tsx (100%)

diff --git a/x-pack/plugins/lens/public/app_plugin/app.tsx b/x-pack/plugins/lens/public/app_plugin/app.tsx
index ba426fa9685a9..29f1c2eec03f0 100644
--- a/x-pack/plugins/lens/public/app_plugin/app.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/app.tsx
@@ -8,7 +8,7 @@ import { I18nProvider } from '@kbn/i18n/react';
 import React from 'react';
 
 import { EditorFrameSetup } from '../types';
-import { NativeRenderer } from '../utils/native_renderer/native_renderer';
+import { NativeRenderer } from '../native_renderer';
 
 export function App({ editorFrame }: { editorFrame: EditorFrameSetup }) {
   return (
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
index bbfd2e6933e3c..c74a107b96313 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
@@ -6,7 +6,7 @@
 
 import React from 'react';
 import { Datasource, Visualization } from '../types';
-import { NativeRenderer } from '../utils/native_renderer/native_renderer';
+import { NativeRenderer } from '../native_renderer';
 
 interface EditorFrameProps {
   datasources: { [key: string]: Datasource };
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts b/x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts
similarity index 100%
rename from x-pack/plugins/lens/public/utils/native_renderer/__mocks__/react.ts
rename to x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts
diff --git a/x-pack/plugins/lens/public/native_renderer/index.ts b/x-pack/plugins/lens/public/native_renderer/index.ts
new file mode 100644
index 0000000000000..076e91d950118
--- /dev/null
+++ b/x-pack/plugins/lens/public/native_renderer/index.ts
@@ -0,0 +1,7 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+ export * from './native_renderer';
\ No newline at end of file
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx b/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
similarity index 100%
rename from x-pack/plugins/lens/public/utils/native_renderer/native_renderer.test.tsx
rename to x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
diff --git a/x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx b/x-pack/plugins/lens/public/native_renderer/native_renderer.tsx
similarity index 100%
rename from x-pack/plugins/lens/public/utils/native_renderer/native_renderer.tsx
rename to x-pack/plugins/lens/public/native_renderer/native_renderer.tsx

From a21b0eaadb129690375afadbd78528f9d4e3e68c Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Fri, 10 May 2019 10:02:11 +0200
Subject: [PATCH 09/41] add test, fix linting error and improve shallow equal
 comparison

---
 .../lens/public/native_renderer/index.ts      |  2 +-
 .../native_renderer/native_renderer.test.tsx  | 12 +++++++++++
 .../native_renderer/native_renderer.tsx       | 21 ++++++++++++++-----
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/x-pack/plugins/lens/public/native_renderer/index.ts b/x-pack/plugins/lens/public/native_renderer/index.ts
index 076e91d950118..0ef9bd8807bc5 100644
--- a/x-pack/plugins/lens/public/native_renderer/index.ts
+++ b/x-pack/plugins/lens/public/native_renderer/index.ts
@@ -4,4 +4,4 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
- export * from './native_renderer';
\ No newline at end of file
+export * from './native_renderer';
diff --git a/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx b/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
index 22c0a8d394043..089b1ed02d166 100644
--- a/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
+++ b/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
@@ -80,6 +80,18 @@ describe('native_renderer', () => {
     expect(renderSpy).lastCalledWith(containerElement, { a: 'def' });
   });
 
+  it('should render again once if props is just a string', () => {
+    const renderSpy = jest.fn();
+    const testProps = 'abc';
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps="def" />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps="def" />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(2);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).lastCalledWith(containerElement, 'def');
+  });
+
   it('should render again if props are extended', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
diff --git a/x-pack/plugins/lens/public/native_renderer/native_renderer.tsx b/x-pack/plugins/lens/public/native_renderer/native_renderer.tsx
index f2693ae23c697..f0eb4b829c153 100644
--- a/x-pack/plugins/lens/public/native_renderer/native_renderer.tsx
+++ b/x-pack/plugins/lens/public/native_renderer/native_renderer.tsx
@@ -13,17 +13,28 @@ export interface NativeRendererProps<T> {
   children?: never;
 }
 
-function isShallowDifferent<T>(a: T, b: T): boolean {
-  if (a === b) {
+function is(x: unknown, y: unknown) {
+  return (x === y && (x !== 0 || 1 / (x as number) === 1 / (y as number))) || (x !== x && y !== y);
+}
+
+function isShallowDifferent<T>(objA: T, objB: T): boolean {
+  if (is(objA, objB)) {
     return false;
   }
 
-  if (Object.keys(a).length !== Object.keys(b).length) {
+  if (typeof objA !== 'object' || objA === null || typeof objB !== 'object' || objB === null) {
+    return true;
+  }
+
+  const keysA = Object.keys(objA) as Array<keyof T>;
+  const keysB = Object.keys(objB) as Array<keyof T>;
+
+  if (keysA.length !== keysB.length) {
     return true;
   }
 
-  for (const key in a) {
-    if (!(key in b) || a[key] !== b[key]) {
+  for (let i = 0; i < keysA.length; i++) {
+    if (!window.hasOwnProperty.call(objB, keysA[i]) || !is(objA[keysA[i]], objB[keysA[i]])) {
       return true;
     }
   }

From 36de9426111d06c4cbc021ca9f1e9c66fa63a34a Mon Sep 17 00:00:00 2001
From: Wylie Conlon <wylieconlon@gmail.com>
Date: Fri, 3 May 2019 14:52:44 -0400
Subject: [PATCH 10/41] Implement basic editor frame state handling

---
 x-pack/plugins/lens/public/app_plugin/app.tsx |  15 +-
 .../plugins/lens/public/app_plugin/plugin.tsx |  13 +-
 .../editor_frame_plugin/editor_frame.test.tsx | 416 ++++++++++++++++++
 .../editor_frame_plugin/editor_frame.tsx      | 227 +++++-----
 .../public/editor_frame_plugin/plugin.tsx     |  48 +-
 .../state_management.test.ts                  | 143 ++++++
 .../editor_frame_plugin/state_management.ts   |  98 +++++
 x-pack/plugins/lens/public/index.ts           |   2 +-
 .../indexpattern_plugin/indexpattern.tsx      |  18 +-
 .../public/native_renderer/__mocks__/react.ts |  14 +
 .../lens/public/native_renderer/index.ts      |   7 +
 .../native_renderer/native_renderer.test.tsx  | 134 ++++++
 .../native_renderer/native_renderer.tsx       |  80 ++++
 x-pack/plugins/lens/public/types.ts           |  14 +-
 14 files changed, 1085 insertions(+), 144 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts
 create mode 100644 x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts
 create mode 100644 x-pack/plugins/lens/public/native_renderer/index.ts
 create mode 100644 x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
 create mode 100644 x-pack/plugins/lens/public/native_renderer/native_renderer.tsx

diff --git a/x-pack/plugins/lens/public/app_plugin/app.tsx b/x-pack/plugins/lens/public/app_plugin/app.tsx
index 3e16a78083092..16860142e28d0 100644
--- a/x-pack/plugins/lens/public/app_plugin/app.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/app.tsx
@@ -5,23 +5,18 @@
  */
 
 import { I18nProvider } from '@kbn/i18n/react';
-import React, { useCallback } from 'react';
+import React from 'react';
 
-import { EditorFrameSetup } from '../types';
-
-export function App({ editorFrame }: { editorFrame: EditorFrameSetup }) {
-  const renderFrame = useCallback(node => {
-    if (node !== null) {
-      editorFrame.render(node);
-    }
-  }, []);
+import { EditorFrameInstance } from '../types';
+import { NativeRenderer } from '../native_renderer';
 
+export function App({ editorFrame }: { editorFrame: EditorFrameInstance }) {
   return (
     <I18nProvider>
       <div>
         <h1>Lens</h1>
 
-        <div ref={renderFrame} />
+        <NativeRenderer render={editorFrame.mount} nativeProps={undefined} />
       </div>
     </I18nProvider>
   );
diff --git a/x-pack/plugins/lens/public/app_plugin/plugin.tsx b/x-pack/plugins/lens/public/app_plugin/plugin.tsx
index 1a096d7c1326c..aeea52c23c0e7 100644
--- a/x-pack/plugins/lens/public/app_plugin/plugin.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/plugin.tsx
@@ -9,8 +9,12 @@ import { editorFrameSetup, editorFrameStop } from '../editor_frame_plugin';
 import { indexPatternDatasourceSetup, indexPatternDatasourceStop } from '../indexpattern_plugin';
 import { xyVisualizationSetup, xyVisualizationStop } from '../xy_visualization_plugin';
 import { App } from './app';
+import { EditorFrameInstance } from '../types';
 
 export class AppPlugin {
+
+  private instance: EditorFrameInstance | null = null;
+
   constructor() {}
 
   setup() {
@@ -23,10 +27,17 @@ export class AppPlugin {
     editorFrame.registerDatasource('indexpattern', indexPattern);
     editorFrame.registerVisualization('xy', xyVisualization);
 
-    return <App editorFrame={editorFrame} />;
+    this.instance = editorFrame.createInstance({});
+
+    return <App editorFrame={this.instance} />;
   }
 
   stop() {
+    if (this.instance) {
+      this.instance.unmount();
+    }
+
+    // TODO this will be handled by the plugin platform itself
     indexPatternDatasourceStop();
     xyVisualizationStop();
     editorFrameStop();
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx
new file mode 100644
index 0000000000000..5d9353113a1c9
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx
@@ -0,0 +1,416 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+import { mount, ReactWrapper } from 'enzyme';
+import { EditorFrame } from './editor_frame';
+import { Visualization, Datasource, VisualizationProps, DatasourcePublicAPI } from '../types';
+
+const nextTick = () => new Promise(resolve => setTimeout(resolve));
+
+describe('editor_frame', () => {
+  const getMockVisualization = () => ({
+    getMappingOfTableToRoles: jest.fn(),
+    getPersistableState: jest.fn(),
+    getSuggestions: jest.fn(),
+    initialize: jest.fn(),
+    renderConfigPanel: jest.fn(),
+    toExpression: jest.fn(),
+  });
+
+  const getMockDatasource = () => ({
+    getDatasourceSuggestionsForField: jest.fn(),
+    getDatasourceSuggestionsFromCurrentState: jest.fn(),
+    getPersistableState: jest.fn(),
+    getPublicAPI: jest.fn(),
+    initialize: jest.fn(() => Promise.resolve()),
+    renderDataPanel: jest.fn(),
+    toExpression: jest.fn(),
+  });
+
+  let mockVisualization: Visualization;
+  let mockDatasource: Datasource;
+
+  let mockVisualization2: Visualization;
+  let mockDatasource2: Datasource;
+
+  beforeEach(() => {
+    mockVisualization = getMockVisualization();
+    mockVisualization2 = getMockVisualization();
+
+    mockDatasource = getMockDatasource();
+    mockDatasource2 = getMockDatasource();
+  });
+
+  describe('initialization', () => {
+    it('should initialize initial datasource and visualization if present', () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      expect(mockVisualization.initialize).toHaveBeenCalled();
+      expect(mockDatasource.initialize).toHaveBeenCalled();
+    });
+
+    it('should not initialize datasource and visualization if no initial one is specificed', () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource={null}
+          initialVisualization={null}
+        />
+      );
+
+      expect(mockVisualization.initialize).not.toHaveBeenCalled();
+      expect(mockDatasource.initialize).not.toHaveBeenCalled();
+    });
+
+    it('should not render something before datasource is initialized', () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      expect(mockVisualization.renderConfigPanel).not.toHaveBeenCalled();
+      expect(mockDatasource.renderDataPanel).not.toHaveBeenCalled();
+    });
+
+    it('should render data panel after initialization is complete', async () => {
+      const initialState = {};
+      let databaseInitialized: ({}) => void;
+
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: {
+              ...mockDatasource,
+              initialize: () =>
+                new Promise(resolve => {
+                  databaseInitialized = resolve;
+                }),
+            },
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+      databaseInitialized!(initialState);
+
+      await nextTick();
+      expect(mockDatasource.renderDataPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({ state: initialState })
+      );
+    });
+
+    it('should initialize visualization state and render config panel', async () => {
+      const initialState = {};
+
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: { ...mockVisualization, initialize: () => initialState },
+          }}
+          datasources={{
+            testDatasource: {
+              ...mockDatasource,
+              initialize: () => Promise.resolve(),
+            },
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      expect(mockVisualization.renderConfigPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({ state: initialState })
+      );
+    });
+  });
+
+  describe('state update', () => {
+    it('should re-render config panel after state update', async () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      const updatedState = {};
+      const setVisualizationState = (mockVisualization.renderConfigPanel as jest.Mock).mock
+        .calls[0][1].setState;
+      setVisualizationState(updatedState);
+
+      expect(mockVisualization.renderConfigPanel).toHaveBeenCalledTimes(2);
+      expect(mockVisualization.renderConfigPanel).toHaveBeenLastCalledWith(
+        expect.any(Element),
+        expect.objectContaining({
+          state: updatedState,
+        })
+      );
+
+      // don't re-render datasource when visulization changes
+      expect(mockDatasource.renderDataPanel).toHaveBeenCalledTimes(1);
+    });
+    
+    it('should re-render data panel after state update', async () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      const updatedState = {};
+      const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock
+        .calls[0][1].setState;
+      setDatasourceState(updatedState);
+
+      expect(mockDatasource.renderDataPanel).toHaveBeenCalledTimes(2);
+      expect(mockDatasource.renderDataPanel).toHaveBeenLastCalledWith(
+        expect.any(Element),
+        expect.objectContaining({
+          state: updatedState,
+        })
+      );
+    });
+    
+    it('should re-render config panel with updated datasource api after datasource state update', async () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      const updatedPublicAPI = {};
+      mockDatasource.getPublicAPI = jest.fn(() => updatedPublicAPI as unknown as DatasourcePublicAPI);
+
+      const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock
+        .calls[0][1].setState;
+      setDatasourceState({});
+
+      expect(mockVisualization.renderConfigPanel).toHaveBeenCalledTimes(2);
+      expect(mockVisualization.renderConfigPanel).toHaveBeenLastCalledWith(
+        expect.any(Element),
+        expect.objectContaining({
+          datasource: updatedPublicAPI,
+        })
+      );
+    });
+  });
+
+  describe('datasource public api communication', () => {
+    it('should pass the datasource api to the visualization', async () => {
+      const publicAPI = ({} as unknown) as DatasourcePublicAPI;
+
+      mockDatasource.getPublicAPI = () => publicAPI;
+
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      expect(mockVisualization.renderConfigPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({ datasource: publicAPI })
+      );
+    });
+
+    it('should give access to the datasource state in the datasource factory function', async () => {
+      const datasourceState = {};
+      mockDatasource.initialize = () => Promise.resolve(datasourceState);
+
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      expect(mockDatasource.getPublicAPI).toHaveBeenCalledWith(
+        datasourceState,
+        expect.any(Function)
+      );
+    });
+
+    it('should re-create the public api after state has been set', async () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      const updatedState = {};
+      const setDatasourceState = (mockDatasource.getPublicAPI as jest.Mock).mock.calls[0][1];
+      setDatasourceState(updatedState);
+
+      expect(mockDatasource.getPublicAPI).toHaveBeenCalledTimes(2);
+      expect(mockDatasource.getPublicAPI).toHaveBeenLastCalledWith(
+        updatedState,
+        expect.any(Function)
+      );
+    });
+  });
+
+  describe('switching', () => {
+    let instance: ReactWrapper;
+    beforeEach(async () => {
+      instance = mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+            testVis2: mockVisualization2,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+            testDatasource2: mockDatasource2,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+      await nextTick();
+
+      // necessary to flush elements to dom synchronously
+      instance.update();
+    });
+
+    it('should have initialized only the initial datasource and visualization', () => {
+      expect(mockDatasource.initialize).toHaveBeenCalled();
+      expect(mockDatasource2.initialize).not.toHaveBeenCalled();
+
+      expect(mockVisualization.initialize).toHaveBeenCalled();
+      expect(mockVisualization2.initialize).not.toHaveBeenCalled();
+    });
+
+    it('should initialize other datasource on switch', async () => {
+      instance
+        .find('select[data-test-subj="datasource-switch"]')
+        .simulate('change', { target: { value: 'testDatasource2' } });
+      expect(mockDatasource2.initialize).toHaveBeenCalled();
+    });
+
+    it('should call datasource render with new state on switch', async () => {
+      const initialState = {};
+      mockDatasource2.initialize = () => Promise.resolve(initialState);
+
+      instance
+        .find('select[data-test-subj="datasource-switch"]')
+        .simulate('change', { target: { value: 'testDatasource2' } });
+
+      await nextTick();
+
+      expect(mockDatasource2.renderDataPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({ state: initialState })
+      );
+    });
+
+    it('should initialize other visualization on switch', async () => {
+      instance
+        .find('select[data-test-subj="visualization-switch"]')
+        .simulate('change', { target: { value: 'testVis2' } });
+      expect(mockVisualization2.initialize).toHaveBeenCalled();
+    });
+
+    it('should call visualization render with new state on switch', async () => {
+      const initialState = {};
+      mockVisualization2.initialize = () => initialState;
+
+      instance
+        .find('select[data-test-subj="visualization-switch"]')
+        .simulate('change', { target: { value: 'testVis2' } });
+
+      expect(mockVisualization2.renderConfigPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({ state: initialState })
+      );
+    });
+  });
+});
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
index 40b9ba40cd430..13e184ccb4d41 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
@@ -4,120 +4,149 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import React, { useReducer, useEffect } from 'react';
-import { Datasource, Visualization } from '../types';
+import React, { useEffect, useReducer, useMemo } from 'react';
+import { EuiSelect, EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
+import { Datasource, Visualization, DatasourceDataPanelProps } from '../types';
+import { NativeRenderer } from '../native_renderer';
+import { reducer, getInitialState } from './state_management';
 
-interface EditorFrameProps {
+export interface EditorFrameProps {
   datasources: { [key: string]: Datasource };
   visualizations: { [key: string]: Visualization };
 
-  initialDatasource?: string;
-}
-
-interface DatasourceState {
-  datasourceName: string;
-  visualizationName: string;
-
-  datasourceState: any;
-  visualizationState: any;
-}
-
-interface UpdateDatasourceAction {
-  type: 'UPDATE_DATASOURCE';
-  payload: any;
-}
-
-interface UpdateVisualizationAction {
-  type: 'UPDATE_VISUALIZATION';
-  payload: any;
-}
-
-type Action = UpdateDatasourceAction | UpdateVisualizationAction;
-
-function stateReducer(state: DatasourceState, action: Action): DatasourceState {
-  switch (action.type) {
-    case 'UPDATE_DATASOURCE':
-      return {
-        ...state,
-        datasourceState: action.payload,
-      };
-    case 'UPDATE_VISUALIZATION':
-      return {
-        ...state,
-        visualizationState: action.payload,
-      };
-  }
-  return state;
+  initialDatasource: string | null;
+  initialVisualization: string | null;
 }
 
 export function EditorFrame(props: EditorFrameProps) {
-  const dsKeys = Object.keys(props.datasources);
-  const vKeys = Object.keys(props.visualizations);
-
-  const [state, dispatch] = useReducer(stateReducer, {
-    datasourceName: props.initialDatasource || dsKeys[0],
-    visualizationName: vKeys[0],
+  const [state, dispatch] = useReducer(reducer, props, getInitialState);
+
+  // Initialize current datasource
+  useEffect(
+    () => {
+      let datasourceGotSwitched = false;
+      if (state.datasourceIsLoading && state.activeDatasource) {
+        props.datasources[state.activeDatasource].initialize().then(datasourceState => {
+          if (!datasourceGotSwitched) {
+            dispatch({
+              type: 'UPDATE_DATASOURCE_STATE',
+              newState: datasourceState,
+            });
+          }
+        });
+
+        return () => {
+          datasourceGotSwitched = true;
+        };
+      }
+    },
+    [state.activeDatasource, state.datasourceIsLoading]
+  );
 
-    datasourceState: null,
-    visualizationState: null,
-  });
+  const datasourceStateUpdater = useMemo(
+    () => (newState: unknown) => {
+      dispatch({
+        type: 'UPDATE_DATASOURCE_STATE',
+        newState,
+      });
+    },
+    [dispatch]
+  );
 
-  useEffect(() => {
-    const vState = props.visualizations[state.visualizationName].initialize();
-    props.datasources[state.datasourceName].initialize().then(dsState => {
+  const visualizationStateUpdater = useMemo(
+    () => (newState: unknown) => {
       dispatch({
-        type: 'UPDATE_DATASOURCE',
-        payload: dsState,
+        type: 'UPDATE_VISUALIZATION_STATE',
+        newState,
       });
-    });
+    },
+    [dispatch]
+  );
 
-    dispatch({
-      type: 'UPDATE_VISUALIZATION',
-      payload: vState,
-    });
-  }, []);
+  const datasourcePublicAPI = useMemo(
+    () =>
+      state.activeDatasource && !state.datasourceIsLoading
+        ? props.datasources[state.activeDatasource].getPublicAPI(
+            state.datasourceState,
+            datasourceStateUpdater
+          )
+        : undefined,
+    [
+      props.datasources,
+      state.datasourceIsLoading,
+      state.activeDatasource,
+      datasourceStateUpdater,
+      state.datasourceState,
+    ]
+  );
 
-  return (
-    <div>
-      <h2>Editor Frame</h2>
+  function renderDatasource() {
+    const datasourceProps: DatasourceDataPanelProps = {
+      state: state.datasourceState,
+      setState: datasourceStateUpdater,
+    };
+
+    return (
+      <>
+        <EuiSelect
+          data-test-subj="datasource-switch"
+          options={Object.keys(props.datasources).map(datasourceId => ({
+            value: datasourceId,
+            text: datasourceId,
+          }))}
+          value={state.activeDatasource || undefined}
+          onChange={e => {
+            dispatch({ type: 'SWITCH_DATASOURCE', newDatasourceId: e.target.value });
+          }}
+        />
+        {state.activeDatasource && !state.datasourceIsLoading && (
+          <NativeRenderer
+            render={props.datasources[state.activeDatasource].renderDataPanel}
+            nativeProps={datasourceProps}
+          />
+        )}
+      </>
+    );
+  }
 
-      <div
-        ref={domElement => {
-          if (domElement) {
-            props.datasources[state.datasourceName].renderDataPanel(domElement, {
-              state: state.datasourceState,
-              setState: newState =>
-                dispatch({
-                  type: 'UPDATE_DATASOURCE',
-                  payload: newState,
-                }),
+  function renderVisualization() {
+    return (
+      <>
+        <EuiSelect
+          data-test-subj="visualization-switch"
+          options={Object.keys(props.visualizations).map(visualizationId => ({
+            value: visualizationId,
+            text: visualizationId,
+          }))}
+          value={state.activeDatasource || undefined}
+          onChange={e => {
+            dispatch({
+              type: 'SWITCH_VISUALIZATION',
+              newVisulizationId: e.target.value,
+              // TODO we probably want to have a separate API to "force" a visualization switch
+              // which isn't a result of a picked suggestion
+              initialState: props.visualizations[e.target.value].initialize(),
             });
-          }
-        }}
-      />
+          }}
+        />
+        {state.activeVisualization && state.activeDatasource && !state.datasourceIsLoading && (
+          <NativeRenderer
+            render={props.visualizations[state.activeVisualization].renderConfigPanel}
+            nativeProps={{
+              state: state.visualizationState[state.activeVisualization],
+              setState: visualizationStateUpdater,
+              datasource: datasourcePublicAPI!,
+            }}
+          />
+        )}
+      </>
+    );
+  }
 
-      <div
-        ref={domElement => {
-          if (domElement) {
-            props.visualizations[state.visualizationName].renderConfigPanel(domElement, {
-              datasource: props.datasources[state.datasourceName].getPublicAPI(
-                state.datasourceState,
-                newState =>
-                  dispatch({
-                    type: 'UPDATE_DATASOURCE',
-                    payload: newState,
-                  })
-              ),
-              state: state.visualizationState,
-              setState: newState =>
-                dispatch({
-                  type: 'UPDATE_VISUALIZATION',
-                  payload: newState,
-                }),
-            });
-          }
-        }}
-      />
-    </div>
+  return (
+    <EuiFlexGroup>
+      <EuiFlexItem>{renderDatasource()}</EuiFlexItem>
+      <EuiFlexItem>{renderVisualization()}</EuiFlexItem>
+    </EuiFlexGroup>
   );
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
index 6a0a82877cefb..1b1957bf8d316 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
@@ -6,7 +6,7 @@
 
 import React from 'react';
 import { render, unmountComponentAtNode } from 'react-dom';
-import { Datasource, Visualization, EditorFrameSetup } from '../types';
+import { Datasource, Visualization, EditorFrameSetup, DatasourcePublicAPI } from '../types';
 
 import { EditorFrame } from './editor_frame';
 
@@ -20,34 +20,39 @@ class EditorFramePlugin {
     [key: string]: Visualization;
   } = {};
 
-  private initialDatasource?: string;
-
-  private element: Element | null = null;
-
   public setup(): EditorFrameSetup {
     return {
-      render: domElement => {
-        this.element = domElement;
-        render(
-          <EditorFrame
-            datasources={this.datasources}
-            visualizations={this.visualizations}
-            initialDatasource={this.initialDatasource}
-          />,
-          domElement
-        );
+      createInstance: () => {
+        let domElement: Element;
+        return {
+          mount: (element: Element) => {
+            domElement = element;
+            render(
+              <EditorFrame
+                datasources={this.datasources}
+                visualizations={this.visualizations}
+                initialDatasource={Object.keys(this.datasources)[0]}
+                initialVisualization={Object.keys(this.visualizations)[0]}
+              />,
+              domElement
+            );
+          },
+          unmount: () => {
+            if (domElement) {
+              unmountComponentAtNode(domElement);
+            }
+          },
+        };
       },
-      registerDatasource: (name, datasource) => {
+      registerDatasource: async (name, datasource) => {
         // casting it to an unknown datasource. This doesn't introduce runtime errors
         // because each type T is always also an unknown, but typescript won't do it
         // on it's own because we are loosing type information here.
         // So it's basically explicitly saying "I'm dropping the information about type T here
         // because this information isn't useful to me." but without using any which can leak
-        this.datasources[name] = datasource as Datasource<unknown>;
+        // const state = await datasource.initialize();
 
-        if (!this.initialDatasource) {
-          this.initialDatasource = name;
-        }
+        this.datasources[name] = datasource as Datasource<unknown>;
       },
       registerVisualization: (name, visualization) => {
         this.visualizations[name] = visualization as Visualization<unknown, unknown>;
@@ -56,9 +61,6 @@ class EditorFramePlugin {
   }
 
   public stop() {
-    if (this.element) {
-      unmountComponentAtNode(this.element);
-    }
     return {};
   }
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
new file mode 100644
index 0000000000000..896818bddfc0f
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
@@ -0,0 +1,143 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { getInitialState, reducer } from './state_management';
+import { EditorFrameProps } from './editor_frame';
+import { Datasource, Visualization } from '../types';
+
+describe('editor_frame state management', () => {
+  describe('initialization', () => {
+    let props: EditorFrameProps;
+
+    beforeEach(() => {
+      props = {
+        datasources: { testDatasource: ({} as unknown) as Datasource },
+        visualizations: { testVis: ({ initialize: jest.fn() } as unknown) as Visualization },
+        initialDatasource: 'testDatasource',
+        initialVisualization: 'testVis',
+      };
+    });
+
+    it('should store initial datasource and visualization', () => {
+      expect(getInitialState(props)).toEqual(
+        expect.objectContaining({
+          activeDatasource: 'testDatasource',
+          activeVisualization: 'testVis',
+        })
+      );
+    });
+
+    it('should initialize visualization', () => {
+      const initialVisState = {};
+      props.visualizations.testVis.initialize = jest.fn(() => initialVisState);
+
+      const initialState = getInitialState(props);
+
+      expect(initialState.visualizationState.testVis).toBe(initialVisState);
+      expect(props.visualizations.testVis.initialize).toHaveBeenCalled();
+    });
+
+    it('should not initialize visualization if no initial visualization is passed in', () => {
+      const initialState = getInitialState({ ...props, initialVisualization: null });
+
+      expect(initialState.visualizationState).toEqual({});
+      expect(props.visualizations.testVis.initialize).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('state update', () => {
+    it('should update the corresponding visualization state on update', () => {
+      const newVisState = {};
+      const newState = reducer(
+        {
+          activeDatasource: 'testDatasource',
+          activeVisualization: 'testVis',
+          datasourceIsLoading: false,
+          datasourceState: {},
+          visualizationState: {
+            testVis: {},
+          },
+        },
+        {
+          type: 'UPDATE_VISUALIZATION_STATE',
+          newState: newVisState,
+        }
+      );
+
+      expect(newState.visualizationState).toEqual({
+        testVis: newVisState,
+      });
+    });
+
+    it('should update the datasource state on update', () => {
+      const newDatasourceState = {};
+      const newState = reducer(
+        {
+          activeDatasource: 'testDatasource',
+          activeVisualization: 'testVis',
+          datasourceIsLoading: false,
+          datasourceState: {},
+          visualizationState: {
+            testVis: {},
+          },
+        },
+        {
+          type: 'UPDATE_DATASOURCE_STATE',
+          newState: newDatasourceState,
+        }
+      );
+
+      expect(newState.datasourceState).toBe(newDatasourceState);
+    });
+
+    it('should should switch active visualization but dont loose old state', () => {
+      const testVisState = {};
+      const newVisState = {};
+      const newState = reducer(
+        {
+          activeDatasource: 'testDatasource',
+          activeVisualization: 'testVis',
+          datasourceIsLoading: false,
+          datasourceState: {},
+          visualizationState: {
+            testVis: testVisState,
+          },
+        },
+        {
+          type: 'SWITCH_VISUALIZATION',
+          newVisulizationId: 'testVis2',
+          initialState: newVisState
+        }
+      );
+
+      expect(newState.visualizationState.testVis).toBe(testVisState);
+      expect(newState.visualizationState.testVis2).toBe(newVisState);
+    });
+
+    it('should should switch active datasource and purge visualization state', () => {
+      const newState = reducer(
+        {
+          activeDatasource: 'testDatasource',
+          activeVisualization: 'testVis',
+          datasourceIsLoading: false,
+          datasourceState: {},
+          visualizationState: {
+            testVis: {},
+          },
+        },
+        {
+          type: 'SWITCH_DATASOURCE',
+          newDatasourceId: 'testDatasource2'
+        }
+      );
+
+      expect(newState.visualizationState).toEqual({});
+      expect(newState.activeVisualization).toBe(null);
+      expect(newState.activeDatasource).toBe('testDatasource2');
+      expect(newState.datasourceState).toBe(null);
+    });
+  });
+});
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts
new file mode 100644
index 0000000000000..e35e258d1153c
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts
@@ -0,0 +1,98 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { EditorFrameProps } from './editor_frame';
+
+export interface EditorFrameState {
+  activeDatasource: string | null;
+  datasourceState: unknown;
+  datasourceIsLoading: boolean;
+
+  activeVisualization: string | null;
+  visualizationState: {
+    [visualizationId: string]: unknown;
+  };
+}
+
+export type Action =
+  | {
+      type: 'UPDATE_DATASOURCE_STATE';
+      newState: unknown;
+    }
+  | {
+      type: 'UPDATE_VISUALIZATION_STATE';
+      newState: unknown;
+    }
+  | {
+      type: 'SWITCH_VISUALIZATION';
+      newVisulizationId: string;
+      initialState: unknown;
+    }
+  | {
+      type: 'SWITCH_DATASOURCE';
+      newDatasourceId: string;
+    };
+
+export const getInitialState = (props: EditorFrameProps) => {
+  return {
+    datasourceState: null,
+    datasourceIsLoading: Boolean(props.initialDatasource),
+    activeDatasource: props.initialDatasource,
+    visualizationState: props.initialVisualization
+      ? {
+          [props.initialVisualization]: props.visualizations[
+            props.initialVisualization
+          ].initialize(),
+        }
+      : {},
+    activeVisualization: props.initialVisualization,
+  };
+};
+
+export const reducer = (state: EditorFrameState, action: Action): EditorFrameState => {
+  switch (action.type) {
+    case 'SWITCH_DATASOURCE':
+      return {
+        ...state,
+        activeDatasource: action.newDatasourceId,
+        // purge all visualizations on datasource switch
+        visualizationState: {},
+        activeVisualization: null,
+        datasourceState: null,
+        datasourceIsLoading: true,
+      };
+    case 'SWITCH_VISUALIZATION':
+      return {
+        ...state,
+        activeVisualization: action.newVisulizationId,
+        visualizationState: {
+          ...state.visualizationState,
+          [action.newVisulizationId]: action.initialState,
+        },
+      };
+    case 'UPDATE_DATASOURCE_STATE':
+      return {
+        ...state,
+        // when the datasource state is updated, the initialization is complete
+        datasourceIsLoading: false,
+        datasourceState: action.newState,
+      };
+    case 'UPDATE_VISUALIZATION_STATE':
+      if (state.activeVisualization) {
+        return {
+          ...state,
+          visualizationState: {
+            ...state.visualizationState,
+            [state.activeVisualization]: action.newState,
+          },
+        };
+      } else {
+        throw new Error('Invariant: visualization state got updated without active visualization');
+      }
+    default:
+      return state;
+  }
+};
diff --git a/x-pack/plugins/lens/public/index.ts b/x-pack/plugins/lens/public/index.ts
index 874249ff75d85..f8f074cdb99bf 100644
--- a/x-pack/plugins/lens/public/index.ts
+++ b/x-pack/plugins/lens/public/index.ts
@@ -6,8 +6,8 @@
 
 export * from './types';
 
-import { IScope } from 'angular';
 import { render, unmountComponentAtNode } from 'react-dom';
+import { IScope } from 'angular';
 import chrome from 'ui/chrome';
 import { appSetup, appStop } from './app_plugin';
 
diff --git a/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx b/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
index 29b018320a6db..567bc6cfd10a8 100644
--- a/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
+++ b/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
@@ -36,6 +36,13 @@ export interface Field {
   searchable: boolean;
 }
 
+export interface Field {
+  name: string;
+  type: string;
+  aggregatable: boolean;
+  searchable: boolean;
+}
+
 export interface IndexPatternPersistedState {
   currentIndexPattern: string;
 
@@ -49,12 +56,13 @@ export type IndexPatternPrivateState = IndexPatternPersistedState & {
   indexPatterns: { [id: string]: IndexPattern };
 };
 
+type PersistedKeys = 'currentIndexPattern' | 'columns' | 'columnOrder';
+
+export type PersistableIndexPatternPrivateState = Pick<IndexPatternPrivateState, PersistedKeys>;
+
 // Not stateful. State is persisted to the frame
-export const indexPatternDatasource: Datasource<
-  IndexPatternPrivateState,
-  IndexPatternPersistedState
-> = {
-  async initialize(state?: IndexPatternPersistedState) {
+export const indexPatternDatasource: Datasource<IndexPatternPrivateState, PersistableIndexPatternPrivateState> = {
+  async initialize(state?: PersistableIndexPatternPrivateState) {
     // TODO: Make fetch request to load indexPatterns from saved objects
     if (state) {
       return {
diff --git a/x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts b/x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts
new file mode 100644
index 0000000000000..6039f0888f596
--- /dev/null
+++ b/x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts
@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+
+// This mock overwrites react to use the layout effect hook instead of the normal one.
+// This is done to have effects executed synchronously to be able to test them without
+// setTimeout hacks.
+
+// eslint-disable-next-line import/no-default-export
+export default { ...React, useEffect: React.useLayoutEffect };
diff --git a/x-pack/plugins/lens/public/native_renderer/index.ts b/x-pack/plugins/lens/public/native_renderer/index.ts
new file mode 100644
index 0000000000000..0ef9bd8807bc5
--- /dev/null
+++ b/x-pack/plugins/lens/public/native_renderer/index.ts
@@ -0,0 +1,7 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+export * from './native_renderer';
diff --git a/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx b/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
new file mode 100644
index 0000000000000..089b1ed02d166
--- /dev/null
+++ b/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
@@ -0,0 +1,134 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+import { render } from 'react-dom';
+import { NativeRenderer } from './native_renderer';
+
+describe('native_renderer', () => {
+  let mountpoint: Element;
+
+  beforeEach(() => {
+    mountpoint = document.createElement('div');
+  });
+
+  afterEach(() => {
+    mountpoint.remove();
+  });
+
+  it('should render element in container', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).toHaveBeenCalledWith(containerElement, testProps);
+  });
+
+  it('should not render again if props do not change', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('should not render again if props do not change shallowly', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ ...testProps }} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('should not render again for unchanged callback functions', () => {
+    const renderSpy = jest.fn();
+    const testCallback = () => {};
+    const testState = { a: 'abc' };
+
+    render(
+      <NativeRenderer
+        render={renderSpy}
+        nativeProps={{ state: testState, setState: testCallback }}
+      />,
+      mountpoint
+    );
+    render(
+      <NativeRenderer
+        render={renderSpy}
+        nativeProps={{ state: testState, setState: testCallback }}
+      />,
+      mountpoint
+    );
+    expect(renderSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it('should render again once if props change', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'def' }} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'def' }} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(2);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).lastCalledWith(containerElement, { a: 'def' });
+  });
+
+  it('should render again once if props is just a string', () => {
+    const renderSpy = jest.fn();
+    const testProps = 'abc';
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps="def" />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps="def" />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(2);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).lastCalledWith(containerElement, 'def');
+  });
+
+  it('should render again if props are extended', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'abc', b: 'def' }} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(2);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).lastCalledWith(containerElement, { a: 'abc', b: 'def' });
+  });
+
+  it('should render again if props are limited', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc', b: 'def' };
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'abc' }} />, mountpoint);
+    expect(renderSpy).toHaveBeenCalledTimes(2);
+    const containerElement = mountpoint.firstElementChild;
+    expect(renderSpy).lastCalledWith(containerElement, { a: 'abc' });
+  });
+
+  it('should render a div as container', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    const containerElement: Element = mountpoint.firstElementChild!;
+    expect(containerElement.nodeName).toBe('DIV');
+  });
+
+  it('should render a specified element as container', () => {
+    const renderSpy = jest.fn();
+    const testProps = { a: 'abc' };
+
+    render(<NativeRenderer render={renderSpy} tag="span" nativeProps={testProps} />, mountpoint);
+    const containerElement: Element = mountpoint.firstElementChild!;
+    expect(containerElement.nodeName).toBe('SPAN');
+  });
+});
diff --git a/x-pack/plugins/lens/public/native_renderer/native_renderer.tsx b/x-pack/plugins/lens/public/native_renderer/native_renderer.tsx
new file mode 100644
index 0000000000000..f0eb4b829c153
--- /dev/null
+++ b/x-pack/plugins/lens/public/native_renderer/native_renderer.tsx
@@ -0,0 +1,80 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React, { useEffect, useRef } from 'react';
+
+export interface NativeRendererProps<T> {
+  render: (domElement: Element, props: T) => void;
+  nativeProps: T;
+  tag?: string;
+  children?: never;
+}
+
+function is(x: unknown, y: unknown) {
+  return (x === y && (x !== 0 || 1 / (x as number) === 1 / (y as number))) || (x !== x && y !== y);
+}
+
+function isShallowDifferent<T>(objA: T, objB: T): boolean {
+  if (is(objA, objB)) {
+    return false;
+  }
+
+  if (typeof objA !== 'object' || objA === null || typeof objB !== 'object' || objB === null) {
+    return true;
+  }
+
+  const keysA = Object.keys(objA) as Array<keyof T>;
+  const keysB = Object.keys(objB) as Array<keyof T>;
+
+  if (keysA.length !== keysB.length) {
+    return true;
+  }
+
+  for (let i = 0; i < keysA.length; i++) {
+    if (!window.hasOwnProperty.call(objB, keysA[i]) || !is(objA[keysA[i]], objB[keysA[i]])) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+/**
+ * A component which takes care of providing a mountpoint for a generic render
+ * function which takes an html element and an optional props object.
+ * It also takes care of calling render again if the props object changes.
+ * By default the mountpoint element will be a div, this can be changed with the
+ * `tag` prop.
+ *
+ * @param props
+ */
+export function NativeRenderer<T>({ render, nativeProps, tag }: NativeRendererProps<T>) {
+  const elementRef = useRef<Element | null>(null);
+  const propsRef = useRef<T | null>(null);
+
+  function renderAndUpdate(element: Element) {
+    elementRef.current = element;
+    propsRef.current = nativeProps;
+    render(element, nativeProps);
+  }
+
+  useEffect(
+    () => {
+      if (elementRef.current && isShallowDifferent(propsRef.current, nativeProps)) {
+        renderAndUpdate(elementRef.current);
+      }
+    },
+    [nativeProps]
+  );
+
+  return React.createElement(tag || 'div', {
+    ref: element => {
+      if (element && element !== elementRef.current) {
+        renderAndUpdate(element);
+      }
+    },
+  });
+}
diff --git a/x-pack/plugins/lens/public/types.ts b/x-pack/plugins/lens/public/types.ts
index ab70ff10e1cb5..67f4d2cdb09dd 100644
--- a/x-pack/plugins/lens/public/types.ts
+++ b/x-pack/plugins/lens/public/types.ts
@@ -4,8 +4,15 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
+export interface EditorFrameOptions {
+}
+
+export interface EditorFrameInstance {
+  mount: (element: Element) => void;
+  unmount: () => void;
+}
 export interface EditorFrameSetup {
-  render: (domElement: Element) => void;
+  createInstance: (options: EditorFrameOptions) => EditorFrameInstance;
   // generic type on the API functions to pull the "unknown vs. specific type" error into the implementation
   registerDatasource: <T, P>(name: string, datasource: Datasource<T, P>) => void;
   registerVisualization: <T, P>(name: string, visualization: Visualization<T, P>) => void;
@@ -139,12 +146,9 @@ export interface VisualizationSuggestion<T = unknown> {
 }
 
 export interface Visualization<T = unknown, P = unknown> {
-  // For initializing, either from an empty state or from persisted state
-  // Because this will be called at runtime, state might have a type of `any` and
-  // visualizations should validate their arguments
+  // For initializing from saved object
   initialize: (state?: P) => T;
 
-  // Given the current state, which parts should be saved?
   getPersistableState: (state: T) => P;
 
   renderConfigPanel: (domElement: Element, props: VisualizationProps<T>) => void;

From 8147aaa957ce0b9505ec3c8cd6c6cb0f9d3b71ec Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Fri, 10 May 2019 15:49:02 +0200
Subject: [PATCH 11/41] fix linting errors

---
 x-pack/plugins/lens/public/app_plugin/plugin.tsx |  1 -
 .../editor_frame_plugin/editor_frame.test.tsx    | 16 +++++++++-------
 .../editor_frame_plugin/state_management.test.ts |  4 ++--
 .../public/indexpattern_plugin/indexpattern.tsx  |  5 ++++-
 x-pack/plugins/lens/public/types.ts              |  4 ++--
 5 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/x-pack/plugins/lens/public/app_plugin/plugin.tsx b/x-pack/plugins/lens/public/app_plugin/plugin.tsx
index aeea52c23c0e7..857cee9adbc64 100644
--- a/x-pack/plugins/lens/public/app_plugin/plugin.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/plugin.tsx
@@ -12,7 +12,6 @@ import { App } from './app';
 import { EditorFrameInstance } from '../types';
 
 export class AppPlugin {
-
   private instance: EditorFrameInstance | null = null;
 
   constructor() {}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx
index 5d9353113a1c9..b1c46afd5e1dc 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx
@@ -194,7 +194,7 @@ describe('editor_frame', () => {
       // don't re-render datasource when visulization changes
       expect(mockDatasource.renderDataPanel).toHaveBeenCalledTimes(1);
     });
-    
+
     it('should re-render data panel after state update', async () => {
       mount(
         <EditorFrame
@@ -212,8 +212,8 @@ describe('editor_frame', () => {
       await nextTick();
 
       const updatedState = {};
-      const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock
-        .calls[0][1].setState;
+      const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock.calls[0][1]
+        .setState;
       setDatasourceState(updatedState);
 
       expect(mockDatasource.renderDataPanel).toHaveBeenCalledTimes(2);
@@ -224,7 +224,7 @@ describe('editor_frame', () => {
         })
       );
     });
-    
+
     it('should re-render config panel with updated datasource api after datasource state update', async () => {
       mount(
         <EditorFrame
@@ -242,10 +242,12 @@ describe('editor_frame', () => {
       await nextTick();
 
       const updatedPublicAPI = {};
-      mockDatasource.getPublicAPI = jest.fn(() => updatedPublicAPI as unknown as DatasourcePublicAPI);
+      mockDatasource.getPublicAPI = jest.fn(
+        () => (updatedPublicAPI as unknown) as DatasourcePublicAPI
+      );
 
-      const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock
-        .calls[0][1].setState;
+      const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock.calls[0][1]
+        .setState;
       setDatasourceState({});
 
       expect(mockVisualization.renderConfigPanel).toHaveBeenCalledTimes(2);
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
index 896818bddfc0f..fb96817685c45 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
@@ -109,7 +109,7 @@ describe('editor_frame state management', () => {
         {
           type: 'SWITCH_VISUALIZATION',
           newVisulizationId: 'testVis2',
-          initialState: newVisState
+          initialState: newVisState,
         }
       );
 
@@ -130,7 +130,7 @@ describe('editor_frame state management', () => {
         },
         {
           type: 'SWITCH_DATASOURCE',
-          newDatasourceId: 'testDatasource2'
+          newDatasourceId: 'testDatasource2',
         }
       );
 
diff --git a/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx b/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
index 567bc6cfd10a8..53522ff39d6fb 100644
--- a/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
+++ b/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
@@ -61,7 +61,10 @@ type PersistedKeys = 'currentIndexPattern' | 'columns' | 'columnOrder';
 export type PersistableIndexPatternPrivateState = Pick<IndexPatternPrivateState, PersistedKeys>;
 
 // Not stateful. State is persisted to the frame
-export const indexPatternDatasource: Datasource<IndexPatternPrivateState, PersistableIndexPatternPrivateState> = {
+export const indexPatternDatasource: Datasource<
+  IndexPatternPrivateState,
+  PersistableIndexPatternPrivateState
+> = {
   async initialize(state?: PersistableIndexPatternPrivateState) {
     // TODO: Make fetch request to load indexPatterns from saved objects
     if (state) {
diff --git a/x-pack/plugins/lens/public/types.ts b/x-pack/plugins/lens/public/types.ts
index 67f4d2cdb09dd..a0328cb3bd988 100644
--- a/x-pack/plugins/lens/public/types.ts
+++ b/x-pack/plugins/lens/public/types.ts
@@ -4,8 +4,8 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-export interface EditorFrameOptions {
-}
+// eslint-disable-next-line
+export interface EditorFrameOptions {}
 
 export interface EditorFrameInstance {
   mount: (element: Element) => void;

From 4baf6b196405d6f563e4078fa8b0fbe0df82d83d Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Fri, 10 May 2019 16:49:44 +0200
Subject: [PATCH 12/41] fix native renderer tests

---
 .../public/native_renderer/__mocks__/react.ts | 14 ---
 .../native_renderer/native_renderer.test.tsx  | 87 +++++++++++++++----
 2 files changed, 70 insertions(+), 31 deletions(-)
 delete mode 100644 x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts

diff --git a/x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts b/x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts
deleted file mode 100644
index 6039f0888f596..0000000000000
--- a/x-pack/plugins/lens/public/native_renderer/__mocks__/react.ts
+++ /dev/null
@@ -1,14 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License;
- * you may not use this file except in compliance with the Elastic License.
- */
-
-import React from 'react';
-
-// This mock overwrites react to use the layout effect hook instead of the normal one.
-// This is done to have effects executed synchronously to be able to test them without
-// setTimeout hacks.
-
-// eslint-disable-next-line import/no-default-export
-export default { ...React, useEffect: React.useLayoutEffect };
diff --git a/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx b/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
index 089b1ed02d166..af5196165f9a5 100644
--- a/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
+++ b/x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx
@@ -7,6 +7,14 @@
 import React from 'react';
 import { render } from 'react-dom';
 import { NativeRenderer } from './native_renderer';
+import { act } from 'react-dom/test-utils';
+
+function renderAndTriggerHooks(element: JSX.Element, mountpoint: Element) {
+  // act takes care of triggering state hooks
+  act(() => {
+    render(element, mountpoint);
+  });
+}
 
 describe('native_renderer', () => {
   let mountpoint: Element;
@@ -23,7 +31,10 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).toHaveBeenCalledWith(containerElement, testProps);
   });
@@ -32,8 +43,14 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
     expect(renderSpy).toHaveBeenCalledTimes(1);
   });
 
@@ -41,8 +58,14 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} nativeProps={{ ...testProps }} />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={{ ...testProps }} />,
+      mountpoint
+    );
     expect(renderSpy).toHaveBeenCalledTimes(1);
   });
 
@@ -72,9 +95,18 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'def' }} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'def' }} />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={{ a: 'def' }} />,
+      mountpoint
+    );
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={{ a: 'def' }} />,
+      mountpoint
+    );
     expect(renderSpy).toHaveBeenCalledTimes(2);
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).lastCalledWith(containerElement, { a: 'def' });
@@ -84,9 +116,12 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = 'abc';
 
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} nativeProps="def" />, mountpoint);
-    render(<NativeRenderer render={renderSpy} nativeProps="def" />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
+    renderAndTriggerHooks(<NativeRenderer render={renderSpy} nativeProps="def" />, mountpoint);
+    renderAndTriggerHooks(<NativeRenderer render={renderSpy} nativeProps="def" />, mountpoint);
     expect(renderSpy).toHaveBeenCalledTimes(2);
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).lastCalledWith(containerElement, 'def');
@@ -96,8 +131,14 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'abc', b: 'def' }} />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={{ a: 'abc', b: 'def' }} />,
+      mountpoint
+    );
     expect(renderSpy).toHaveBeenCalledTimes(2);
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).lastCalledWith(containerElement, { a: 'abc', b: 'def' });
@@ -107,8 +148,14 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc', b: 'def' };
 
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
-    render(<NativeRenderer render={renderSpy} nativeProps={{ a: 'abc' }} />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={{ a: 'abc' }} />,
+      mountpoint
+    );
     expect(renderSpy).toHaveBeenCalledTimes(2);
     const containerElement = mountpoint.firstElementChild;
     expect(renderSpy).lastCalledWith(containerElement, { a: 'abc' });
@@ -118,7 +165,10 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} nativeProps={testProps} />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} nativeProps={testProps} />,
+      mountpoint
+    );
     const containerElement: Element = mountpoint.firstElementChild!;
     expect(containerElement.nodeName).toBe('DIV');
   });
@@ -127,7 +177,10 @@ describe('native_renderer', () => {
     const renderSpy = jest.fn();
     const testProps = { a: 'abc' };
 
-    render(<NativeRenderer render={renderSpy} tag="span" nativeProps={testProps} />, mountpoint);
+    renderAndTriggerHooks(
+      <NativeRenderer render={renderSpy} tag="span" nativeProps={testProps} />,
+      mountpoint
+    );
     const containerElement: Element = mountpoint.firstElementChild!;
     expect(containerElement.nodeName).toBe('SPAN');
   });

From 033a632ba6b7c4f92e22370fe4460685d3dd647c Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Fri, 10 May 2019 17:31:15 +0200
Subject: [PATCH 13/41] re structure editor frame

---
 .../editor_frame_plugin/editor_frame.tsx      | 152 ------------------
 .../editor_frame/config_panel_wrapper.tsx     |  63 ++++++++
 .../editor_frame/data_panel_wrapper.tsx       |  58 +++++++
 .../{ => editor_frame}/editor_frame.test.tsx  |   2 +-
 .../editor_frame/editor_frame.tsx             |  88 ++++++++++
 .../editor_frame_plugin/editor_frame/index.ts |   7 +
 .../public/editor_frame_plugin/plugin.tsx     |   2 +-
 7 files changed, 218 insertions(+), 154 deletions(-)
 delete mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
 rename x-pack/plugins/lens/public/editor_frame_plugin/{ => editor_frame}/editor_frame.test.tsx (99%)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/index.ts

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
deleted file mode 100644
index 13e184ccb4d41..0000000000000
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.tsx
+++ /dev/null
@@ -1,152 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License;
- * you may not use this file except in compliance with the Elastic License.
- */
-
-import React, { useEffect, useReducer, useMemo } from 'react';
-import { EuiSelect, EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
-import { Datasource, Visualization, DatasourceDataPanelProps } from '../types';
-import { NativeRenderer } from '../native_renderer';
-import { reducer, getInitialState } from './state_management';
-
-export interface EditorFrameProps {
-  datasources: { [key: string]: Datasource };
-  visualizations: { [key: string]: Visualization };
-
-  initialDatasource: string | null;
-  initialVisualization: string | null;
-}
-
-export function EditorFrame(props: EditorFrameProps) {
-  const [state, dispatch] = useReducer(reducer, props, getInitialState);
-
-  // Initialize current datasource
-  useEffect(
-    () => {
-      let datasourceGotSwitched = false;
-      if (state.datasourceIsLoading && state.activeDatasource) {
-        props.datasources[state.activeDatasource].initialize().then(datasourceState => {
-          if (!datasourceGotSwitched) {
-            dispatch({
-              type: 'UPDATE_DATASOURCE_STATE',
-              newState: datasourceState,
-            });
-          }
-        });
-
-        return () => {
-          datasourceGotSwitched = true;
-        };
-      }
-    },
-    [state.activeDatasource, state.datasourceIsLoading]
-  );
-
-  const datasourceStateUpdater = useMemo(
-    () => (newState: unknown) => {
-      dispatch({
-        type: 'UPDATE_DATASOURCE_STATE',
-        newState,
-      });
-    },
-    [dispatch]
-  );
-
-  const visualizationStateUpdater = useMemo(
-    () => (newState: unknown) => {
-      dispatch({
-        type: 'UPDATE_VISUALIZATION_STATE',
-        newState,
-      });
-    },
-    [dispatch]
-  );
-
-  const datasourcePublicAPI = useMemo(
-    () =>
-      state.activeDatasource && !state.datasourceIsLoading
-        ? props.datasources[state.activeDatasource].getPublicAPI(
-            state.datasourceState,
-            datasourceStateUpdater
-          )
-        : undefined,
-    [
-      props.datasources,
-      state.datasourceIsLoading,
-      state.activeDatasource,
-      datasourceStateUpdater,
-      state.datasourceState,
-    ]
-  );
-
-  function renderDatasource() {
-    const datasourceProps: DatasourceDataPanelProps = {
-      state: state.datasourceState,
-      setState: datasourceStateUpdater,
-    };
-
-    return (
-      <>
-        <EuiSelect
-          data-test-subj="datasource-switch"
-          options={Object.keys(props.datasources).map(datasourceId => ({
-            value: datasourceId,
-            text: datasourceId,
-          }))}
-          value={state.activeDatasource || undefined}
-          onChange={e => {
-            dispatch({ type: 'SWITCH_DATASOURCE', newDatasourceId: e.target.value });
-          }}
-        />
-        {state.activeDatasource && !state.datasourceIsLoading && (
-          <NativeRenderer
-            render={props.datasources[state.activeDatasource].renderDataPanel}
-            nativeProps={datasourceProps}
-          />
-        )}
-      </>
-    );
-  }
-
-  function renderVisualization() {
-    return (
-      <>
-        <EuiSelect
-          data-test-subj="visualization-switch"
-          options={Object.keys(props.visualizations).map(visualizationId => ({
-            value: visualizationId,
-            text: visualizationId,
-          }))}
-          value={state.activeDatasource || undefined}
-          onChange={e => {
-            dispatch({
-              type: 'SWITCH_VISUALIZATION',
-              newVisulizationId: e.target.value,
-              // TODO we probably want to have a separate API to "force" a visualization switch
-              // which isn't a result of a picked suggestion
-              initialState: props.visualizations[e.target.value].initialize(),
-            });
-          }}
-        />
-        {state.activeVisualization && state.activeDatasource && !state.datasourceIsLoading && (
-          <NativeRenderer
-            render={props.visualizations[state.activeVisualization].renderConfigPanel}
-            nativeProps={{
-              state: state.visualizationState[state.activeVisualization],
-              setState: visualizationStateUpdater,
-              datasource: datasourcePublicAPI!,
-            }}
-          />
-        )}
-      </>
-    );
-  }
-
-  return (
-    <EuiFlexGroup>
-      <EuiFlexItem>{renderDatasource()}</EuiFlexItem>
-      <EuiFlexItem>{renderVisualization()}</EuiFlexItem>
-    </EuiFlexGroup>
-  );
-}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
new file mode 100644
index 0000000000000..e494b8f81d786
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
@@ -0,0 +1,63 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React, { useMemo } from 'react';
+import { EuiSelect } from '@elastic/eui';
+import { NativeRenderer } from '../../native_renderer';
+import { Action } from '../state_management';
+import { Visualization, DatasourcePublicAPI } from '../../types';
+
+interface ConfigPanelWrapperProps {
+  visualizationState: Record<string, unknown>;
+  visualizations: Record<string, Visualization>;
+  activeVisualization: string | null;
+  dispatch: (action: Action) => void;
+  datasourcePublicAPI: DatasourcePublicAPI;
+}
+
+export function ConfigPanelWrapper(props: ConfigPanelWrapperProps) {
+  const setVisualizationState = useMemo(
+    () => (newState: unknown) => {
+      props.dispatch({
+        type: 'UPDATE_VISUALIZATION_STATE',
+        newState,
+      });
+    },
+    [props.dispatch]
+  );
+
+  return (
+    <>
+      <EuiSelect
+        data-test-subj="visualization-switch"
+        options={Object.keys(props.visualizations).map(visualizationId => ({
+          value: visualizationId,
+          text: visualizationId,
+        }))}
+        value={props.activeVisualization || undefined}
+        onChange={e => {
+          props.dispatch({
+            type: 'SWITCH_VISUALIZATION',
+            newVisulizationId: e.target.value,
+            // TODO we probably want to have a separate API to "force" a visualization switch
+            // which isn't a result of a picked suggestion
+            initialState: props.visualizations[e.target.value].initialize(),
+          });
+        }}
+      />
+      {props.activeVisualization && (
+        <NativeRenderer
+          render={props.visualizations[props.activeVisualization].renderConfigPanel}
+          nativeProps={{
+            state: props.visualizationState[props.activeVisualization],
+            setState: setVisualizationState,
+            datasource: props.datasourcePublicAPI,
+          }}
+        />
+      )}
+    </>
+  );
+}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
new file mode 100644
index 0000000000000..2fdf509c6716e
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
@@ -0,0 +1,58 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React, { useMemo } from 'react';
+import { EuiSelect } from '@elastic/eui';
+import { DatasourceDataPanelProps, Datasource } from '../..';
+import { NativeRenderer } from '../../native_renderer';
+import { Action } from '../state_management';
+
+interface DataPanelWrapperProps {
+  datasourceState: unknown;
+  datasources: Record<string, Datasource>;
+  activeDatasource: string | null;
+  datasourceIsLoading: boolean;
+  dispatch: (action: Action) => void;
+}
+
+export function DataPanelWrapper(props: DataPanelWrapperProps) {
+  const setDatasourceState = useMemo(
+    () => (newState: unknown) => {
+      props.dispatch({
+        type: 'UPDATE_DATASOURCE_STATE',
+        newState,
+      });
+    },
+    [props.dispatch]
+  );
+
+  const datasourceProps: DatasourceDataPanelProps = {
+    state: props.datasourceState,
+    setState: setDatasourceState,
+  };
+
+  return (
+    <>
+      <EuiSelect
+        data-test-subj="datasource-switch"
+        options={Object.keys(props.datasources).map(datasourceId => ({
+          value: datasourceId,
+          text: datasourceId,
+        }))}
+        value={props.activeDatasource || undefined}
+        onChange={e => {
+          props.dispatch({ type: 'SWITCH_DATASOURCE', newDatasourceId: e.target.value });
+        }}
+      />
+      {props.activeDatasource && !props.datasourceIsLoading && (
+        <NativeRenderer
+          render={props.datasources[props.activeDatasource].renderDataPanel}
+          nativeProps={datasourceProps}
+        />
+      )}
+    </>
+  );
+}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
similarity index 99%
rename from x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx
rename to x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
index b1c46afd5e1dc..1b2496b570d82 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
@@ -7,7 +7,7 @@
 import React from 'react';
 import { mount, ReactWrapper } from 'enzyme';
 import { EditorFrame } from './editor_frame';
-import { Visualization, Datasource, VisualizationProps, DatasourcePublicAPI } from '../types';
+import { Visualization, Datasource, DatasourcePublicAPI } from '../../types';
 
 const nextTick = () => new Promise(resolve => setTimeout(resolve));
 
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
new file mode 100644
index 0000000000000..3f92021f8c2ad
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -0,0 +1,88 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React, { useEffect, useReducer, useMemo } from 'react';
+import { EuiSelect, EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
+import { Datasource, Visualization, DatasourceDataPanelProps } from '../../types';
+import { NativeRenderer } from '../../native_renderer';
+import { reducer, getInitialState } from '../state_management';
+import { DataPanelWrapper } from './data_panel_wrapper';
+import { ConfigPanelWrapper } from './config_panel_wrapper';
+
+export interface EditorFrameProps {
+  datasources: { [key: string]: Datasource };
+  visualizations: { [key: string]: Visualization };
+
+  initialDatasource: string | null;
+  initialVisualization: string | null;
+}
+
+export function EditorFrame(props: EditorFrameProps) {
+  const [state, dispatch] = useReducer(reducer, props, getInitialState);
+
+  // Initialize current datasource
+  useEffect(
+    () => {
+      let datasourceGotSwitched = false;
+      if (state.datasourceIsLoading && state.activeDatasource) {
+        props.datasources[state.activeDatasource].initialize().then(datasourceState => {
+          if (!datasourceGotSwitched) {
+            dispatch({
+              type: 'UPDATE_DATASOURCE_STATE',
+              newState: datasourceState,
+            });
+          }
+        });
+
+        return () => {
+          datasourceGotSwitched = true;
+        };
+      }
+    },
+    [state.activeDatasource, state.datasourceIsLoading]
+  );
+
+  const datasourcePublicAPI = useMemo(
+    () =>
+      state.activeDatasource && !state.datasourceIsLoading
+        ? props.datasources[state.activeDatasource].getPublicAPI(
+            state.datasourceState,
+            (newState: unknown) => {
+              dispatch({
+                type: 'UPDATE_DATASOURCE_STATE',
+                newState,
+              });
+            }
+          )
+        : undefined,
+    [props.datasources, state.datasourceIsLoading, state.activeDatasource, state.datasourceState]
+  );
+
+  return (
+    <EuiFlexGroup>
+      <EuiFlexItem>
+        <DataPanelWrapper
+          activeDatasource={state.activeDatasource}
+          datasourceIsLoading={state.datasourceIsLoading}
+          datasourceState={state.datasourceState}
+          datasources={props.datasources}
+          dispatch={dispatch}
+        />
+      </EuiFlexItem>
+      <EuiFlexItem>
+        {state.activeDatasource && !state.datasourceIsLoading && (
+          <ConfigPanelWrapper
+            visualizations={props.visualizations}
+            activeVisualization={state.activeVisualization}
+            datasourcePublicAPI={datasourcePublicAPI!}
+            dispatch={dispatch}
+            visualizationState={state.visualizationState}
+          />
+        )}
+      </EuiFlexItem>
+    </EuiFlexGroup>
+  );
+}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/index.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/index.ts
new file mode 100644
index 0000000000000..41558caafc64c
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/index.ts
@@ -0,0 +1,7 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+export * from './editor_frame';
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
index 1b1957bf8d316..34923e98ad770 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
@@ -6,7 +6,7 @@
 
 import React from 'react';
 import { render, unmountComponentAtNode } from 'react-dom';
-import { Datasource, Visualization, EditorFrameSetup, DatasourcePublicAPI } from '../types';
+import { Datasource, Visualization, EditorFrameSetup } from '../types';
 
 import { EditorFrame } from './editor_frame';
 

From 345e20f534524a70a78d4002184d4218bfae8537 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Fri, 10 May 2019 18:15:49 +0200
Subject: [PATCH 14/41] fix tests

---
 .../editor_frame/editor_frame.test.tsx        | 157 ++++++++++--------
 1 file changed, 90 insertions(+), 67 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
index 1b2496b570d82..4b334a541c984 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
@@ -8,6 +8,7 @@ import React from 'react';
 import { mount, ReactWrapper } from 'enzyme';
 import { EditorFrame } from './editor_frame';
 import { Visualization, Datasource, DatasourcePublicAPI } from '../../types';
+import { act } from 'react-dom/test-utils';
 
 const nextTick = () => new Promise(resolve => setTimeout(resolve));
 
@@ -47,54 +48,60 @@ describe('editor_frame', () => {
 
   describe('initialization', () => {
     it('should initialize initial datasource and visualization if present', () => {
-      mount(
-        <EditorFrame
-          visualizations={{
-            testVis: mockVisualization,
-          }}
-          datasources={{
-            testDatasource: mockDatasource,
-          }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
-        />
-      );
+      act(() => {
+        mount(
+          <EditorFrame
+            visualizations={{
+              testVis: mockVisualization,
+            }}
+            datasources={{
+              testDatasource: mockDatasource,
+            }}
+            initialDatasource="testDatasource"
+            initialVisualization="testVis"
+          />
+        );
+      });
 
       expect(mockVisualization.initialize).toHaveBeenCalled();
       expect(mockDatasource.initialize).toHaveBeenCalled();
     });
 
     it('should not initialize datasource and visualization if no initial one is specificed', () => {
-      mount(
-        <EditorFrame
-          visualizations={{
-            testVis: mockVisualization,
-          }}
-          datasources={{
-            testDatasource: mockDatasource,
-          }}
-          initialDatasource={null}
-          initialVisualization={null}
-        />
-      );
+      act(() => {
+        mount(
+          <EditorFrame
+            visualizations={{
+              testVis: mockVisualization,
+            }}
+            datasources={{
+              testDatasource: mockDatasource,
+            }}
+            initialDatasource={null}
+            initialVisualization={null}
+          />
+        );
+      });
 
       expect(mockVisualization.initialize).not.toHaveBeenCalled();
       expect(mockDatasource.initialize).not.toHaveBeenCalled();
     });
 
     it('should not render something before datasource is initialized', () => {
-      mount(
-        <EditorFrame
-          visualizations={{
-            testVis: mockVisualization,
-          }}
-          datasources={{
-            testDatasource: mockDatasource,
-          }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
-        />
-      );
+      act(() => {
+        mount(
+          <EditorFrame
+            visualizations={{
+              testVis: mockVisualization,
+            }}
+            datasources={{
+              testDatasource: mockDatasource,
+            }}
+            initialDatasource="testDatasource"
+            initialVisualization="testVis"
+          />
+        );
+      });
 
       expect(mockVisualization.renderConfigPanel).not.toHaveBeenCalled();
       expect(mockDatasource.renderDataPanel).not.toHaveBeenCalled();
@@ -104,24 +111,26 @@ describe('editor_frame', () => {
       const initialState = {};
       let databaseInitialized: ({}) => void;
 
-      mount(
-        <EditorFrame
-          visualizations={{
-            testVis: mockVisualization,
-          }}
-          datasources={{
-            testDatasource: {
-              ...mockDatasource,
-              initialize: () =>
-                new Promise(resolve => {
-                  databaseInitialized = resolve;
-                }),
-            },
-          }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
-        />
-      );
+      act(() => {
+        mount(
+          <EditorFrame
+            visualizations={{
+              testVis: mockVisualization,
+            }}
+            datasources={{
+              testDatasource: {
+                ...mockDatasource,
+                initialize: () =>
+                  new Promise(resolve => {
+                    databaseInitialized = resolve;
+                  }),
+              },
+            }}
+            initialDatasource="testDatasource"
+            initialVisualization="testVis"
+          />
+        );
+      });
 
       await nextTick();
       databaseInitialized!(initialState);
@@ -181,7 +190,9 @@ describe('editor_frame', () => {
       const updatedState = {};
       const setVisualizationState = (mockVisualization.renderConfigPanel as jest.Mock).mock
         .calls[0][1].setState;
-      setVisualizationState(updatedState);
+      act(() => {
+        setVisualizationState(updatedState);
+      });
 
       expect(mockVisualization.renderConfigPanel).toHaveBeenCalledTimes(2);
       expect(mockVisualization.renderConfigPanel).toHaveBeenLastCalledWith(
@@ -214,7 +225,9 @@ describe('editor_frame', () => {
       const updatedState = {};
       const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock.calls[0][1]
         .setState;
-      setDatasourceState(updatedState);
+      act(() => {
+        setDatasourceState(updatedState);
+      });
 
       expect(mockDatasource.renderDataPanel).toHaveBeenCalledTimes(2);
       expect(mockDatasource.renderDataPanel).toHaveBeenLastCalledWith(
@@ -248,7 +261,9 @@ describe('editor_frame', () => {
 
       const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock.calls[0][1]
         .setState;
-      setDatasourceState({});
+      act(() => {
+        setDatasourceState({});
+      });
 
       expect(mockVisualization.renderConfigPanel).toHaveBeenCalledTimes(2);
       expect(mockVisualization.renderConfigPanel).toHaveBeenLastCalledWith(
@@ -330,7 +345,9 @@ describe('editor_frame', () => {
 
       const updatedState = {};
       const setDatasourceState = (mockDatasource.getPublicAPI as jest.Mock).mock.calls[0][1];
-      setDatasourceState(updatedState);
+      act(() => {
+        setDatasourceState(updatedState);
+      });
 
       expect(mockDatasource.getPublicAPI).toHaveBeenCalledTimes(2);
       expect(mockDatasource.getPublicAPI).toHaveBeenLastCalledWith(
@@ -372,9 +389,11 @@ describe('editor_frame', () => {
     });
 
     it('should initialize other datasource on switch', async () => {
-      instance
-        .find('select[data-test-subj="datasource-switch"]')
-        .simulate('change', { target: { value: 'testDatasource2' } });
+      act(() => {
+        instance
+          .find('select[data-test-subj="datasource-switch"]')
+          .simulate('change', { target: { value: 'testDatasource2' } });
+      });
       expect(mockDatasource2.initialize).toHaveBeenCalled();
     });
 
@@ -395,9 +414,11 @@ describe('editor_frame', () => {
     });
 
     it('should initialize other visualization on switch', async () => {
-      instance
-        .find('select[data-test-subj="visualization-switch"]')
-        .simulate('change', { target: { value: 'testVis2' } });
+      act(() => {
+        instance
+          .find('select[data-test-subj="visualization-switch"]')
+          .simulate('change', { target: { value: 'testVis2' } });
+      });
       expect(mockVisualization2.initialize).toHaveBeenCalled();
     });
 
@@ -405,9 +426,11 @@ describe('editor_frame', () => {
       const initialState = {};
       mockVisualization2.initialize = () => initialState;
 
-      instance
-        .find('select[data-test-subj="visualization-switch"]')
-        .simulate('change', { target: { value: 'testVis2' } });
+      act(() => {
+        instance
+          .find('select[data-test-subj="visualization-switch"]')
+          .simulate('change', { target: { value: 'testVis2' } });
+      });
 
       expect(mockVisualization2.renderConfigPanel).toHaveBeenCalledWith(
         expect.any(Element),

From 48203ddc3cf51627b0c8d1431ed23f540c192b9b Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Fri, 10 May 2019 18:22:52 +0200
Subject: [PATCH 15/41] move layout to a separate component

---
 .../editor_frame/editor_frame.tsx             | 22 +++++++++---------
 .../editor_frame/frame_layout.tsx             | 23 +++++++++++++++++++
 2 files changed, 34 insertions(+), 11 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 3f92021f8c2ad..03acbc69a69c4 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -5,12 +5,11 @@
  */
 
 import React, { useEffect, useReducer, useMemo } from 'react';
-import { EuiSelect, EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
-import { Datasource, Visualization, DatasourceDataPanelProps } from '../../types';
-import { NativeRenderer } from '../../native_renderer';
+import { Datasource, Visualization } from '../../types';
 import { reducer, getInitialState } from '../state_management';
 import { DataPanelWrapper } from './data_panel_wrapper';
 import { ConfigPanelWrapper } from './config_panel_wrapper';
+import { FrameLayout } from './frame_layout';
 
 export interface EditorFrameProps {
   datasources: { [key: string]: Datasource };
@@ -62,8 +61,8 @@ export function EditorFrame(props: EditorFrameProps) {
   );
 
   return (
-    <EuiFlexGroup>
-      <EuiFlexItem>
+    <FrameLayout
+      dataPanel={
         <DataPanelWrapper
           activeDatasource={state.activeDatasource}
           datasourceIsLoading={state.datasourceIsLoading}
@@ -71,9 +70,10 @@ export function EditorFrame(props: EditorFrameProps) {
           datasources={props.datasources}
           dispatch={dispatch}
         />
-      </EuiFlexItem>
-      <EuiFlexItem>
-        {state.activeDatasource && !state.datasourceIsLoading && (
+      }
+      configPanel={
+        state.activeDatasource &&
+        !state.datasourceIsLoading && (
           <ConfigPanelWrapper
             visualizations={props.visualizations}
             activeVisualization={state.activeVisualization}
@@ -81,8 +81,8 @@ export function EditorFrame(props: EditorFrameProps) {
             dispatch={dispatch}
             visualizationState={state.visualizationState}
           />
-        )}
-      </EuiFlexItem>
-    </EuiFlexGroup>
+        )
+      }
+    />
   );
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx
new file mode 100644
index 0000000000000..182d0c9f0b592
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx
@@ -0,0 +1,23 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
+
+export interface FrameLayoutProps {
+  dataPanel: React.ReactNode;
+  configPanel: React.ReactNode;
+}
+
+export function FrameLayout(props: FrameLayoutProps) {
+  return (
+    <EuiFlexGroup>
+      {/* TODO style this and add workspace prop and loading flags */}
+      <EuiFlexItem>{props.dataPanel}</EuiFlexItem>
+      <EuiFlexItem>{props.configPanel}</EuiFlexItem>
+    </EuiFlexGroup>
+  );
+}

From ee6fdbcee6b72a6dffbab294dfca4b46680f6596 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Sun, 12 May 2019 13:11:04 +0200
Subject: [PATCH 16/41] clean up and improve typings

---
 .../editor_frame/editor_frame.tsx             |  4 ++--
 .../public/editor_frame_plugin/plugin.tsx     | 24 +++++++------------
 2 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 03acbc69a69c4..09eb5a7fcb0c7 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -12,8 +12,8 @@ import { ConfigPanelWrapper } from './config_panel_wrapper';
 import { FrameLayout } from './frame_layout';
 
 export interface EditorFrameProps {
-  datasources: { [key: string]: Datasource };
-  visualizations: { [key: string]: Visualization };
+  datasources: Record<string, Datasource>;
+  visualizations: Record<string, Visualization>;
 
   initialDatasource: string | null;
   initialVisualization: string | null;
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
index 34923e98ad770..8787eadd7dccc 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
@@ -10,15 +10,11 @@ import { Datasource, Visualization, EditorFrameSetup } from '../types';
 
 import { EditorFrame } from './editor_frame';
 
-class EditorFramePlugin {
+export class EditorFramePlugin {
   constructor() {}
 
-  private datasources: {
-    [key: string]: Datasource;
-  } = {};
-  private visualizations: {
-    [key: string]: Visualization;
-  } = {};
+  private datasources: Record<string, Datasource> = {};
+  private visualizations: Record<string, Visualization> = {};
 
   public setup(): EditorFrameSetup {
     return {
@@ -26,6 +22,9 @@ class EditorFramePlugin {
         let domElement: Element;
         return {
           mount: (element: Element) => {
+            if (domElement) {
+              unmountComponentAtNode(domElement);
+            }
             domElement = element;
             render(
               <EditorFrame
@@ -44,15 +43,8 @@ class EditorFramePlugin {
           },
         };
       },
-      registerDatasource: async (name, datasource) => {
-        // casting it to an unknown datasource. This doesn't introduce runtime errors
-        // because each type T is always also an unknown, but typescript won't do it
-        // on it's own because we are loosing type information here.
-        // So it's basically explicitly saying "I'm dropping the information about type T here
-        // because this information isn't useful to me." but without using any which can leak
-        // const state = await datasource.initialize();
-
-        this.datasources[name] = datasource as Datasource<unknown>;
+      registerDatasource: (name, datasource) => {
+        this.datasources[name] = datasource as Datasource<unknown, unknown>;
       },
       registerVisualization: (name, visualization) => {
         this.visualizations[name] = visualization as Visualization<unknown, unknown>;

From 396cf1c7bd217f1bb53be5f8f44652a75535aeff Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Sun, 12 May 2019 13:11:45 +0200
Subject: [PATCH 17/41] add tests for plugin itself

---
 .../editor_frame_plugin/plugin.test.tsx       | 112 ++++++++++++++++++
 1 file changed, 112 insertions(+)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx
new file mode 100644
index 0000000000000..1ca641b2e6e37
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx
@@ -0,0 +1,112 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { EditorFramePlugin } from './plugin';
+import { Visualization, Datasource } from '../types';
+
+const nextTick = () => new Promise(resolve => setTimeout(resolve));
+
+describe('editor_frame plugin', () => {
+  let pluginInstance: EditorFramePlugin;
+  let mountpoint: Element;
+
+  beforeEach(() => {
+    pluginInstance = new EditorFramePlugin();
+    mountpoint = document.createElement('div');
+  });
+
+  afterEach(() => {
+    mountpoint.remove();
+  });
+
+  it('should create an editor frame instance which mounts and unmounts', () => {
+    expect(() => {
+      const publicAPI = pluginInstance.setup();
+      const instance = publicAPI.createInstance({});
+      instance.mount(mountpoint);
+      instance.unmount();
+    }).not.toThrowError();
+  });
+
+  it('should render something in the provided dom element', () => {
+    const publicAPI = pluginInstance.setup();
+    const instance = publicAPI.createInstance({});
+    instance.mount(mountpoint);
+
+    expect(mountpoint.hasChildNodes()).toBe(true);
+
+    instance.unmount();
+  });
+
+  it('should not have child nodes after unmount', () => {
+    const publicAPI = pluginInstance.setup();
+    const instance = publicAPI.createInstance({});
+    instance.mount(mountpoint);
+    instance.unmount();
+
+    expect(mountpoint.hasChildNodes()).toBe(false);
+  });
+
+  it('should initialize and render provided datasource', async () => {
+    const publicAPI = pluginInstance.setup();
+    const mockDatasource = {
+      getDatasourceSuggestionsForField: jest.fn(),
+      getDatasourceSuggestionsFromCurrentState: jest.fn(),
+      getPersistableState: jest.fn(),
+      getPublicAPI: jest.fn(),
+      initialize: jest.fn(() => Promise.resolve()),
+      renderDataPanel: jest.fn(),
+      toExpression: jest.fn(),
+    };
+
+    publicAPI.registerDatasource('test', mockDatasource);
+
+    const instance = publicAPI.createInstance({});
+    instance.mount(mountpoint);
+
+    await nextTick();
+
+    expect(mockDatasource.initialize).toHaveBeenCalled();
+    expect(mockDatasource.renderDataPanel).toHaveBeenCalled();
+
+    instance.unmount();
+  });
+
+  it('should initialize visualization and render config panel', async () => {
+    const publicAPI = pluginInstance.setup();
+    const mockDatasource: Datasource = {
+      getDatasourceSuggestionsForField: jest.fn(),
+      getDatasourceSuggestionsFromCurrentState: jest.fn(),
+      getPersistableState: jest.fn(),
+      getPublicAPI: jest.fn(),
+      initialize: jest.fn(() => Promise.resolve()),
+      renderDataPanel: jest.fn(),
+      toExpression: jest.fn(),
+    };
+
+    const mockVisualization: Visualization = {
+      getMappingOfTableToRoles: jest.fn(),
+      getPersistableState: jest.fn(),
+      getSuggestions: jest.fn(),
+      initialize: jest.fn(),
+      renderConfigPanel: jest.fn(),
+      toExpression: jest.fn(),
+    };
+
+    publicAPI.registerDatasource('test', mockDatasource);
+    publicAPI.registerVisualization('test', mockVisualization);
+
+    const instance = publicAPI.createInstance({});
+    instance.mount(mountpoint);
+
+    await nextTick();
+
+    expect(mockVisualization.initialize).toHaveBeenCalled();
+    expect(mockVisualization.renderConfigPanel).toHaveBeenCalled();
+
+    instance.unmount();
+  });
+});

From 6f8c2950552f23d201174c4594269b5159261307 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Sun, 12 May 2019 13:40:48 +0200
Subject: [PATCH 18/41] start implementing suggestion control flow

---
 .../editor_frame/editor_frame.tsx             | 29 ++++++++++++++++++-
 .../editor_frame/frame_layout.tsx             |  6 +++-
 x-pack/plugins/lens/public/types.ts           |  8 ++---
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 09eb5a7fcb0c7..657958e16edc1 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -5,7 +5,7 @@
  */
 
 import React, { useEffect, useReducer, useMemo } from 'react';
-import { Datasource, Visualization } from '../../types';
+import { Datasource, Visualization, TableColumn } from '../../types';
 import { reducer, getInitialState } from '../state_management';
 import { DataPanelWrapper } from './data_panel_wrapper';
 import { ConfigPanelWrapper } from './config_panel_wrapper';
@@ -60,6 +60,32 @@ export function EditorFrame(props: EditorFrameProps) {
     [props.datasources, state.datasourceIsLoading, state.activeDatasource, state.datasourceState]
   );
 
+  if (state.activeDatasource && !state.datasourceIsLoading) {
+    const datasourceTableSuggestions = props.datasources[
+      state.activeDatasource
+    ].getDatasourceSuggestionsFromCurrentState(state.datasourceState);
+
+    const visualizationSuggestions = Object.values(props.visualizations)
+      .flatMap(visualization => {
+        const datasourceTableMetas: Record<string, TableColumn[]> = {};
+
+        datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
+          datasourceTableMetas[datasourceSuggestionId] = tableColumns;
+        });
+
+        return visualization.getSuggestions({
+          tableColumns: datasourceTableMetas,
+          roles: state.activeVisualization
+            ? props.visualizations[state.activeVisualization].getMappingOfTableToRoles(
+                state.visualizationState[state.activeVisualization],
+                datasourcePublicAPI!
+              )
+            : [],
+        });
+      })
+      .sort(({ score: scoreA }, { score: scoreB }) => scoreA - scoreB);
+  }
+
   return (
     <FrameLayout
       dataPanel={
@@ -83,6 +109,7 @@ export function EditorFrame(props: EditorFrameProps) {
           />
         )
       }
+      suggestionsPanel={<div>Suggestions</div>}
     />
   );
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx
index 182d0c9f0b592..7c7d0f23ca9cd 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx
@@ -10,6 +10,7 @@ import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
 export interface FrameLayoutProps {
   dataPanel: React.ReactNode;
   configPanel: React.ReactNode;
+  suggestionsPanel?: React.ReactNode;
 }
 
 export function FrameLayout(props: FrameLayoutProps) {
@@ -17,7 +18,10 @@ export function FrameLayout(props: FrameLayoutProps) {
     <EuiFlexGroup>
       {/* TODO style this and add workspace prop and loading flags */}
       <EuiFlexItem>{props.dataPanel}</EuiFlexItem>
-      <EuiFlexItem>{props.configPanel}</EuiFlexItem>
+      <EuiFlexItem>
+        {props.configPanel}
+        {props.suggestionsPanel}
+      </EuiFlexItem>
     </EuiFlexGroup>
   );
 }
diff --git a/x-pack/plugins/lens/public/types.ts b/x-pack/plugins/lens/public/types.ts
index a0328cb3bd988..4943d9cdf80d2 100644
--- a/x-pack/plugins/lens/public/types.ts
+++ b/x-pack/plugins/lens/public/types.ts
@@ -31,14 +31,14 @@ export type DimensionRole =
   | 'size'
   | string; // Some visualizations will use custom names that have other meaning
 
-export interface TableColumns {
+export interface TableColumn {
   columnId: string;
   operation: Operation;
 }
 
 export interface DatasourceSuggestion<T = unknown> {
   state: T;
-  tableColumns: TableColumns[];
+  tableColumns: TableColumn[];
 }
 
 /**
@@ -132,9 +132,9 @@ export interface VisualizationProps<T = unknown> {
 
 export interface SuggestionRequest<T = unknown> {
   // Roles currently being used
-  roles: DimensionRole[];
+  roles?: DimensionRole[];
   // It is up to the Visualization to rank these tables
-  tableColumns: { [datasourceSuggestionId: string]: TableColumns };
+  tableColumns: { [datasourceSuggestionId: string]: TableColumn[] };
   state?: T; // State is only passed if the visualization is active
 }
 

From b61008c0cdfa87399343c822999e63ceaac2ef5e Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Sun, 12 May 2019 17:50:05 +0200
Subject: [PATCH 19/41] refactor control flow and add visualization switch
 action dispatch

---
 .../editor_frame/editor_frame.tsx             | 128 ++++++++++++------
 .../editor_frame/suggestion_panel_wrapper.tsx |  37 +++++
 2 files changed, 121 insertions(+), 44 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 657958e16edc1..3733bcb5d1278 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -5,11 +5,19 @@
  */
 
 import React, { useEffect, useReducer, useMemo } from 'react';
-import { Datasource, Visualization, TableColumn } from '../../types';
+import {
+  Datasource,
+  Visualization,
+  TableColumn,
+  DatasourcePublicAPI,
+  DatasourceSuggestion,
+  DimensionRole,
+} from '../../types';
 import { reducer, getInitialState } from '../state_management';
 import { DataPanelWrapper } from './data_panel_wrapper';
 import { ConfigPanelWrapper } from './config_panel_wrapper';
 import { FrameLayout } from './frame_layout';
+import { SuggestionPanelWrapper } from './suggestion_panel_wrapper';
 
 export interface EditorFrameProps {
   datasources: Record<string, Datasource>;
@@ -60,46 +68,37 @@ export function EditorFrame(props: EditorFrameProps) {
     [props.datasources, state.datasourceIsLoading, state.activeDatasource, state.datasourceState]
   );
 
-  if (state.activeDatasource && !state.datasourceIsLoading) {
-    const datasourceTableSuggestions = props.datasources[
-      state.activeDatasource
-    ].getDatasourceSuggestionsFromCurrentState(state.datasourceState);
+  const datasourceIsActive = Boolean(state.activeDatasource && !state.datasourceIsLoading);
 
-    const visualizationSuggestions = Object.values(props.visualizations)
-      .flatMap(visualization => {
-        const datasourceTableMetas: Record<string, TableColumn[]> = {};
-
-        datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
-          datasourceTableMetas[datasourceSuggestionId] = tableColumns;
-        });
-
-        return visualization.getSuggestions({
-          tableColumns: datasourceTableMetas,
-          roles: state.activeVisualization
-            ? props.visualizations[state.activeVisualization].getMappingOfTableToRoles(
-                state.visualizationState[state.activeVisualization],
-                datasourcePublicAPI!
-              )
-            : [],
-        });
-      })
-      .sort(({ score: scoreA }, { score: scoreB }) => scoreA - scoreB);
-  }
+  if (datasourceIsActive) {
+    const suggestions =
+      state.activeDatasource && !state.datasourceIsLoading
+        ? getSuggestions(
+            props.datasources[state.activeDatasource].getDatasourceSuggestionsFromCurrentState(
+              state.activeDatasource
+            ),
+            props.visualizations,
+            state.activeVisualization
+              ? props.visualizations[state.activeVisualization].getMappingOfTableToRoles(
+                  state.visualizationState[state.activeVisualization],
+                  datasourcePublicAPI!
+                )
+              : []
+          )
+        : [];
 
-  return (
-    <FrameLayout
-      dataPanel={
-        <DataPanelWrapper
-          activeDatasource={state.activeDatasource}
-          datasourceIsLoading={state.datasourceIsLoading}
-          datasourceState={state.datasourceState}
-          datasources={props.datasources}
-          dispatch={dispatch}
-        />
-      }
-      configPanel={
-        state.activeDatasource &&
-        !state.datasourceIsLoading && (
+    return (
+      <FrameLayout
+        dataPanel={
+          <DataPanelWrapper
+            activeDatasource={state.activeDatasource}
+            datasourceIsLoading={state.datasourceIsLoading}
+            datasourceState={state.datasourceState}
+            datasources={props.datasources}
+            dispatch={dispatch}
+          />
+        }
+        configPanel={
           <ConfigPanelWrapper
             visualizations={props.visualizations}
             activeVisualization={state.activeVisualization}
@@ -107,9 +106,50 @@ export function EditorFrame(props: EditorFrameProps) {
             dispatch={dispatch}
             visualizationState={state.visualizationState}
           />
-        )
-      }
-      suggestionsPanel={<div>Suggestions</div>}
-    />
-  );
+        }
+        suggestionsPanel={<SuggestionPanelWrapper suggestions={suggestions} dispatch={dispatch} />}
+      />
+    );
+  } else {
+    return (
+      <FrameLayout
+        dataPanel={
+          <DataPanelWrapper
+            activeDatasource={state.activeDatasource}
+            datasourceIsLoading={state.datasourceIsLoading}
+            datasourceState={state.datasourceState}
+            datasources={props.datasources}
+            dispatch={dispatch}
+          />
+        }
+        configPanel={null}
+        suggestionsPanel={null}
+      />
+    );
+  }
+}
+
+function getSuggestions(
+  datasourceTableSuggestions: DatasourceSuggestion[],
+  visualizations: Record<string, Visualization>,
+  currentColumnRoles: DimensionRole[]
+) {
+  return Object.entries(visualizations)
+    .map(([visualizationId, visualization]) => {
+      const datasourceTableMetas: Record<string, TableColumn[]> = {};
+      datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
+        datasourceTableMetas[datasourceSuggestionId] = tableColumns;
+      });
+      return visualization
+        .getSuggestions({
+          tableColumns: datasourceTableMetas,
+          roles: currentColumnRoles,
+        })
+        .map(suggestion => ({
+          ...suggestion,
+          visualizationId,
+        }));
+    })
+    .flat()
+    .sort(({ score: scoreA }, { score: scoreB }) => scoreA - scoreB);
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
new file mode 100644
index 0000000000000..5e3df79f12dd0
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+import { Action } from '../state_management';
+import { VisualizationSuggestion } from '../../types';
+
+interface SuggestionPanelWrapperProps {
+  suggestions: Array<VisualizationSuggestion & { visualizationId: string }>;
+  dispatch: (action: Action) => void;
+}
+
+export function SuggestionPanelWrapper(props: SuggestionPanelWrapperProps) {
+  return (
+    <>
+      <h2>Suggestions</h2>
+      {(props.suggestions || []).map(suggestion => {
+        return (
+          <button
+            onClick={() => {
+              props.dispatch({
+                type: 'SWITCH_VISUALIZATION',
+                newVisulizationId: suggestion.visualizationId,
+                initialState: suggestion.state,
+              });
+            }}
+          >
+            {suggestion.title}
+          </button>
+        );
+      })}
+    </>
+  );
+}

From cae9fe0042308a6edd48466c10dc78e95714f6cc Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Sun, 12 May 2019 19:02:12 +0200
Subject: [PATCH 20/41] move mock builder into separate file and fix tests

---
 .../editor_frame/editor_frame.test.tsx        | 28 +++----------
 .../editor_frame/editor_frame.tsx             | 39 ++++++++++--------
 .../editor_frame_plugin/mock_extensions.ts    | 41 +++++++++++++++++++
 .../editor_frame_plugin/plugin.test.tsx       | 33 ++-------------
 4 files changed, 71 insertions(+), 70 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
index 4b334a541c984..2de698e996ef5 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
@@ -9,29 +9,11 @@ import { mount, ReactWrapper } from 'enzyme';
 import { EditorFrame } from './editor_frame';
 import { Visualization, Datasource, DatasourcePublicAPI } from '../../types';
 import { act } from 'react-dom/test-utils';
+import { createMockVisualization, createMockDatasource } from '../mock_extensions';
 
 const nextTick = () => new Promise(resolve => setTimeout(resolve));
 
 describe('editor_frame', () => {
-  const getMockVisualization = () => ({
-    getMappingOfTableToRoles: jest.fn(),
-    getPersistableState: jest.fn(),
-    getSuggestions: jest.fn(),
-    initialize: jest.fn(),
-    renderConfigPanel: jest.fn(),
-    toExpression: jest.fn(),
-  });
-
-  const getMockDatasource = () => ({
-    getDatasourceSuggestionsForField: jest.fn(),
-    getDatasourceSuggestionsFromCurrentState: jest.fn(),
-    getPersistableState: jest.fn(),
-    getPublicAPI: jest.fn(),
-    initialize: jest.fn(() => Promise.resolve()),
-    renderDataPanel: jest.fn(),
-    toExpression: jest.fn(),
-  });
-
   let mockVisualization: Visualization;
   let mockDatasource: Datasource;
 
@@ -39,11 +21,11 @@ describe('editor_frame', () => {
   let mockDatasource2: Datasource;
 
   beforeEach(() => {
-    mockVisualization = getMockVisualization();
-    mockVisualization2 = getMockVisualization();
+    mockVisualization = createMockVisualization();
+    mockVisualization2 = createMockVisualization();
 
-    mockDatasource = getMockDatasource();
-    mockDatasource2 = getMockDatasource();
+    mockDatasource = createMockDatasource();
+    mockDatasource2 = createMockDatasource();
   });
 
   describe('initialization', () => {
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 3733bcb5d1278..210d1736e15bf 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -134,22 +134,25 @@ function getSuggestions(
   visualizations: Record<string, Visualization>,
   currentColumnRoles: DimensionRole[]
 ) {
-  return Object.entries(visualizations)
-    .map(([visualizationId, visualization]) => {
-      const datasourceTableMetas: Record<string, TableColumn[]> = {};
-      datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
-        datasourceTableMetas[datasourceSuggestionId] = tableColumns;
-      });
-      return visualization
-        .getSuggestions({
-          tableColumns: datasourceTableMetas,
-          roles: currentColumnRoles,
-        })
-        .map(suggestion => ({
-          ...suggestion,
-          visualizationId,
-        }));
-    })
-    .flat()
-    .sort(({ score: scoreA }, { score: scoreB }) => scoreA - scoreB);
+  return (
+    Object.entries(visualizations)
+      .map(([visualizationId, visualization]) => {
+        const datasourceTableMetas: Record<string, TableColumn[]> = {};
+        datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
+          datasourceTableMetas[datasourceSuggestionId] = tableColumns;
+        });
+        return visualization
+          .getSuggestions({
+            tableColumns: datasourceTableMetas,
+            roles: currentColumnRoles,
+          })
+          .map(suggestion => ({
+            ...suggestion,
+            visualizationId,
+          }));
+      })
+      // TODO why is flatMap not available here?
+      .reduce((globalList, currentList) => [...globalList, ...currentList], [])
+      .sort(({ score: scoreA }, { score: scoreB }) => scoreA - scoreB)
+  );
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
new file mode 100644
index 0000000000000..a61b531b4494a
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { DatasourcePublicAPI } from '../types';
+
+export const createMockVisualization = () => ({
+  getMappingOfTableToRoles: jest.fn(() => []),
+  getPersistableState: jest.fn(() => ({})),
+  getSuggestions: jest.fn(() => []),
+  initialize: jest.fn(() => {}),
+  renderConfigPanel: jest.fn(),
+  toExpression: jest.fn(() => ''),
+});
+
+export const createMockDatasource = () => {
+  const publicAPIMock: DatasourcePublicAPI = {
+    getTableSpec: jest.fn(() => []),
+    getOperationForColumnId: jest.fn(() => null),
+    renderDimensionPanel: jest.fn(),
+    removeColumnInTableSpec: jest.fn(),
+    moveColumnTo: jest.fn(),
+    duplicateColumn: jest.fn(),
+  };
+
+  return {
+    getDatasourceSuggestionsForField: jest.fn(() => []),
+    getDatasourceSuggestionsFromCurrentState: jest.fn(() => []),
+    getPersistableState: jest.fn(),
+    getPublicAPI: jest.fn(() => publicAPIMock),
+    initialize: jest.fn(() => Promise.resolve()),
+    renderDataPanel: jest.fn(),
+    toExpression: jest.fn(() => ''),
+
+    // this is an additional property which doesn't exist on real datasources
+    // but can be used to validate whether specific API mock functions are called
+    publicAPIMock,
+  };
+};
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx
index 1ca641b2e6e37..0ba4bd168216a 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx
@@ -5,7 +5,7 @@
  */
 
 import { EditorFramePlugin } from './plugin';
-import { Visualization, Datasource } from '../types';
+import { createMockDatasource, createMockVisualization } from './mock_extensions';
 
 const nextTick = () => new Promise(resolve => setTimeout(resolve));
 
@@ -51,17 +51,8 @@ describe('editor_frame plugin', () => {
   });
 
   it('should initialize and render provided datasource', async () => {
+    const mockDatasource = createMockDatasource();
     const publicAPI = pluginInstance.setup();
-    const mockDatasource = {
-      getDatasourceSuggestionsForField: jest.fn(),
-      getDatasourceSuggestionsFromCurrentState: jest.fn(),
-      getPersistableState: jest.fn(),
-      getPublicAPI: jest.fn(),
-      initialize: jest.fn(() => Promise.resolve()),
-      renderDataPanel: jest.fn(),
-      toExpression: jest.fn(),
-    };
-
     publicAPI.registerDatasource('test', mockDatasource);
 
     const instance = publicAPI.createInstance({});
@@ -76,25 +67,9 @@ describe('editor_frame plugin', () => {
   });
 
   it('should initialize visualization and render config panel', async () => {
+    const mockDatasource = createMockDatasource();
+    const mockVisualization = createMockVisualization();
     const publicAPI = pluginInstance.setup();
-    const mockDatasource: Datasource = {
-      getDatasourceSuggestionsForField: jest.fn(),
-      getDatasourceSuggestionsFromCurrentState: jest.fn(),
-      getPersistableState: jest.fn(),
-      getPublicAPI: jest.fn(),
-      initialize: jest.fn(() => Promise.resolve()),
-      renderDataPanel: jest.fn(),
-      toExpression: jest.fn(),
-    };
-
-    const mockVisualization: Visualization = {
-      getMappingOfTableToRoles: jest.fn(),
-      getPersistableState: jest.fn(),
-      getSuggestions: jest.fn(),
-      initialize: jest.fn(),
-      renderConfigPanel: jest.fn(),
-      toExpression: jest.fn(),
-    };
 
     publicAPI.registerDatasource('test', mockDatasource);
     publicAPI.registerVisualization('test', mockVisualization);

From 3f48c98d48cccd713ea58c3d6cbeceb323c641b8 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Mon, 13 May 2019 18:40:08 +0200
Subject: [PATCH 21/41] remove unused import

---
 .../public/editor_frame_plugin/editor_frame/editor_frame.tsx     | 1 -
 1 file changed, 1 deletion(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 210d1736e15bf..70d8f0c150526 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -9,7 +9,6 @@ import {
   Datasource,
   Visualization,
   TableColumn,
-  DatasourcePublicAPI,
   DatasourceSuggestion,
   DimensionRole,
 } from '../../types';

From 7cdea68ad3bebd689d48799eb666a66d0437fbf2 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Mon, 13 May 2019 14:09:33 -0400
Subject: [PATCH 22/41] add tests for suggestions

---
 .../editor_frame/editor_frame.test.tsx        | 164 ++++++++++++++++++
 .../editor_frame/editor_frame.tsx             |   5 +-
 .../editor_frame/suggestion_panel_wrapper.tsx |  13 +-
 .../state_management.test.ts                  |  27 +++
 .../editor_frame_plugin/state_management.ts   |   2 +
 x-pack/plugins/lens/public/types.ts           |   4 +-
 6 files changed, 209 insertions(+), 6 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
index 2de698e996ef5..bf53762826b97 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
@@ -420,4 +420,168 @@ describe('editor_frame', () => {
       );
     });
   });
+
+  describe('suggestions', () => {
+    it('should fetch suggestions of currently active datasource', async () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+            testDatasource2: mockDatasource2,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      expect(mockDatasource.getDatasourceSuggestionsFromCurrentState).toHaveBeenCalled();
+      expect(mockDatasource2.getDatasourceSuggestionsFromCurrentState).not.toHaveBeenCalled();
+    });
+
+    it('should fetch suggestions of all visualizations', async () => {
+      mount(
+        <EditorFrame
+          visualizations={{
+            testVis: mockVisualization,
+            testVis2: mockVisualization2,
+          }}
+          datasources={{
+            testDatasource: mockDatasource,
+            testDatasource2: mockDatasource2,
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      expect(mockVisualization.getSuggestions).toHaveBeenCalled();
+      expect(mockVisualization2.getSuggestions).toHaveBeenCalled();
+    });
+
+    it('should display suggestions in descending order', async () => {
+      const instance = mount(
+        <EditorFrame
+          visualizations={{
+            testVis: {
+              ...mockVisualization,
+              getSuggestions: () => [
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.5,
+                  state: {},
+                  title: 'Suggestion2',
+                },
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.8,
+                  state: {},
+                  title: 'Suggestion1',
+                },
+              ],
+            },
+            testVis2: {
+              ...mockVisualization,
+              getSuggestions: () => [
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.4,
+                  state: {},
+                  title: 'Suggestion4',
+                },
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.45,
+                  state: {},
+                  title: 'Suggestion3',
+                },
+              ],
+            },
+          }}
+          datasources={{
+            testDatasource: {
+              ...mockDatasource,
+              getDatasourceSuggestionsFromCurrentState: () => [{ state: {}, tableColumns: [] }],
+            },
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis"
+        />
+      );
+
+      await nextTick();
+
+      // TODO why is this necessary?
+      instance.update();
+      const suggestions = instance.find('[data-test-subj="suggestion"]');
+      expect(suggestions.map(el => el.text())).toEqual([
+        'Suggestion1',
+        'Suggestion2',
+        'Suggestion3',
+        'Suggestion4',
+      ]);
+    });
+
+    it('should switch to suggested visualization', async () => {
+      const newDatasourceState = {};
+      const suggestionVisState = {};
+      const instance = mount(
+        <EditorFrame
+          visualizations={{
+            testVis: {
+              ...mockVisualization,
+              getSuggestions: () => [
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.8,
+                  state: suggestionVisState,
+                  title: 'Suggestion1',
+                },
+              ],
+            },
+            testVis2: mockVisualization2,
+          }}
+          datasources={{
+            testDatasource: {
+              ...mockDatasource,
+              getDatasourceSuggestionsFromCurrentState: () => [
+                { state: newDatasourceState, tableColumns: [] },
+              ],
+            },
+          }}
+          initialDatasource="testDatasource"
+          initialVisualization="testVis2"
+        />
+      );
+
+      await nextTick();
+
+      // TODO why is this necessary?
+      instance.update();
+
+      act(() => {
+        instance.find('[data-test-subj="suggestion"]').simulate('click');
+      });
+
+      expect(mockVisualization.renderConfigPanel).toHaveBeenCalledTimes(1);
+      expect(mockVisualization.renderConfigPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({
+          state: suggestionVisState,
+        })
+      );
+      expect(mockDatasource.renderDataPanel).toHaveBeenLastCalledWith(
+        expect.any(Element),
+        expect.objectContaining({
+          state: newDatasourceState,
+        })
+      );
+    });
+  });
 });
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 70d8f0c150526..d2b86c0cd4f05 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -145,13 +145,14 @@ function getSuggestions(
             tableColumns: datasourceTableMetas,
             roles: currentColumnRoles,
           })
-          .map(suggestion => ({
+          .map(({ datasourceSuggestionId, ...suggestion }) => ({
             ...suggestion,
             visualizationId,
+            datasourceState: datasourceTableSuggestions[datasourceSuggestionId].state,
           }));
       })
       // TODO why is flatMap not available here?
       .reduce((globalList, currentList) => [...globalList, ...currentList], [])
-      .sort(({ score: scoreA }, { score: scoreB }) => scoreA - scoreB)
+      .sort(({ score: scoreA }, { score: scoreB }) => scoreB - scoreA)
   );
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
index 5e3df79f12dd0..623eb1257551c 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
@@ -8,8 +8,13 @@ import React from 'react';
 import { Action } from '../state_management';
 import { VisualizationSuggestion } from '../../types';
 
+export type Suggestion = Pick<
+  VisualizationSuggestion,
+  Exclude<keyof VisualizationSuggestion, 'datasourceSuggestionId'>
+> & { visualizationId: string; datasourceState: unknown };
+
 interface SuggestionPanelWrapperProps {
-  suggestions: Array<VisualizationSuggestion & { visualizationId: string }>;
+  suggestions: Suggestion[];
   dispatch: (action: Action) => void;
 }
 
@@ -17,14 +22,18 @@ export function SuggestionPanelWrapper(props: SuggestionPanelWrapperProps) {
   return (
     <>
       <h2>Suggestions</h2>
-      {(props.suggestions || []).map(suggestion => {
+      {(props.suggestions || []).map((suggestion, index) => {
         return (
           <button
+            key={index}
+            data-test-subj={`suggestion`}
             onClick={() => {
+              // TODO single action for that?
               props.dispatch({
                 type: 'SWITCH_VISUALIZATION',
                 newVisulizationId: suggestion.visualizationId,
                 initialState: suggestion.state,
+                datasourceState: suggestion.datasourceState,
               });
             }}
           >
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
index fb96817685c45..a2de69977d458 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
@@ -117,6 +117,33 @@ describe('editor_frame state management', () => {
       expect(newState.visualizationState.testVis2).toBe(newVisState);
     });
 
+    it('should should switch active visualization and update datasource state', () => {
+      const testVisState = {};
+      const newVisState = {};
+      const newDatasourceState = {};
+      const newState = reducer(
+        {
+          activeDatasource: 'testDatasource',
+          activeVisualization: 'testVis',
+          datasourceIsLoading: false,
+          datasourceState: {},
+          visualizationState: {
+            testVis: testVisState,
+          },
+        },
+        {
+          type: 'SWITCH_VISUALIZATION',
+          newVisulizationId: 'testVis2',
+          initialState: newVisState,
+          datasourceState: newDatasourceState,
+        }
+      );
+
+      expect(newState.visualizationState.testVis).toBe(testVisState);
+      expect(newState.visualizationState.testVis2).toBe(newVisState);
+      expect(newState.datasourceState).toBe(newDatasourceState);
+    });
+
     it('should should switch active datasource and purge visualization state', () => {
       const newState = reducer(
         {
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts
index e35e258d1153c..824c5fbf94f29 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts
@@ -30,6 +30,7 @@ export type Action =
       type: 'SWITCH_VISUALIZATION';
       newVisulizationId: string;
       initialState: unknown;
+      datasourceState?: unknown;
     }
   | {
       type: 'SWITCH_DATASOURCE';
@@ -72,6 +73,7 @@ export const reducer = (state: EditorFrameState, action: Action): EditorFrameSta
           ...state.visualizationState,
           [action.newVisulizationId]: action.initialState,
         },
+        datasourceState: action.datasourceState ? action.datasourceState : state.datasourceState,
       };
     case 'UPDATE_DATASOURCE_STATE':
       return {
diff --git a/x-pack/plugins/lens/public/types.ts b/x-pack/plugins/lens/public/types.ts
index 4943d9cdf80d2..27fe9542987c0 100644
--- a/x-pack/plugins/lens/public/types.ts
+++ b/x-pack/plugins/lens/public/types.ts
@@ -134,7 +134,7 @@ export interface SuggestionRequest<T = unknown> {
   // Roles currently being used
   roles?: DimensionRole[];
   // It is up to the Visualization to rank these tables
-  tableColumns: { [datasourceSuggestionId: string]: TableColumn[] };
+  tableColumns: Record<number, TableColumn[]>;
   state?: T; // State is only passed if the visualization is active
 }
 
@@ -142,7 +142,7 @@ export interface VisualizationSuggestion<T = unknown> {
   score: number;
   title: string;
   state: T;
-  datasourceSuggestionId: string;
+  datasourceSuggestionId: number;
 }
 
 export interface Visualization<T = unknown, P = unknown> {

From d63d2f905d24140edc167e8a332e169193c0d5d9 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Mon, 13 May 2019 14:29:00 -0400
Subject: [PATCH 23/41] correct typo

---
 .../editor_frame/suggestion_panel_wrapper.tsx                   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
index 623eb1257551c..438d6ae3eac44 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
@@ -26,7 +26,7 @@ export function SuggestionPanelWrapper(props: SuggestionPanelWrapperProps) {
         return (
           <button
             key={index}
-            data-test-subj={`suggestion`}
+            data-test-subj="suggestion"
             onClick={() => {
               // TODO single action for that?
               props.dispatch({

From 52cad4f424b89cc3616c6eca7af26c51a417414a Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 14 May 2019 18:13:19 -0400
Subject: [PATCH 24/41] rework suggestion stuff

---
 .../editor_frame/editor_frame.tsx             | 79 ++++------------
 .../editor_frame/suggestion_panel.tsx         | 94 +++++++++++++++++++
 .../editor_frame/suggestion_panel_wrapper.tsx | 46 ---------
 3 files changed, 113 insertions(+), 106 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
 delete mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index d2b86c0cd4f05..835935bb3aa12 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -5,18 +5,12 @@
  */
 
 import React, { useEffect, useReducer, useMemo } from 'react';
-import {
-  Datasource,
-  Visualization,
-  TableColumn,
-  DatasourceSuggestion,
-  DimensionRole,
-} from '../../types';
+import { Datasource, Visualization } from '../../types';
 import { reducer, getInitialState } from '../state_management';
 import { DataPanelWrapper } from './data_panel_wrapper';
 import { ConfigPanelWrapper } from './config_panel_wrapper';
 import { FrameLayout } from './frame_layout';
-import { SuggestionPanelWrapper } from './suggestion_panel_wrapper';
+import { SuggestionPanel } from './suggestion_panel';
 
 export interface EditorFrameProps {
   datasources: Record<string, Datasource>;
@@ -67,33 +61,15 @@ export function EditorFrame(props: EditorFrameProps) {
     [props.datasources, state.datasourceIsLoading, state.activeDatasource, state.datasourceState]
   );
 
-  const datasourceIsActive = Boolean(state.activeDatasource && !state.datasourceIsLoading);
-
-  if (datasourceIsActive) {
-    const suggestions =
-      state.activeDatasource && !state.datasourceIsLoading
-        ? getSuggestions(
-            props.datasources[state.activeDatasource].getDatasourceSuggestionsFromCurrentState(
-              state.activeDatasource
-            ),
-            props.visualizations,
-            state.activeVisualization
-              ? props.visualizations[state.activeVisualization].getMappingOfTableToRoles(
-                  state.visualizationState[state.activeVisualization],
-                  datasourcePublicAPI!
-                )
-              : []
-          )
-        : [];
-
+  if (state.activeDatasource && !state.datasourceIsLoading) {
     return (
       <FrameLayout
         dataPanel={
           <DataPanelWrapper
+            datasources={props.datasources}
             activeDatasource={state.activeDatasource}
-            datasourceIsLoading={state.datasourceIsLoading}
             datasourceState={state.datasourceState}
-            datasources={props.datasources}
+            datasourceIsLoading={state.datasourceIsLoading}
             dispatch={dispatch}
           />
         }
@@ -101,12 +77,24 @@ export function EditorFrame(props: EditorFrameProps) {
           <ConfigPanelWrapper
             visualizations={props.visualizations}
             activeVisualization={state.activeVisualization}
+            visualizationState={state.visualizationState}
             datasourcePublicAPI={datasourcePublicAPI!}
             dispatch={dispatch}
-            visualizationState={state.visualizationState}
           />
         }
-        suggestionsPanel={<SuggestionPanelWrapper suggestions={suggestions} dispatch={dispatch} />}
+        suggestionsPanel={
+          <SuggestionPanel
+            activeDatasource={props.datasources[state.activeDatasource]}
+            activeVisualization={state.activeVisualization}
+            datasourcePublicAPI={datasourcePublicAPI!}
+            datasourceState={state.datasourceState}
+            visualizationState={
+              state.activeVisualization && state.visualizationState[state.activeVisualization]
+            }
+            visualizations={props.visualizations}
+            dispatch={dispatch}
+          />
+        }
       />
     );
   } else {
@@ -127,32 +115,3 @@ export function EditorFrame(props: EditorFrameProps) {
     );
   }
 }
-
-function getSuggestions(
-  datasourceTableSuggestions: DatasourceSuggestion[],
-  visualizations: Record<string, Visualization>,
-  currentColumnRoles: DimensionRole[]
-) {
-  return (
-    Object.entries(visualizations)
-      .map(([visualizationId, visualization]) => {
-        const datasourceTableMetas: Record<string, TableColumn[]> = {};
-        datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
-          datasourceTableMetas[datasourceSuggestionId] = tableColumns;
-        });
-        return visualization
-          .getSuggestions({
-            tableColumns: datasourceTableMetas,
-            roles: currentColumnRoles,
-          })
-          .map(({ datasourceSuggestionId, ...suggestion }) => ({
-            ...suggestion,
-            visualizationId,
-            datasourceState: datasourceTableSuggestions[datasourceSuggestionId].state,
-          }));
-      })
-      // TODO why is flatMap not available here?
-      .reduce((globalList, currentList) => [...globalList, ...currentList], [])
-      .sort(({ score: scoreA }, { score: scoreB }) => scoreB - scoreA)
-  );
-}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
new file mode 100644
index 0000000000000..68162eb3ad95f
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
@@ -0,0 +1,94 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+import { Action } from '../state_management';
+import {
+  Datasource,
+  Visualization,
+  DatasourcePublicAPI,
+  DatasourceSuggestion,
+  DimensionRole,
+  TableColumn,
+} from '../../types';
+
+interface SuggestionPanelWrapper {
+  activeDatasource: Datasource;
+  datasourceState: unknown;
+  activeVisualization: string | null;
+  visualizations: Record<string, Visualization>;
+  visualizationState: unknown;
+  datasourcePublicAPI: DatasourcePublicAPI;
+  dispatch: (action: Action) => void;
+}
+
+export function SuggestionPanel(props: SuggestionPanelWrapper) {
+  const currentDatasource = props.activeDatasource;
+  const datasourceSuggestions = currentDatasource.getDatasourceSuggestionsFromCurrentState(
+    props.datasourceState
+  );
+  const roleMapping = props.activeVisualization
+    ? props.visualizations[props.activeVisualization].getMappingOfTableToRoles(
+        props.visualizationState,
+        props.datasourcePublicAPI
+      )
+    : [];
+  const suggestions = getSuggestions(datasourceSuggestions, props.visualizations, roleMapping);
+  return (
+    <>
+      <h2>Suggestions</h2>
+      {suggestions.map((suggestion, index) => {
+        return (
+          <button
+            key={index}
+            data-test-subj="suggestion"
+            onClick={() => {
+              // TODO single action for that?
+              props.dispatch({
+                type: 'SWITCH_VISUALIZATION',
+                newVisulizationId: suggestion.visualizationId,
+                initialState: suggestion.state,
+                datasourceState: suggestion.datasourceState,
+              });
+            }}
+          >
+            {suggestion.title}
+          </button>
+        );
+      })}
+    </>
+  );
+}
+
+function getSuggestions(
+  datasourceTableSuggestions: DatasourceSuggestion[],
+  visualizations: Record<string, Visualization>,
+  currentColumnRoles: DimensionRole[]
+) {
+  const datasourceTableMetas: Record<string, TableColumn[]> = {};
+  datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
+    datasourceTableMetas[datasourceSuggestionId] = tableColumns;
+  });
+
+  return (
+    Object.entries(visualizations)
+      .map(([visualizationId, visualization]) => {
+        return visualization
+          .getSuggestions({
+            tableColumns: datasourceTableMetas,
+            roles: currentColumnRoles,
+          })
+          .map(({ datasourceSuggestionId, ...suggestion }) => ({
+            ...suggestion,
+            visualizationId,
+            datasourceState: datasourceTableSuggestions[datasourceSuggestionId].state,
+          }));
+      })
+      // TODO why is flatMap not available here?
+      .reduce((globalList, currentList) => [...globalList, ...currentList], [])
+      .sort(({ score: scoreA }, { score: scoreB }) => scoreB - scoreA)
+  );
+}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
deleted file mode 100644
index 438d6ae3eac44..0000000000000
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel_wrapper.tsx
+++ /dev/null
@@ -1,46 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License;
- * you may not use this file except in compliance with the Elastic License.
- */
-
-import React from 'react';
-import { Action } from '../state_management';
-import { VisualizationSuggestion } from '../../types';
-
-export type Suggestion = Pick<
-  VisualizationSuggestion,
-  Exclude<keyof VisualizationSuggestion, 'datasourceSuggestionId'>
-> & { visualizationId: string; datasourceState: unknown };
-
-interface SuggestionPanelWrapperProps {
-  suggestions: Suggestion[];
-  dispatch: (action: Action) => void;
-}
-
-export function SuggestionPanelWrapper(props: SuggestionPanelWrapperProps) {
-  return (
-    <>
-      <h2>Suggestions</h2>
-      {(props.suggestions || []).map((suggestion, index) => {
-        return (
-          <button
-            key={index}
-            data-test-subj="suggestion"
-            onClick={() => {
-              // TODO single action for that?
-              props.dispatch({
-                type: 'SWITCH_VISUALIZATION',
-                newVisulizationId: suggestion.visualizationId,
-                initialState: suggestion.state,
-                datasourceState: suggestion.datasourceState,
-              });
-            }}
-          >
-            {suggestion.title}
-          </button>
-        );
-      })}
-    </>
-  );
-}

From f05024b64142ad0538137305698b1e20f1030f47 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Mon, 20 May 2019 06:27:42 -0400
Subject: [PATCH 25/41] resolve conflicts

---
 .../public/indexpattern_plugin/indexpattern.tsx   | 15 ++-------------
 1 file changed, 2 insertions(+), 13 deletions(-)

diff --git a/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx b/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
index 53522ff39d6fb..29b018320a6db 100644
--- a/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
+++ b/x-pack/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
@@ -36,13 +36,6 @@ export interface Field {
   searchable: boolean;
 }
 
-export interface Field {
-  name: string;
-  type: string;
-  aggregatable: boolean;
-  searchable: boolean;
-}
-
 export interface IndexPatternPersistedState {
   currentIndexPattern: string;
 
@@ -56,16 +49,12 @@ export type IndexPatternPrivateState = IndexPatternPersistedState & {
   indexPatterns: { [id: string]: IndexPattern };
 };
 
-type PersistedKeys = 'currentIndexPattern' | 'columns' | 'columnOrder';
-
-export type PersistableIndexPatternPrivateState = Pick<IndexPatternPrivateState, PersistedKeys>;
-
 // Not stateful. State is persisted to the frame
 export const indexPatternDatasource: Datasource<
   IndexPatternPrivateState,
-  PersistableIndexPatternPrivateState
+  IndexPatternPersistedState
 > = {
-  async initialize(state?: PersistableIndexPatternPrivateState) {
+  async initialize(state?: IndexPatternPersistedState) {
     // TODO: Make fetch request to load indexPatterns from saved objects
     if (state) {
       return {

From bc0c85aa2c65ed108cc87902732283f2c59e8026 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Mon, 20 May 2019 06:47:35 -0400
Subject: [PATCH 26/41] start implementing drop fiel control flow

---
 .../editor_frame/editor_frame.tsx             |   2 +-
 .../editor_frame/suggestion_panel.tsx         |  12 +-
 .../editor_frame/workspace_panel.tsx          | 111 ++++++++++++++++++
 3 files changed, 118 insertions(+), 7 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 835935bb3aa12..3765edc070b05 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -85,7 +85,7 @@ export function EditorFrame(props: EditorFrameProps) {
         suggestionsPanel={
           <SuggestionPanel
             activeDatasource={props.datasources[state.activeDatasource]}
-            activeVisualization={state.activeVisualization}
+            activeVisualizationId={state.activeVisualization}
             datasourcePublicAPI={datasourcePublicAPI!}
             datasourceState={state.datasourceState}
             visualizationState={
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
index 68162eb3ad95f..b6345c17754f6 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
@@ -15,23 +15,23 @@ import {
   TableColumn,
 } from '../../types';
 
-interface SuggestionPanelWrapper {
+interface SuggestionPanelProps {
   activeDatasource: Datasource;
   datasourceState: unknown;
-  activeVisualization: string | null;
+  activeVisualizationId: string | null;
   visualizations: Record<string, Visualization>;
   visualizationState: unknown;
   datasourcePublicAPI: DatasourcePublicAPI;
   dispatch: (action: Action) => void;
 }
 
-export function SuggestionPanel(props: SuggestionPanelWrapper) {
+export function SuggestionPanel(props: SuggestionPanelProps) {
   const currentDatasource = props.activeDatasource;
   const datasourceSuggestions = currentDatasource.getDatasourceSuggestionsFromCurrentState(
     props.datasourceState
   );
-  const roleMapping = props.activeVisualization
-    ? props.visualizations[props.activeVisualization].getMappingOfTableToRoles(
+  const roleMapping = props.activeVisualizationId
+    ? props.visualizations[props.activeVisualizationId].getMappingOfTableToRoles(
         props.visualizationState,
         props.datasourcePublicAPI
       )
@@ -39,6 +39,7 @@ export function SuggestionPanel(props: SuggestionPanelWrapper) {
   const suggestions = getSuggestions(datasourceSuggestions, props.visualizations, roleMapping);
   return (
     <>
+      {/* TODO: I18N */}
       <h2>Suggestions</h2>
       {suggestions.map((suggestion, index) => {
         return (
@@ -46,7 +47,6 @@ export function SuggestionPanel(props: SuggestionPanelWrapper) {
             key={index}
             data-test-subj="suggestion"
             onClick={() => {
-              // TODO single action for that?
               props.dispatch({
                 type: 'SWITCH_VISUALIZATION',
                 newVisulizationId: suggestion.visualizationId,
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
new file mode 100644
index 0000000000000..c57e0252d1952
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
@@ -0,0 +1,111 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React from 'react';
+import { Action } from '../state_management';
+import {
+  Datasource,
+  Visualization,
+  DatasourcePublicAPI,
+  DatasourceSuggestion,
+  DimensionRole,
+  TableColumn,
+} from '../../types';
+import { DragDrop } from '../../drag_drop';
+
+interface WorkspacePanelProps {
+  activeDatasource: Datasource;
+  datasourceState: unknown;
+  activeVisualizationId: string | null;
+  visualizations: Record<string, Visualization>;
+  visualizationState: unknown;
+  datasourcePublicAPI: DatasourcePublicAPI;
+  dispatch: (action: Action) => void;
+}
+
+interface ExpressionRendererProps {
+  expression: string;
+}
+
+function ExpressionRenderer(props: ExpressionRendererProps) {
+  // TODO: actually render the expression and move this to a generic folder as it can be re-used for
+  // suggestion rendering
+  return <span>{props.expression}</span>;
+}
+
+export function WorkspacePanel(props: WorkspacePanelProps) {
+  const currentDatasource = props.activeDatasource;
+  const expression = props.activeVisualizationId
+    ? `${props.activeDatasource.toExpression(props.datasourceState)} | ${props.visualizations[
+        props.activeVisualizationId
+      ].toExpression(props.visualizationState, props.datasourcePublicAPI)}`
+    : null;
+  return (
+    <DragDrop
+      draggable={false}
+      droppable={true}
+      onDrop={() => {
+        const datasourceSuggestions = currentDatasource.getDatasourceSuggestionsForField(
+          props.datasourceState
+        );
+        const roleMapping = props.activeVisualizationId
+          ? props.visualizations[props.activeVisualizationId].getMappingOfTableToRoles(
+              props.visualizationState,
+              props.datasourcePublicAPI
+            )
+          : [];
+        const suggestion = getSuggestions(
+          datasourceSuggestions,
+          props.visualizations,
+          roleMapping
+        )[0];
+        // TODO heuristically present the suggestions in a modal instead of just picking the first one
+        props.dispatch({
+          type: 'SWITCH_VISUALIZATION',
+          newVisulizationId: suggestion.visualizationId,
+          initialState: suggestion.state,
+          datasourceState: suggestion.datasourceState,
+        });
+      }}
+    >
+      {expression !== null ? (
+        <ExpressionRenderer expression={expression} />
+      ) : (
+        <p>{/* TODO: I18N */}This is the workspace panel. Drop fields here</p>
+      )}
+    </DragDrop>
+  );
+}
+
+function getSuggestions(
+  datasourceTableSuggestions: DatasourceSuggestion[],
+  visualizations: Record<string, Visualization>,
+  currentColumnRoles: DimensionRole[]
+) {
+  const datasourceTableMetas: Record<string, TableColumn[]> = {};
+  datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
+    datasourceTableMetas[datasourceSuggestionId] = tableColumns;
+  });
+
+  return (
+    Object.entries(visualizations)
+      .map(([visualizationId, visualization]) => {
+        return visualization
+          .getSuggestions({
+            tableColumns: datasourceTableMetas,
+            roles: currentColumnRoles,
+          })
+          .map(({ datasourceSuggestionId, ...suggestion }) => ({
+            ...suggestion,
+            visualizationId,
+            datasourceState: datasourceTableSuggestions[datasourceSuggestionId].state,
+          }));
+      })
+      // TODO why is flatMap not available here?
+      .reduce((globalList, currentList) => [...globalList, ...currentList], [])
+      .sort(({ score: scoreA }, { score: scoreB }) => scoreB - scoreA)
+  );
+}

From 3d055ee79b20431b8d1b61d960cc5a6ae003af86 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 12:29:19 +0200
Subject: [PATCH 27/41] adress review and restructure state

---
 .../editor_frame/config_panel_wrapper.tsx     |  22 +--
 .../editor_frame/data_panel_wrapper.tsx       |   8 +-
 .../editor_frame/editor_frame.test.tsx        |  73 ++++----
 .../editor_frame/editor_frame.tsx             |  43 ++---
 .../editor_frame/state_management.test.ts     | 156 ++++++++++++++++++
 .../editor_frame/state_management.ts          | 119 +++++++++++++
 .../public/editor_frame_plugin/plugin.tsx     |  55 +++---
 .../state_management.test.ts                  | 143 ----------------
 .../editor_frame_plugin/state_management.ts   |  98 -----------
 9 files changed, 380 insertions(+), 337 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
 delete mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
 delete mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
index e494b8f81d786..1a8c9bfe3b298 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
@@ -7,13 +7,13 @@
 import React, { useMemo } from 'react';
 import { EuiSelect } from '@elastic/eui';
 import { NativeRenderer } from '../../native_renderer';
-import { Action } from '../state_management';
+import { Action } from './state_management';
 import { Visualization, DatasourcePublicAPI } from '../../types';
 
 interface ConfigPanelWrapperProps {
-  visualizationState: Record<string, unknown>;
-  visualizations: Record<string, Visualization>;
-  activeVisualization: string | null;
+  visualizationStateMap: Record<string, unknown>;
+  visualizationMap: Record<string, Visualization>;
+  activeVisualizationId: string | null;
   dispatch: (action: Action) => void;
   datasourcePublicAPI: DatasourcePublicAPI;
 }
@@ -33,26 +33,26 @@ export function ConfigPanelWrapper(props: ConfigPanelWrapperProps) {
     <>
       <EuiSelect
         data-test-subj="visualization-switch"
-        options={Object.keys(props.visualizations).map(visualizationId => ({
+        options={Object.keys(props.visualizationMap).map(visualizationId => ({
           value: visualizationId,
           text: visualizationId,
         }))}
-        value={props.activeVisualization || undefined}
+        value={props.activeVisualizationId || undefined}
         onChange={e => {
           props.dispatch({
             type: 'SWITCH_VISUALIZATION',
-            newVisulizationId: e.target.value,
+            newVisualizationId: e.target.value,
             // TODO we probably want to have a separate API to "force" a visualization switch
             // which isn't a result of a picked suggestion
-            initialState: props.visualizations[e.target.value].initialize(),
+            initialState: props.visualizationMap[e.target.value].initialize(),
           });
         }}
       />
-      {props.activeVisualization && (
+      {props.activeVisualizationId && (
         <NativeRenderer
-          render={props.visualizations[props.activeVisualization].renderConfigPanel}
+          render={props.visualizationMap[props.activeVisualizationId].renderConfigPanel}
           nativeProps={{
-            state: props.visualizationState[props.activeVisualization],
+            state: props.visualizationStateMap[props.activeVisualizationId],
             setState: setVisualizationState,
             datasource: props.datasourcePublicAPI,
           }}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
index 2fdf509c6716e..4ad0fe3285802 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
@@ -8,11 +8,11 @@ import React, { useMemo } from 'react';
 import { EuiSelect } from '@elastic/eui';
 import { DatasourceDataPanelProps, Datasource } from '../..';
 import { NativeRenderer } from '../../native_renderer';
-import { Action } from '../state_management';
+import { Action } from './state_management';
 
 interface DataPanelWrapperProps {
   datasourceState: unknown;
-  datasources: Record<string, Datasource>;
+  datasourceMap: Record<string, Datasource>;
   activeDatasource: string | null;
   datasourceIsLoading: boolean;
   dispatch: (action: Action) => void;
@@ -38,7 +38,7 @@ export function DataPanelWrapper(props: DataPanelWrapperProps) {
     <>
       <EuiSelect
         data-test-subj="datasource-switch"
-        options={Object.keys(props.datasources).map(datasourceId => ({
+        options={Object.keys(props.datasourceMap).map(datasourceId => ({
           value: datasourceId,
           text: datasourceId,
         }))}
@@ -49,7 +49,7 @@ export function DataPanelWrapper(props: DataPanelWrapperProps) {
       />
       {props.activeDatasource && !props.datasourceIsLoading && (
         <NativeRenderer
-          render={props.datasources[props.activeDatasource].renderDataPanel}
+          render={props.datasourceMap[props.activeDatasource].renderDataPanel}
           nativeProps={datasourceProps}
         />
       )}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
index 4b334a541c984..860d5bb1d8008 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
@@ -10,7 +10,9 @@ import { EditorFrame } from './editor_frame';
 import { Visualization, Datasource, DatasourcePublicAPI } from '../../types';
 import { act } from 'react-dom/test-utils';
 
-const nextTick = () => new Promise(resolve => setTimeout(resolve));
+// calling this function will wait for all pending Promises from mock
+// datasources to be processed by its callers.
+const waitForPromises = () => new Promise(resolve => setImmediate(resolve));
 
 describe('editor_frame', () => {
   const getMockVisualization = () => ({
@@ -51,10 +53,10 @@ describe('editor_frame', () => {
       act(() => {
         mount(
           <EditorFrame
-            visualizations={{
+            visualizationMap={{
               testVis: mockVisualization,
             }}
-            datasources={{
+            datasourceMap={{
               testDatasource: mockDatasource,
             }}
             initialDatasource="testDatasource"
@@ -71,10 +73,10 @@ describe('editor_frame', () => {
       act(() => {
         mount(
           <EditorFrame
-            visualizations={{
+            visualizationMap={{
               testVis: mockVisualization,
             }}
-            datasources={{
+            datasourceMap={{
               testDatasource: mockDatasource,
             }}
             initialDatasource={null}
@@ -91,10 +93,10 @@ describe('editor_frame', () => {
       act(() => {
         mount(
           <EditorFrame
-            visualizations={{
+            visualizationMap={{
               testVis: mockVisualization,
             }}
-            datasources={{
+            datasourceMap={{
               testDatasource: mockDatasource,
             }}
             initialDatasource="testDatasource"
@@ -114,10 +116,10 @@ describe('editor_frame', () => {
       act(() => {
         mount(
           <EditorFrame
-            visualizations={{
+            visualizationMap={{
               testVis: mockVisualization,
             }}
-            datasources={{
+            datasourceMap={{
               testDatasource: {
                 ...mockDatasource,
                 initialize: () =>
@@ -132,10 +134,9 @@ describe('editor_frame', () => {
         );
       });
 
-      await nextTick();
       databaseInitialized!(initialState);
 
-      await nextTick();
+      await waitForPromises();
       expect(mockDatasource.renderDataPanel).toHaveBeenCalledWith(
         expect.any(Element),
         expect.objectContaining({ state: initialState })
@@ -147,10 +148,10 @@ describe('editor_frame', () => {
 
       mount(
         <EditorFrame
-          visualizations={{
+          visualizationMap={{
             testVis: { ...mockVisualization, initialize: () => initialState },
           }}
-          datasources={{
+          datasourceMap={{
             testDatasource: {
               ...mockDatasource,
               initialize: () => Promise.resolve(),
@@ -161,7 +162,7 @@ describe('editor_frame', () => {
         />
       );
 
-      await nextTick();
+      await waitForPromises();
 
       expect(mockVisualization.renderConfigPanel).toHaveBeenCalledWith(
         expect.any(Element),
@@ -174,10 +175,10 @@ describe('editor_frame', () => {
     it('should re-render config panel after state update', async () => {
       mount(
         <EditorFrame
-          visualizations={{
+          visualizationMap={{
             testVis: mockVisualization,
           }}
-          datasources={{
+          datasourceMap={{
             testDatasource: mockDatasource,
           }}
           initialDatasource="testDatasource"
@@ -185,7 +186,7 @@ describe('editor_frame', () => {
         />
       );
 
-      await nextTick();
+      await waitForPromises();
 
       const updatedState = {};
       const setVisualizationState = (mockVisualization.renderConfigPanel as jest.Mock).mock
@@ -209,10 +210,10 @@ describe('editor_frame', () => {
     it('should re-render data panel after state update', async () => {
       mount(
         <EditorFrame
-          visualizations={{
+          visualizationMap={{
             testVis: mockVisualization,
           }}
-          datasources={{
+          datasourceMap={{
             testDatasource: mockDatasource,
           }}
           initialDatasource="testDatasource"
@@ -220,7 +221,7 @@ describe('editor_frame', () => {
         />
       );
 
-      await nextTick();
+      await waitForPromises();
 
       const updatedState = {};
       const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock.calls[0][1]
@@ -241,10 +242,10 @@ describe('editor_frame', () => {
     it('should re-render config panel with updated datasource api after datasource state update', async () => {
       mount(
         <EditorFrame
-          visualizations={{
+          visualizationMap={{
             testVis: mockVisualization,
           }}
-          datasources={{
+          datasourceMap={{
             testDatasource: mockDatasource,
           }}
           initialDatasource="testDatasource"
@@ -252,7 +253,7 @@ describe('editor_frame', () => {
         />
       );
 
-      await nextTick();
+      await waitForPromises();
 
       const updatedPublicAPI = {};
       mockDatasource.getPublicAPI = jest.fn(
@@ -283,10 +284,10 @@ describe('editor_frame', () => {
 
       mount(
         <EditorFrame
-          visualizations={{
+          visualizationMap={{
             testVis: mockVisualization,
           }}
-          datasources={{
+          datasourceMap={{
             testDatasource: mockDatasource,
           }}
           initialDatasource="testDatasource"
@@ -294,7 +295,7 @@ describe('editor_frame', () => {
         />
       );
 
-      await nextTick();
+      await waitForPromises();
 
       expect(mockVisualization.renderConfigPanel).toHaveBeenCalledWith(
         expect.any(Element),
@@ -308,10 +309,10 @@ describe('editor_frame', () => {
 
       mount(
         <EditorFrame
-          visualizations={{
+          visualizationMap={{
             testVis: mockVisualization,
           }}
-          datasources={{
+          datasourceMap={{
             testDatasource: mockDatasource,
           }}
           initialDatasource="testDatasource"
@@ -319,7 +320,7 @@ describe('editor_frame', () => {
         />
       );
 
-      await nextTick();
+      await waitForPromises();
 
       expect(mockDatasource.getPublicAPI).toHaveBeenCalledWith(
         datasourceState,
@@ -330,10 +331,10 @@ describe('editor_frame', () => {
     it('should re-create the public api after state has been set', async () => {
       mount(
         <EditorFrame
-          visualizations={{
+          visualizationMap={{
             testVis: mockVisualization,
           }}
-          datasources={{
+          datasourceMap={{
             testDatasource: mockDatasource,
           }}
           initialDatasource="testDatasource"
@@ -341,7 +342,7 @@ describe('editor_frame', () => {
         />
       );
 
-      await nextTick();
+      await waitForPromises();
 
       const updatedState = {};
       const setDatasourceState = (mockDatasource.getPublicAPI as jest.Mock).mock.calls[0][1];
@@ -362,11 +363,11 @@ describe('editor_frame', () => {
     beforeEach(async () => {
       instance = mount(
         <EditorFrame
-          visualizations={{
+          visualizationMap={{
             testVis: mockVisualization,
             testVis2: mockVisualization2,
           }}
-          datasources={{
+          datasourceMap={{
             testDatasource: mockDatasource,
             testDatasource2: mockDatasource2,
           }}
@@ -374,7 +375,7 @@ describe('editor_frame', () => {
           initialVisualization="testVis"
         />
       );
-      await nextTick();
+      await waitForPromises();
 
       // necessary to flush elements to dom synchronously
       instance.update();
@@ -405,7 +406,7 @@ describe('editor_frame', () => {
         .find('select[data-test-subj="datasource-switch"]')
         .simulate('change', { target: { value: 'testDatasource2' } });
 
-      await nextTick();
+      await waitForPromises();
 
       expect(mockDatasource2.renderDataPanel).toHaveBeenCalledWith(
         expect.any(Element),
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 09eb5a7fcb0c7..c6e70ea22425f 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -6,14 +6,14 @@
 
 import React, { useEffect, useReducer, useMemo } from 'react';
 import { Datasource, Visualization } from '../../types';
-import { reducer, getInitialState } from '../state_management';
+import { reducer, getInitialState } from './state_management';
 import { DataPanelWrapper } from './data_panel_wrapper';
 import { ConfigPanelWrapper } from './config_panel_wrapper';
 import { FrameLayout } from './frame_layout';
 
 export interface EditorFrameProps {
-  datasources: Record<string, Datasource>;
-  visualizations: Record<string, Visualization>;
+  datasourceMap: Record<string, Datasource>;
+  visualizationMap: Record<string, Visualization>;
 
   initialDatasource: string | null;
   initialVisualization: string | null;
@@ -26,8 +26,8 @@ export function EditorFrame(props: EditorFrameProps) {
   useEffect(
     () => {
       let datasourceGotSwitched = false;
-      if (state.datasourceIsLoading && state.activeDatasource) {
-        props.datasources[state.activeDatasource].initialize().then(datasourceState => {
+      if (state.datasource.isLoading && state.datasource.activeId) {
+        props.datasourceMap[state.datasource.activeId].initialize().then(datasourceState => {
           if (!datasourceGotSwitched) {
             dispatch({
               type: 'UPDATE_DATASOURCE_STATE',
@@ -41,14 +41,14 @@ export function EditorFrame(props: EditorFrameProps) {
         };
       }
     },
-    [state.activeDatasource, state.datasourceIsLoading]
+    [state.datasource.activeId, state.datasource.isLoading]
   );
 
   const datasourcePublicAPI = useMemo(
     () =>
-      state.activeDatasource && !state.datasourceIsLoading
-        ? props.datasources[state.activeDatasource].getPublicAPI(
-            state.datasourceState,
+      state.datasource.activeId && !state.datasource.isLoading
+        ? props.datasourceMap[state.datasource.activeId].getPublicAPI(
+            state.datasource.state,
             (newState: unknown) => {
               dispatch({
                 type: 'UPDATE_DATASOURCE_STATE',
@@ -57,29 +57,34 @@ export function EditorFrame(props: EditorFrameProps) {
             }
           )
         : undefined,
-    [props.datasources, state.datasourceIsLoading, state.activeDatasource, state.datasourceState]
+    [
+      props.datasourceMap,
+      state.datasource.isLoading,
+      state.datasource.activeId,
+      state.datasource.state,
+    ]
   );
 
   return (
     <FrameLayout
       dataPanel={
         <DataPanelWrapper
-          activeDatasource={state.activeDatasource}
-          datasourceIsLoading={state.datasourceIsLoading}
-          datasourceState={state.datasourceState}
-          datasources={props.datasources}
+          activeDatasource={state.datasource.activeId}
+          datasourceIsLoading={state.datasource.isLoading}
+          datasourceState={state.datasource.state}
+          datasourceMap={props.datasourceMap}
           dispatch={dispatch}
         />
       }
       configPanel={
-        state.activeDatasource &&
-        !state.datasourceIsLoading && (
+        state.datasource.activeId &&
+        !state.datasource.isLoading && (
           <ConfigPanelWrapper
-            visualizations={props.visualizations}
-            activeVisualization={state.activeVisualization}
+            visualizationMap={props.visualizationMap}
+            activeVisualizationId={state.visualization.activeId}
             datasourcePublicAPI={datasourcePublicAPI!}
             dispatch={dispatch}
-            visualizationState={state.visualizationState}
+            visualizationStateMap={state.visualization.stateMap}
           />
         )
       }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
new file mode 100644
index 0000000000000..c610a74368a01
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
@@ -0,0 +1,156 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { getInitialState, reducer } from './state_management';
+import { EditorFrameProps } from '.';
+import { Datasource, Visualization } from '../../types';
+
+describe('editor_frame state management', () => {
+  describe('initialization', () => {
+    let props: EditorFrameProps;
+
+    beforeEach(() => {
+      props = {
+        datasourceMap: { testDatasource: ({} as unknown) as Datasource },
+        visualizationMap: { testVis: ({ initialize: jest.fn() } as unknown) as Visualization },
+        initialDatasource: 'testDatasource',
+        initialVisualization: 'testVis',
+      };
+    });
+
+    it('should store initial datasource and visualization', () => {
+      const initialState = getInitialState(props);
+      expect(initialState.datasource.activeId).toEqual('testDatasource');
+      expect(initialState.visualization.activeId).toEqual('testVis');
+    });
+
+    it('should initialize visualization', () => {
+      const initialVisState = {};
+      props.visualizationMap.testVis.initialize = jest.fn(() => initialVisState);
+
+      const initialState = getInitialState(props);
+
+      expect(initialState.visualization.stateMap.testVis).toBe(initialVisState);
+      expect(props.visualizationMap.testVis.initialize).toHaveBeenCalled();
+    });
+
+    it('should not initialize visualization if no initial visualization is passed in', () => {
+      const initialState = getInitialState({ ...props, initialVisualization: null });
+
+      expect(initialState.visualization.stateMap).toEqual({});
+      expect(props.visualizationMap.testVis.initialize).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('state update', () => {
+    it('should update the corresponding visualization state on update', () => {
+      const newVisState = {};
+      const newState = reducer(
+        {
+          datasource: {
+            activeId: 'testDatasource',
+            state: {},
+            isLoading: false,
+          },
+          visualization: {
+            activeId: 'testVis',
+            stateMap: {
+              testVis: {},
+            },
+          },
+        },
+        {
+          type: 'UPDATE_VISUALIZATION_STATE',
+          newState: newVisState,
+        }
+      );
+
+      expect(newState.visualization.stateMap).toEqual({
+        testVis: newVisState,
+      });
+    });
+
+    it('should update the datasource state on update', () => {
+      const newDatasourceState = {};
+      const newState = reducer(
+        {
+          datasource: {
+            activeId: 'testDatasource',
+            state: {},
+            isLoading: false,
+          },
+          visualization: {
+            activeId: 'testVis',
+            stateMap: {
+              testVis: {},
+            },
+          },
+        },
+        {
+          type: 'UPDATE_DATASOURCE_STATE',
+          newState: newDatasourceState,
+        }
+      );
+
+      expect(newState.datasource.state).toBe(newDatasourceState);
+    });
+
+    it('should should switch active visualization but dont loose old state', () => {
+      const testVisState = {};
+      const newVisState = {};
+      const newState = reducer(
+        {
+          datasource: {
+            activeId: 'testDatasource',
+            state: {},
+            isLoading: false,
+          },
+          visualization: {
+            activeId: 'testVis',
+            stateMap: {
+              testVis: testVisState,
+            },
+          },
+        },
+        {
+          type: 'SWITCH_VISUALIZATION',
+          newVisualizationId: 'testVis2',
+          initialState: newVisState,
+        }
+      );
+
+      expect(newState.visualization.stateMap.testVis).toBe(testVisState);
+      expect(newState.visualization.stateMap.testVis2).toBe(newVisState);
+    });
+
+    it('should should switch active datasource and purge visualization state', () => {
+      const newState = reducer(
+        {
+          datasource: {
+            activeId: 'testDatasource',
+            state: {},
+            isLoading: false,
+          },
+          visualization: {
+            activeId: 'testVis',
+            stateMap: {
+              testVis: {},
+            },
+          },
+        },
+        {
+          type: 'SWITCH_DATASOURCE',
+          newDatasourceId: 'testDatasource2',
+        }
+      );
+
+      expect(newState.visualization.stateMap).toEqual({});
+      expect(newState.visualization.activeId).toBe(null);
+      expect(newState.datasource.activeId).toBe('testDatasource2');
+      expect(newState.datasource.state).toBe(null);
+    });
+  });
+});
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
new file mode 100644
index 0000000000000..444551374389f
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
@@ -0,0 +1,119 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { EditorFrameProps } from '.';
+
+export interface EditorFrameState {
+  visualization: {
+    activeId: string | null;
+    stateMap: {
+      [visualizationId: string]: unknown;
+    };
+  };
+  datasource: {
+    activeId: string | null;
+    state: unknown;
+    isLoading: boolean;
+  };
+}
+
+export type Action =
+  | {
+      type: 'UPDATE_DATASOURCE_STATE';
+      newState: unknown;
+    }
+  | {
+      type: 'UPDATE_VISUALIZATION_STATE';
+      newState: unknown;
+    }
+  | {
+      type: 'SWITCH_VISUALIZATION';
+      newVisualizationId: string;
+      initialState: unknown;
+    }
+  | {
+      type: 'SWITCH_DATASOURCE';
+      newDatasourceId: string;
+    };
+
+export const getInitialState = (props: EditorFrameProps): EditorFrameState => {
+  return {
+    datasource: {
+      state: null,
+      isLoading: Boolean(props.initialDatasource),
+      activeId: props.initialDatasource,
+    },
+    visualization: {
+      stateMap: props.initialVisualization
+        ? {
+            [props.initialVisualization]: props.visualizationMap[
+              props.initialVisualization
+            ].initialize(),
+          }
+        : {},
+      activeId: props.initialVisualization,
+    },
+  };
+};
+
+export const reducer = (state: EditorFrameState, action: Action): EditorFrameState => {
+  switch (action.type) {
+    case 'SWITCH_DATASOURCE':
+      return {
+        ...state,
+        datasource: {
+          ...state.datasource,
+          isLoading: true,
+          state: null,
+          activeId: action.newDatasourceId,
+        },
+        visualization: {
+          ...state.visualization,
+          // purge all visualizations on datasource switch
+          stateMap: {},
+          activeId: null,
+        },
+      };
+    case 'SWITCH_VISUALIZATION':
+      return {
+        ...state,
+        visualization: {
+          ...state.visualization,
+          activeId: action.newVisualizationId,
+          stateMap: {
+            ...state.visualization.stateMap,
+            [action.newVisualizationId]: action.initialState,
+          },
+        },
+      };
+    case 'UPDATE_DATASOURCE_STATE':
+      return {
+        ...state,
+        datasource: {
+          ...state.datasource,
+          // when the datasource state is updated, the initialization is complete
+          isLoading: false,
+          state: action.newState,
+        },
+      };
+    case 'UPDATE_VISUALIZATION_STATE':
+      if (!state.visualization.activeId) {
+        throw new Error('Invariant: visualization state got updated without active visualization');
+      }
+      return {
+        ...state,
+        visualization: {
+          ...state.visualization,
+          stateMap: {
+            ...state.visualization.stateMap,
+            [state.visualization.activeId]: action.newState,
+          },
+        },
+      };
+    default:
+      return state;
+  }
+};
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
index 8787eadd7dccc..d358e46dbe0f9 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
@@ -6,7 +6,7 @@
 
 import React from 'react';
 import { render, unmountComponentAtNode } from 'react-dom';
-import { Datasource, Visualization, EditorFrameSetup } from '../types';
+import { Datasource, Visualization, EditorFrameSetup, EditorFrameInstance } from '../types';
 
 import { EditorFrame } from './editor_frame';
 
@@ -16,33 +16,36 @@ export class EditorFramePlugin {
   private datasources: Record<string, Datasource> = {};
   private visualizations: Record<string, Visualization> = {};
 
-  public setup(): EditorFrameSetup {
+  private createInstance(): EditorFrameInstance {
+    let domElement: Element;
+
+    function unmount() {
+      if (domElement) {
+        unmountComponentAtNode(domElement);
+      }
+    }
+
     return {
-      createInstance: () => {
-        let domElement: Element;
-        return {
-          mount: (element: Element) => {
-            if (domElement) {
-              unmountComponentAtNode(domElement);
-            }
-            domElement = element;
-            render(
-              <EditorFrame
-                datasources={this.datasources}
-                visualizations={this.visualizations}
-                initialDatasource={Object.keys(this.datasources)[0]}
-                initialVisualization={Object.keys(this.visualizations)[0]}
-              />,
-              domElement
-            );
-          },
-          unmount: () => {
-            if (domElement) {
-              unmountComponentAtNode(domElement);
-            }
-          },
-        };
+      mount: element => {
+        unmount();
+        domElement = element;
+        render(
+          <EditorFrame
+            datasourceMap={this.datasources}
+            visualizationMap={this.visualizations}
+            initialDatasource={Object.keys(this.datasources)[0]}
+            initialVisualization={Object.keys(this.visualizations)[0]}
+          />,
+          domElement
+        );
       },
+      unmount,
+    };
+  }
+
+  public setup(): EditorFrameSetup {
+    return {
+      createInstance: this.createInstance.bind(this),
       registerDatasource: (name, datasource) => {
         this.datasources[name] = datasource as Datasource<unknown, unknown>;
       },
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
deleted file mode 100644
index fb96817685c45..0000000000000
--- a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.test.ts
+++ /dev/null
@@ -1,143 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License;
- * you may not use this file except in compliance with the Elastic License.
- */
-
-import { getInitialState, reducer } from './state_management';
-import { EditorFrameProps } from './editor_frame';
-import { Datasource, Visualization } from '../types';
-
-describe('editor_frame state management', () => {
-  describe('initialization', () => {
-    let props: EditorFrameProps;
-
-    beforeEach(() => {
-      props = {
-        datasources: { testDatasource: ({} as unknown) as Datasource },
-        visualizations: { testVis: ({ initialize: jest.fn() } as unknown) as Visualization },
-        initialDatasource: 'testDatasource',
-        initialVisualization: 'testVis',
-      };
-    });
-
-    it('should store initial datasource and visualization', () => {
-      expect(getInitialState(props)).toEqual(
-        expect.objectContaining({
-          activeDatasource: 'testDatasource',
-          activeVisualization: 'testVis',
-        })
-      );
-    });
-
-    it('should initialize visualization', () => {
-      const initialVisState = {};
-      props.visualizations.testVis.initialize = jest.fn(() => initialVisState);
-
-      const initialState = getInitialState(props);
-
-      expect(initialState.visualizationState.testVis).toBe(initialVisState);
-      expect(props.visualizations.testVis.initialize).toHaveBeenCalled();
-    });
-
-    it('should not initialize visualization if no initial visualization is passed in', () => {
-      const initialState = getInitialState({ ...props, initialVisualization: null });
-
-      expect(initialState.visualizationState).toEqual({});
-      expect(props.visualizations.testVis.initialize).not.toHaveBeenCalled();
-    });
-  });
-
-  describe('state update', () => {
-    it('should update the corresponding visualization state on update', () => {
-      const newVisState = {};
-      const newState = reducer(
-        {
-          activeDatasource: 'testDatasource',
-          activeVisualization: 'testVis',
-          datasourceIsLoading: false,
-          datasourceState: {},
-          visualizationState: {
-            testVis: {},
-          },
-        },
-        {
-          type: 'UPDATE_VISUALIZATION_STATE',
-          newState: newVisState,
-        }
-      );
-
-      expect(newState.visualizationState).toEqual({
-        testVis: newVisState,
-      });
-    });
-
-    it('should update the datasource state on update', () => {
-      const newDatasourceState = {};
-      const newState = reducer(
-        {
-          activeDatasource: 'testDatasource',
-          activeVisualization: 'testVis',
-          datasourceIsLoading: false,
-          datasourceState: {},
-          visualizationState: {
-            testVis: {},
-          },
-        },
-        {
-          type: 'UPDATE_DATASOURCE_STATE',
-          newState: newDatasourceState,
-        }
-      );
-
-      expect(newState.datasourceState).toBe(newDatasourceState);
-    });
-
-    it('should should switch active visualization but dont loose old state', () => {
-      const testVisState = {};
-      const newVisState = {};
-      const newState = reducer(
-        {
-          activeDatasource: 'testDatasource',
-          activeVisualization: 'testVis',
-          datasourceIsLoading: false,
-          datasourceState: {},
-          visualizationState: {
-            testVis: testVisState,
-          },
-        },
-        {
-          type: 'SWITCH_VISUALIZATION',
-          newVisulizationId: 'testVis2',
-          initialState: newVisState,
-        }
-      );
-
-      expect(newState.visualizationState.testVis).toBe(testVisState);
-      expect(newState.visualizationState.testVis2).toBe(newVisState);
-    });
-
-    it('should should switch active datasource and purge visualization state', () => {
-      const newState = reducer(
-        {
-          activeDatasource: 'testDatasource',
-          activeVisualization: 'testVis',
-          datasourceIsLoading: false,
-          datasourceState: {},
-          visualizationState: {
-            testVis: {},
-          },
-        },
-        {
-          type: 'SWITCH_DATASOURCE',
-          newDatasourceId: 'testDatasource2',
-        }
-      );
-
-      expect(newState.visualizationState).toEqual({});
-      expect(newState.activeVisualization).toBe(null);
-      expect(newState.activeDatasource).toBe('testDatasource2');
-      expect(newState.datasourceState).toBe(null);
-    });
-  });
-});
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts b/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts
deleted file mode 100644
index e35e258d1153c..0000000000000
--- a/x-pack/plugins/lens/public/editor_frame_plugin/state_management.ts
+++ /dev/null
@@ -1,98 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License;
- * you may not use this file except in compliance with the Elastic License.
- */
-
-import { EditorFrameProps } from './editor_frame';
-
-export interface EditorFrameState {
-  activeDatasource: string | null;
-  datasourceState: unknown;
-  datasourceIsLoading: boolean;
-
-  activeVisualization: string | null;
-  visualizationState: {
-    [visualizationId: string]: unknown;
-  };
-}
-
-export type Action =
-  | {
-      type: 'UPDATE_DATASOURCE_STATE';
-      newState: unknown;
-    }
-  | {
-      type: 'UPDATE_VISUALIZATION_STATE';
-      newState: unknown;
-    }
-  | {
-      type: 'SWITCH_VISUALIZATION';
-      newVisulizationId: string;
-      initialState: unknown;
-    }
-  | {
-      type: 'SWITCH_DATASOURCE';
-      newDatasourceId: string;
-    };
-
-export const getInitialState = (props: EditorFrameProps) => {
-  return {
-    datasourceState: null,
-    datasourceIsLoading: Boolean(props.initialDatasource),
-    activeDatasource: props.initialDatasource,
-    visualizationState: props.initialVisualization
-      ? {
-          [props.initialVisualization]: props.visualizations[
-            props.initialVisualization
-          ].initialize(),
-        }
-      : {},
-    activeVisualization: props.initialVisualization,
-  };
-};
-
-export const reducer = (state: EditorFrameState, action: Action): EditorFrameState => {
-  switch (action.type) {
-    case 'SWITCH_DATASOURCE':
-      return {
-        ...state,
-        activeDatasource: action.newDatasourceId,
-        // purge all visualizations on datasource switch
-        visualizationState: {},
-        activeVisualization: null,
-        datasourceState: null,
-        datasourceIsLoading: true,
-      };
-    case 'SWITCH_VISUALIZATION':
-      return {
-        ...state,
-        activeVisualization: action.newVisulizationId,
-        visualizationState: {
-          ...state.visualizationState,
-          [action.newVisulizationId]: action.initialState,
-        },
-      };
-    case 'UPDATE_DATASOURCE_STATE':
-      return {
-        ...state,
-        // when the datasource state is updated, the initialization is complete
-        datasourceIsLoading: false,
-        datasourceState: action.newState,
-      };
-    case 'UPDATE_VISUALIZATION_STATE':
-      if (state.activeVisualization) {
-        return {
-          ...state,
-          visualizationState: {
-            ...state.visualizationState,
-            [state.activeVisualization]: action.newState,
-          },
-        };
-      } else {
-        throw new Error('Invariant: visualization state got updated without active visualization');
-      }
-    default:
-      return state;
-  }
-};

From 7c3a2c849dfb11fc6a40858935d76e81a22893f2 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 12:34:52 +0200
Subject: [PATCH 28/41] align naming

---
 .../editor_frame/editor_frame.test.tsx        | 48 +++++++++----------
 .../editor_frame/editor_frame.tsx             |  4 +-
 .../editor_frame/state_management.test.ts     |  6 +--
 .../editor_frame/state_management.ts          | 12 ++---
 .../public/editor_frame_plugin/plugin.tsx     |  4 +-
 5 files changed, 37 insertions(+), 37 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
index 860d5bb1d8008..32785d2d95753 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
@@ -59,8 +59,8 @@ describe('editor_frame', () => {
             datasourceMap={{
               testDatasource: mockDatasource,
             }}
-            initialDatasource="testDatasource"
-            initialVisualization="testVis"
+            initialDatasourceId="testDatasource"
+            initialVisualizationId="testVis"
           />
         );
       });
@@ -79,8 +79,8 @@ describe('editor_frame', () => {
             datasourceMap={{
               testDatasource: mockDatasource,
             }}
-            initialDatasource={null}
-            initialVisualization={null}
+            initialDatasourceId={null}
+            initialVisualizationId={null}
           />
         );
       });
@@ -99,8 +99,8 @@ describe('editor_frame', () => {
             datasourceMap={{
               testDatasource: mockDatasource,
             }}
-            initialDatasource="testDatasource"
-            initialVisualization="testVis"
+            initialDatasourceId="testDatasource"
+            initialVisualizationId="testVis"
           />
         );
       });
@@ -128,8 +128,8 @@ describe('editor_frame', () => {
                   }),
               },
             }}
-            initialDatasource="testDatasource"
-            initialVisualization="testVis"
+            initialDatasourceId="testDatasource"
+            initialVisualizationId="testVis"
           />
         );
       });
@@ -157,8 +157,8 @@ describe('editor_frame', () => {
               initialize: () => Promise.resolve(),
             },
           }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
         />
       );
 
@@ -181,8 +181,8 @@ describe('editor_frame', () => {
           datasourceMap={{
             testDatasource: mockDatasource,
           }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
         />
       );
 
@@ -216,8 +216,8 @@ describe('editor_frame', () => {
           datasourceMap={{
             testDatasource: mockDatasource,
           }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
         />
       );
 
@@ -248,8 +248,8 @@ describe('editor_frame', () => {
           datasourceMap={{
             testDatasource: mockDatasource,
           }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
         />
       );
 
@@ -290,8 +290,8 @@ describe('editor_frame', () => {
           datasourceMap={{
             testDatasource: mockDatasource,
           }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
         />
       );
 
@@ -315,8 +315,8 @@ describe('editor_frame', () => {
           datasourceMap={{
             testDatasource: mockDatasource,
           }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
         />
       );
 
@@ -337,8 +337,8 @@ describe('editor_frame', () => {
           datasourceMap={{
             testDatasource: mockDatasource,
           }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
         />
       );
 
@@ -371,8 +371,8 @@ describe('editor_frame', () => {
             testDatasource: mockDatasource,
             testDatasource2: mockDatasource2,
           }}
-          initialDatasource="testDatasource"
-          initialVisualization="testVis"
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
         />
       );
       await waitForPromises();
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index c6e70ea22425f..069320ba55d84 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -15,8 +15,8 @@ export interface EditorFrameProps {
   datasourceMap: Record<string, Datasource>;
   visualizationMap: Record<string, Visualization>;
 
-  initialDatasource: string | null;
-  initialVisualization: string | null;
+  initialDatasourceId: string | null;
+  initialVisualizationId: string | null;
 }
 
 export function EditorFrame(props: EditorFrameProps) {
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
index c610a74368a01..ea40f4e203eb0 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
@@ -16,8 +16,8 @@ describe('editor_frame state management', () => {
       props = {
         datasourceMap: { testDatasource: ({} as unknown) as Datasource },
         visualizationMap: { testVis: ({ initialize: jest.fn() } as unknown) as Visualization },
-        initialDatasource: 'testDatasource',
-        initialVisualization: 'testVis',
+        initialDatasourceId: 'testDatasource',
+        initialVisualizationId: 'testVis',
       };
     });
 
@@ -38,7 +38,7 @@ describe('editor_frame state management', () => {
     });
 
     it('should not initialize visualization if no initial visualization is passed in', () => {
-      const initialState = getInitialState({ ...props, initialVisualization: null });
+      const initialState = getInitialState({ ...props, initialVisualizationId: null });
 
       expect(initialState.visualization.stateMap).toEqual({});
       expect(props.visualizationMap.testVis.initialize).not.toHaveBeenCalled();
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
index 444551374389f..e8b8dac09d5d3 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
@@ -43,18 +43,18 @@ export const getInitialState = (props: EditorFrameProps): EditorFrameState => {
   return {
     datasource: {
       state: null,
-      isLoading: Boolean(props.initialDatasource),
-      activeId: props.initialDatasource,
+      isLoading: Boolean(props.initialDatasourceId),
+      activeId: props.initialDatasourceId,
     },
     visualization: {
-      stateMap: props.initialVisualization
+      stateMap: props.initialVisualizationId
         ? {
-            [props.initialVisualization]: props.visualizationMap[
-              props.initialVisualization
+            [props.initialVisualizationId]: props.visualizationMap[
+              props.initialVisualizationId
             ].initialize(),
           }
         : {},
-      activeId: props.initialVisualization,
+      activeId: props.initialVisualizationId,
     },
   };
 };
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
index d358e46dbe0f9..07c1841601140 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
@@ -33,8 +33,8 @@ export class EditorFramePlugin {
           <EditorFrame
             datasourceMap={this.datasources}
             visualizationMap={this.visualizations}
-            initialDatasource={Object.keys(this.datasources)[0]}
-            initialVisualization={Object.keys(this.visualizations)[0]}
+            initialDatasourceId={Object.keys(this.datasources)[0]}
+            initialVisualizationId={Object.keys(this.visualizations)[0]}
           />,
           domElement
         );

From 849f3a11dc50acda1c44912160b3c65db8cf13be Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 12:37:53 +0200
Subject: [PATCH 29/41] add comment

---
 .../public/editor_frame_plugin/editor_frame/editor_frame.tsx    | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 069320ba55d84..159762939a58b 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -44,6 +44,8 @@ export function EditorFrame(props: EditorFrameProps) {
     [state.datasource.activeId, state.datasource.isLoading]
   );
 
+  // create public datasource api for current state
+  // as soon as datasource is available and memoize it
   const datasourcePublicAPI = useMemo(
     () =>
       state.datasource.activeId && !state.datasource.isLoading

From 754091670dd14f11945c7efb9a38ba30cfceb554 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 12:54:32 +0200
Subject: [PATCH 30/41] fix merge conflict

---
 .../editor_frame_plugin/editor_frame/suggestion_panel.tsx     | 4 ++--
 .../editor_frame_plugin/editor_frame/workspace_panel.tsx      | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
index b6345c17754f6..3ad3d7fccc00f 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
@@ -5,7 +5,7 @@
  */
 
 import React from 'react';
-import { Action } from '../state_management';
+import { Action } from './state_management';
 import {
   Datasource,
   Visualization,
@@ -49,7 +49,7 @@ export function SuggestionPanel(props: SuggestionPanelProps) {
             onClick={() => {
               props.dispatch({
                 type: 'SWITCH_VISUALIZATION',
-                newVisulizationId: suggestion.visualizationId,
+                newVisualizationId: suggestion.visualizationId,
                 initialState: suggestion.state,
                 datasourceState: suggestion.datasourceState,
               });
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
index c57e0252d1952..194c282214781 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
@@ -5,7 +5,7 @@
  */
 
 import React from 'react';
-import { Action } from '../state_management';
+import { Action } from './state_management';
 import {
   Datasource,
   Visualization,
@@ -65,7 +65,7 @@ export function WorkspacePanel(props: WorkspacePanelProps) {
         // TODO heuristically present the suggestions in a modal instead of just picking the first one
         props.dispatch({
           type: 'SWITCH_VISUALIZATION',
-          newVisulizationId: suggestion.visualizationId,
+          newVisualizationId: suggestion.visualizationId,
           initialState: suggestion.state,
           datasourceState: suggestion.datasourceState,
         });

From 7b8c21a333d948ad7fbfec72138e11b6e2acc95d Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 15:31:19 +0200
Subject: [PATCH 31/41] refactor and add tests

---
 .../editor_frame/editor_frame.test.tsx        | 118 ++++++++++++++-
 .../editor_frame/editor_frame.tsx             |  21 ++-
 .../editor_frame/frame_layout.tsx             |   4 +-
 .../editor_frame/suggestion_helpers.ts        |  45 ++++++
 .../editor_frame/suggestion_panel.tsx         |  77 ++++------
 .../editor_frame/workspace_panel.tsx          | 138 ++++++++----------
 6 files changed, 267 insertions(+), 136 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
index 22462d25f30db..c92f29eb2684c 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
@@ -13,7 +13,7 @@ import { createMockVisualization, createMockDatasource } from '../mock_extension
 
 // calling this function will wait for all pending Promises from mock
 // datasources to be processed by its callers.
-const waitForPromises = () => new Promise(resolve => setImmediate(resolve));
+const waitForPromises = () => new Promise(resolve => setTimeout(resolve));
 
 describe('editor_frame', () => {
   let mockVisualization: Visualization;
@@ -584,5 +584,121 @@ describe('editor_frame', () => {
         })
       );
     });
+
+    it('should switch to best suggested visualization on field drop', async () => {
+      const suggestionVisState = {};
+      const instance = mount(
+        <EditorFrame
+          visualizationMap={{
+            testVis: {
+              ...mockVisualization,
+              getSuggestions: () => [
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.2,
+                  state: {},
+                  title: 'Suggestion1',
+                },
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.8,
+                  state: suggestionVisState,
+                  title: 'Suggestion2',
+                },
+              ],
+            },
+            testVis2: mockVisualization2,
+          }}
+          datasourceMap={{
+            testDatasource: {
+              ...mockDatasource,
+              getDatasourceSuggestionsForField: () => [{ state: {}, tableColumns: [] }],
+              getDatasourceSuggestionsFromCurrentState: () => [{ state: {}, tableColumns: [] }],
+            },
+          }}
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
+        />
+      );
+
+      await waitForPromises();
+
+      // TODO why is this necessary?
+      instance.update();
+
+      act(() => {
+        instance.find('[data-test-subj="lnsDragDrop"]').simulate('drop');
+      });
+
+      expect(mockVisualization.renderConfigPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({
+          state: suggestionVisState,
+        })
+      );
+    });
+
+    it('should switch to best suggested visualization regardless extension on field drop', async () => {
+      const suggestionVisState = {};
+      const instance = mount(
+        <EditorFrame
+          visualizationMap={{
+            testVis: {
+              ...mockVisualization,
+              getSuggestions: () => [
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.2,
+                  state: {},
+                  title: 'Suggestion1',
+                },
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.6,
+                  state: {},
+                  title: 'Suggestion2',
+                },
+              ],
+            },
+            testVis2: {
+              ...mockVisualization2,
+              getSuggestions: () => [
+                {
+                  datasourceSuggestionId: 0,
+                  score: 0.8,
+                  state: suggestionVisState,
+                  title: 'Suggestion3',
+                },
+              ],
+            },
+          }}
+          datasourceMap={{
+            testDatasource: {
+              ...mockDatasource,
+              getDatasourceSuggestionsForField: () => [{ state: {}, tableColumns: [] }],
+              getDatasourceSuggestionsFromCurrentState: () => [{ state: {}, tableColumns: [] }],
+            },
+          }}
+          initialDatasourceId="testDatasource"
+          initialVisualizationId="testVis"
+        />
+      );
+
+      await waitForPromises();
+
+      // TODO why is this necessary?
+      instance.update();
+
+      act(() => {
+        instance.find('[data-test-subj="lnsDragDrop"]').simulate('drop');
+      });
+
+      expect(mockVisualization2.renderConfigPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({
+          state: suggestionVisState,
+        })
+      );
+    });
   });
 });
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 80907a8e84a1d..4595ace7477c5 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -11,6 +11,7 @@ import { DataPanelWrapper } from './data_panel_wrapper';
 import { ConfigPanelWrapper } from './config_panel_wrapper';
 import { FrameLayout } from './frame_layout';
 import { SuggestionPanel } from './suggestion_panel';
+import { WorkspacePanel } from './workspace_panel';
 
 export interface EditorFrameProps {
   datasourceMap: Record<string, Datasource>;
@@ -89,17 +90,31 @@ export function EditorFrame(props: EditorFrameProps) {
             dispatch={dispatch}
           />
         }
+        workspacePanel={
+          <WorkspacePanel
+            activeDatasource={props.datasourceMap[state.datasource.activeId]}
+            activeVisualizationId={state.visualization.activeId}
+            datasourcePublicAPI={datasourcePublicAPI!}
+            datasourceState={state.datasource.state}
+            activeVisualizationState={
+              state.visualization.activeId &&
+              state.visualization.stateMap[state.visualization.activeId]
+            }
+            visualizationMap={props.visualizationMap}
+            dispatch={dispatch}
+          />
+        }
         suggestionsPanel={
           <SuggestionPanel
             activeDatasource={props.datasourceMap[state.datasource.activeId]}
             activeVisualizationId={state.visualization.activeId}
             datasourcePublicAPI={datasourcePublicAPI!}
             datasourceState={state.datasource.state}
-            visualizationState={
+            activeVisualizationState={
               state.visualization.activeId &&
               state.visualization.stateMap[state.visualization.activeId]
             }
-            visualizations={props.visualizationMap}
+            visualizationMap={props.visualizationMap}
             dispatch={dispatch}
           />
         }
@@ -117,8 +132,6 @@ export function EditorFrame(props: EditorFrameProps) {
             dispatch={dispatch}
           />
         }
-        configPanel={null}
-        suggestionsPanel={null}
       />
     );
   }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx
index 7c7d0f23ca9cd..f62722bf71b85 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/frame_layout.tsx
@@ -9,8 +9,9 @@ import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
 
 export interface FrameLayoutProps {
   dataPanel: React.ReactNode;
-  configPanel: React.ReactNode;
+  configPanel?: React.ReactNode;
   suggestionsPanel?: React.ReactNode;
+  workspacePanel?: React.ReactNode;
 }
 
 export function FrameLayout(props: FrameLayoutProps) {
@@ -18,6 +19,7 @@ export function FrameLayout(props: FrameLayoutProps) {
     <EuiFlexGroup>
       {/* TODO style this and add workspace prop and loading flags */}
       <EuiFlexItem>{props.dataPanel}</EuiFlexItem>
+      <EuiFlexItem>{props.workspacePanel}</EuiFlexItem>
       <EuiFlexItem>
         {props.configPanel}
         {props.suggestionsPanel}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
new file mode 100644
index 0000000000000..67192321095f3
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
@@ -0,0 +1,45 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { DatasourcePublicAPI, TableColumn, Visualization, DatasourceSuggestion } from '../../types';
+
+export function getSuggestions(
+  datasourceTableSuggestions: DatasourceSuggestion[],
+  visualizationMap: Record<string, Visualization>,
+  activeVisualizationId: string | null,
+  activeVisualizationState: unknown,
+  datasourcePublicAPI: DatasourcePublicAPI
+) {
+  const roleMapping = activeVisualizationId
+    ? visualizationMap[activeVisualizationId].getMappingOfTableToRoles(
+        activeVisualizationState,
+        datasourcePublicAPI
+      )
+    : [];
+  const datasourceTableMetas: Record<string, TableColumn[]> = {};
+  datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
+    datasourceTableMetas[datasourceSuggestionId] = tableColumns;
+  });
+
+  return (
+    Object.entries(visualizationMap)
+      .map(([visualizationId, visualization]) => {
+        return visualization
+          .getSuggestions({
+            tableColumns: datasourceTableMetas,
+            roles: roleMapping,
+          })
+          .map(({ datasourceSuggestionId, ...suggestion }) => ({
+            ...suggestion,
+            visualizationId,
+            datasourceState: datasourceTableSuggestions[datasourceSuggestionId].state,
+          }));
+      })
+      // TODO why is flatMap not available here?
+      .reduce((globalList, currentList) => [...globalList, ...currentList], [])
+      .sort(({ score: scoreA }, { score: scoreB }) => scoreB - scoreA)
+  );
+}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
index 3ad3d7fccc00f..9766b6e2b6357 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
@@ -6,37 +6,40 @@
 
 import React from 'react';
 import { Action } from './state_management';
-import {
-  Datasource,
-  Visualization,
-  DatasourcePublicAPI,
-  DatasourceSuggestion,
-  DimensionRole,
-  TableColumn,
-} from '../../types';
+import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
+import { getSuggestions } from './suggestion_helpers';
 
 interface SuggestionPanelProps {
   activeDatasource: Datasource;
   datasourceState: unknown;
   activeVisualizationId: string | null;
-  visualizations: Record<string, Visualization>;
-  visualizationState: unknown;
+  visualizationMap: Record<string, Visualization>;
+  activeVisualizationState: unknown;
   datasourcePublicAPI: DatasourcePublicAPI;
   dispatch: (action: Action) => void;
 }
 
-export function SuggestionPanel(props: SuggestionPanelProps) {
-  const currentDatasource = props.activeDatasource;
-  const datasourceSuggestions = currentDatasource.getDatasourceSuggestionsFromCurrentState(
-    props.datasourceState
+export function SuggestionPanel({
+  activeDatasource,
+  datasourceState,
+  activeVisualizationId,
+  visualizationMap,
+  activeVisualizationState,
+  datasourcePublicAPI,
+  dispatch,
+}: SuggestionPanelProps) {
+  const datasourceSuggestions = activeDatasource.getDatasourceSuggestionsFromCurrentState(
+    datasourceState
   );
-  const roleMapping = props.activeVisualizationId
-    ? props.visualizations[props.activeVisualizationId].getMappingOfTableToRoles(
-        props.visualizationState,
-        props.datasourcePublicAPI
-      )
-    : [];
-  const suggestions = getSuggestions(datasourceSuggestions, props.visualizations, roleMapping);
+
+  const suggestions = getSuggestions(
+    datasourceSuggestions,
+    visualizationMap,
+    activeVisualizationId,
+    activeVisualizationState,
+    datasourcePublicAPI
+  );
+
   return (
     <>
       {/* TODO: I18N */}
@@ -47,7 +50,7 @@ export function SuggestionPanel(props: SuggestionPanelProps) {
             key={index}
             data-test-subj="suggestion"
             onClick={() => {
-              props.dispatch({
+              dispatch({
                 type: 'SWITCH_VISUALIZATION',
                 newVisualizationId: suggestion.visualizationId,
                 initialState: suggestion.state,
@@ -62,33 +65,3 @@ export function SuggestionPanel(props: SuggestionPanelProps) {
     </>
   );
 }
-
-function getSuggestions(
-  datasourceTableSuggestions: DatasourceSuggestion[],
-  visualizations: Record<string, Visualization>,
-  currentColumnRoles: DimensionRole[]
-) {
-  const datasourceTableMetas: Record<string, TableColumn[]> = {};
-  datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
-    datasourceTableMetas[datasourceSuggestionId] = tableColumns;
-  });
-
-  return (
-    Object.entries(visualizations)
-      .map(([visualizationId, visualization]) => {
-        return visualization
-          .getSuggestions({
-            tableColumns: datasourceTableMetas,
-            roles: currentColumnRoles,
-          })
-          .map(({ datasourceSuggestionId, ...suggestion }) => ({
-            ...suggestion,
-            visualizationId,
-            datasourceState: datasourceTableSuggestions[datasourceSuggestionId].state,
-          }));
-      })
-      // TODO why is flatMap not available here?
-      .reduce((globalList, currentList) => [...globalList, ...currentList], [])
-      .sort(({ score: scoreA }, { score: scoreB }) => scoreB - scoreA)
-  );
-}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
index 194c282214781..e2571a27d5f94 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
@@ -6,22 +6,16 @@
 
 import React from 'react';
 import { Action } from './state_management';
-import {
-  Datasource,
-  Visualization,
-  DatasourcePublicAPI,
-  DatasourceSuggestion,
-  DimensionRole,
-  TableColumn,
-} from '../../types';
+import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
 import { DragDrop } from '../../drag_drop';
+import { getSuggestions } from './suggestion_helpers';
 
 interface WorkspacePanelProps {
   activeDatasource: Datasource;
   datasourceState: unknown;
   activeVisualizationId: string | null;
-  visualizations: Record<string, Visualization>;
-  visualizationState: unknown;
+  visualizationMap: Record<string, Visualization>;
+  activeVisualizationState: unknown;
   datasourcePublicAPI: DatasourcePublicAPI;
   dispatch: (action: Action) => void;
 }
@@ -36,76 +30,64 @@ function ExpressionRenderer(props: ExpressionRendererProps) {
   return <span>{props.expression}</span>;
 }
 
-export function WorkspacePanel(props: WorkspacePanelProps) {
-  const currentDatasource = props.activeDatasource;
-  const expression = props.activeVisualizationId
-    ? `${props.activeDatasource.toExpression(props.datasourceState)} | ${props.visualizations[
-        props.activeVisualizationId
-      ].toExpression(props.visualizationState, props.datasourcePublicAPI)}`
-    : null;
-  return (
-    <DragDrop
-      draggable={false}
-      droppable={true}
-      onDrop={() => {
-        const datasourceSuggestions = currentDatasource.getDatasourceSuggestionsForField(
-          props.datasourceState
-        );
-        const roleMapping = props.activeVisualizationId
-          ? props.visualizations[props.activeVisualizationId].getMappingOfTableToRoles(
-              props.visualizationState,
-              props.datasourcePublicAPI
-            )
-          : [];
-        const suggestion = getSuggestions(
-          datasourceSuggestions,
-          props.visualizations,
-          roleMapping
-        )[0];
-        // TODO heuristically present the suggestions in a modal instead of just picking the first one
-        props.dispatch({
-          type: 'SWITCH_VISUALIZATION',
-          newVisualizationId: suggestion.visualizationId,
-          initialState: suggestion.state,
-          datasourceState: suggestion.datasourceState,
-        });
-      }}
-    >
-      {expression !== null ? (
-        <ExpressionRenderer expression={expression} />
-      ) : (
-        <p>{/* TODO: I18N */}This is the workspace panel. Drop fields here</p>
-      )}
-    </DragDrop>
-  );
-}
+export function WorkspacePanel({
+  activeDatasource,
+  activeVisualizationId,
+  datasourceState,
+  visualizationMap,
+  activeVisualizationState,
+  datasourcePublicAPI,
+  dispatch,
+}: WorkspacePanelProps) {
+  function onDrop() {
+    const datasourceSuggestions = activeDatasource.getDatasourceSuggestionsForField(
+      datasourceState
+    );
+
+    const suggestions = getSuggestions(
+      datasourceSuggestions,
+      visualizationMap,
+      activeVisualizationId,
+      activeVisualizationState,
+      datasourcePublicAPI
+    );
+
+    if (suggestions.length === 0) {
+      // TODO specify and implement behavior in case
+      // of no valid suggestions
+      return;
+    }
 
-function getSuggestions(
-  datasourceTableSuggestions: DatasourceSuggestion[],
-  visualizations: Record<string, Visualization>,
-  currentColumnRoles: DimensionRole[]
-) {
-  const datasourceTableMetas: Record<string, TableColumn[]> = {};
-  datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
-    datasourceTableMetas[datasourceSuggestionId] = tableColumns;
-  });
+    const suggestion = suggestions[0];
+
+    // TODO heuristically present the suggestions in a modal instead of just picking the first one
+    dispatch({
+      type: 'SWITCH_VISUALIZATION',
+      newVisualizationId: suggestion.visualizationId,
+      initialState: suggestion.state,
+      datasourceState: suggestion.datasourceState,
+    });
+  }
+
+  function renderVisualization() {
+    if (activeVisualizationId === null) {
+      return <p>{/* TODO: I18N */}This is the workspace panel. Drop fields here</p>;
+    }
+
+    const activeVisualization = visualizationMap[activeVisualizationId];
+    const datasourceExpression = activeDatasource.toExpression(datasourceState);
+    const visualizationExpression = activeVisualization.toExpression(
+      activeVisualizationState,
+      datasourcePublicAPI
+    );
+    const expression = `${datasourceExpression} | ${visualizationExpression}`;
+
+    return <ExpressionRenderer expression={expression} />;
+  }
 
   return (
-    Object.entries(visualizations)
-      .map(([visualizationId, visualization]) => {
-        return visualization
-          .getSuggestions({
-            tableColumns: datasourceTableMetas,
-            roles: currentColumnRoles,
-          })
-          .map(({ datasourceSuggestionId, ...suggestion }) => ({
-            ...suggestion,
-            visualizationId,
-            datasourceState: datasourceTableSuggestions[datasourceSuggestionId].state,
-          }));
-      })
-      // TODO why is flatMap not available here?
-      .reduce((globalList, currentList) => [...globalList, ...currentList], [])
-      .sort(({ score: scoreA }, { score: scoreB }) => scoreB - scoreA)
+    <DragDrop draggable={false} droppable={true} onDrop={onDrop}>
+      {renderVisualization()}
+    </DragDrop>
   );
 }

From 0c721a72e6b01f531725b36cdb2897997733ef82 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 15:43:38 +0200
Subject: [PATCH 32/41] more refactoring

---
 .../editor_frame/suggestion_helpers.ts        | 28 ++++++++++++++++++-
 .../editor_frame/suggestion_panel.tsx         |  9 ++----
 .../editor_frame/workspace_panel.tsx          |  9 ++----
 3 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
index 67192321095f3..671361722d138 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
@@ -5,14 +5,31 @@
  */
 
 import { DatasourcePublicAPI, TableColumn, Visualization, DatasourceSuggestion } from '../../types';
+import { Action } from './state_management';
 
+export interface Suggestion {
+  visualizationId: string;
+  datasourceState: unknown;
+  score: number;
+  title: string;
+  state: unknown;
+}
+
+/**
+ * This function takes a list of available data tables and a list of viusalization
+ * extensions and creates a ranked list of suggestions which contain a pair of a data table
+ * and a visualization.
+ *
+ * Each suggestion represents a valid state of the editor and can be applied by creating an
+ * action with `toSwitchAction` and dispatching it
+ */
 export function getSuggestions(
   datasourceTableSuggestions: DatasourceSuggestion[],
   visualizationMap: Record<string, Visualization>,
   activeVisualizationId: string | null,
   activeVisualizationState: unknown,
   datasourcePublicAPI: DatasourcePublicAPI
-) {
+): Suggestion[] {
   const roleMapping = activeVisualizationId
     ? visualizationMap[activeVisualizationId].getMappingOfTableToRoles(
         activeVisualizationState,
@@ -43,3 +60,12 @@ export function getSuggestions(
       .sort(({ score: scoreA }, { score: scoreB }) => scoreB - scoreA)
   );
 }
+
+export function toSwitchAction(suggestion: Suggestion): Action {
+  return {
+    type: 'SWITCH_VISUALIZATION',
+    newVisualizationId: suggestion.visualizationId,
+    initialState: suggestion.state,
+    datasourceState: suggestion.datasourceState,
+  };
+}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
index 9766b6e2b6357..bde7e56e9bcb4 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
@@ -7,7 +7,7 @@
 import React from 'react';
 import { Action } from './state_management';
 import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
-import { getSuggestions } from './suggestion_helpers';
+import { getSuggestions, toSwitchAction } from './suggestion_helpers';
 
 interface SuggestionPanelProps {
   activeDatasource: Datasource;
@@ -50,12 +50,7 @@ export function SuggestionPanel({
             key={index}
             data-test-subj="suggestion"
             onClick={() => {
-              dispatch({
-                type: 'SWITCH_VISUALIZATION',
-                newVisualizationId: suggestion.visualizationId,
-                initialState: suggestion.state,
-                datasourceState: suggestion.datasourceState,
-              });
+              dispatch(toSwitchAction(suggestion));
             }}
           >
             {suggestion.title}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
index e2571a27d5f94..5da08828b8f7e 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
@@ -8,7 +8,7 @@ import React from 'react';
 import { Action } from './state_management';
 import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
 import { DragDrop } from '../../drag_drop';
-import { getSuggestions } from './suggestion_helpers';
+import { getSuggestions, toSwitchAction } from './suggestion_helpers';
 
 interface WorkspacePanelProps {
   activeDatasource: Datasource;
@@ -61,12 +61,7 @@ export function WorkspacePanel({
     const suggestion = suggestions[0];
 
     // TODO heuristically present the suggestions in a modal instead of just picking the first one
-    dispatch({
-      type: 'SWITCH_VISUALIZATION',
-      newVisualizationId: suggestion.visualizationId,
-      initialState: suggestion.state,
-      datasourceState: suggestion.datasourceState,
-    });
+    dispatch(toSwitchAction(suggestion));
   }
 
   function renderVisualization() {

From 21c8756f5ae20e251ea30fee19c3ff8eccb1d83e Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 16:39:37 +0200
Subject: [PATCH 33/41] only save single visualization state

---
 .../editor_frame/config_panel_wrapper.tsx     |  4 +--
 .../editor_frame/editor_frame.tsx             |  2 +-
 .../editor_frame/state_management.test.ts     | 31 ++++++-------------
 .../editor_frame/state_management.ts          | 28 +++++------------
 4 files changed, 21 insertions(+), 44 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
index 1a8c9bfe3b298..b1329ee6fc2a8 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/config_panel_wrapper.tsx
@@ -11,7 +11,7 @@ import { Action } from './state_management';
 import { Visualization, DatasourcePublicAPI } from '../../types';
 
 interface ConfigPanelWrapperProps {
-  visualizationStateMap: Record<string, unknown>;
+  visualizationState: unknown;
   visualizationMap: Record<string, Visualization>;
   activeVisualizationId: string | null;
   dispatch: (action: Action) => void;
@@ -52,7 +52,7 @@ export function ConfigPanelWrapper(props: ConfigPanelWrapperProps) {
         <NativeRenderer
           render={props.visualizationMap[props.activeVisualizationId].renderConfigPanel}
           nativeProps={{
-            state: props.visualizationStateMap[props.activeVisualizationId],
+            state: props.visualizationState,
             setState: setVisualizationState,
             datasource: props.datasourcePublicAPI,
           }}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 159762939a58b..ee5023c06f3f7 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -86,7 +86,7 @@ export function EditorFrame(props: EditorFrameProps) {
             activeVisualizationId={state.visualization.activeId}
             datasourcePublicAPI={datasourcePublicAPI!}
             dispatch={dispatch}
-            visualizationStateMap={state.visualization.stateMap}
+            visualizationState={state.visualization.state}
           />
         )
       }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
index ea40f4e203eb0..373b321309586 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
@@ -33,14 +33,14 @@ describe('editor_frame state management', () => {
 
       const initialState = getInitialState(props);
 
-      expect(initialState.visualization.stateMap.testVis).toBe(initialVisState);
+      expect(initialState.visualization.state).toBe(initialVisState);
       expect(props.visualizationMap.testVis.initialize).toHaveBeenCalled();
     });
 
     it('should not initialize visualization if no initial visualization is passed in', () => {
       const initialState = getInitialState({ ...props, initialVisualizationId: null });
 
-      expect(initialState.visualization.stateMap).toEqual({});
+      expect(initialState.visualization.state).toEqual(null);
       expect(props.visualizationMap.testVis.initialize).not.toHaveBeenCalled();
     });
   });
@@ -57,9 +57,7 @@ describe('editor_frame state management', () => {
           },
           visualization: {
             activeId: 'testVis',
-            stateMap: {
-              testVis: {},
-            },
+            state: {},
           },
         },
         {
@@ -68,9 +66,7 @@ describe('editor_frame state management', () => {
         }
       );
 
-      expect(newState.visualization.stateMap).toEqual({
-        testVis: newVisState,
-      });
+      expect(newState.visualization.state).toBe(newVisState);
     });
 
     it('should update the datasource state on update', () => {
@@ -84,9 +80,7 @@ describe('editor_frame state management', () => {
           },
           visualization: {
             activeId: 'testVis',
-            stateMap: {
-              testVis: {},
-            },
+            state: {},
           },
         },
         {
@@ -98,7 +92,7 @@ describe('editor_frame state management', () => {
       expect(newState.datasource.state).toBe(newDatasourceState);
     });
 
-    it('should should switch active visualization but dont loose old state', () => {
+    it('should should switch active visualization', () => {
       const testVisState = {};
       const newVisState = {};
       const newState = reducer(
@@ -110,9 +104,7 @@ describe('editor_frame state management', () => {
           },
           visualization: {
             activeId: 'testVis',
-            stateMap: {
-              testVis: testVisState,
-            },
+            state: testVisState,
           },
         },
         {
@@ -122,8 +114,7 @@ describe('editor_frame state management', () => {
         }
       );
 
-      expect(newState.visualization.stateMap.testVis).toBe(testVisState);
-      expect(newState.visualization.stateMap.testVis2).toBe(newVisState);
+      expect(newState.visualization.state).toBe(newVisState);
     });
 
     it('should should switch active datasource and purge visualization state', () => {
@@ -136,9 +127,7 @@ describe('editor_frame state management', () => {
           },
           visualization: {
             activeId: 'testVis',
-            stateMap: {
-              testVis: {},
-            },
+            state: {},
           },
         },
         {
@@ -147,7 +136,7 @@ describe('editor_frame state management', () => {
         }
       );
 
-      expect(newState.visualization.stateMap).toEqual({});
+      expect(newState.visualization.state).toEqual(null);
       expect(newState.visualization.activeId).toBe(null);
       expect(newState.datasource.activeId).toBe('testDatasource2');
       expect(newState.datasource.state).toBe(null);
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
index e8b8dac09d5d3..2358da104378b 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
@@ -9,9 +9,7 @@ import { EditorFrameProps } from '.';
 export interface EditorFrameState {
   visualization: {
     activeId: string | null;
-    stateMap: {
-      [visualizationId: string]: unknown;
-    };
+    state: unknown;
   };
   datasource: {
     activeId: string | null;
@@ -47,13 +45,9 @@ export const getInitialState = (props: EditorFrameProps): EditorFrameState => {
       activeId: props.initialDatasourceId,
     },
     visualization: {
-      stateMap: props.initialVisualizationId
-        ? {
-            [props.initialVisualizationId]: props.visualizationMap[
-              props.initialVisualizationId
-            ].initialize(),
-          }
-        : {},
+      state: props.initialVisualizationId
+        ? props.visualizationMap[props.initialVisualizationId].initialize()
+        : null,
       activeId: props.initialVisualizationId,
     },
   };
@@ -72,8 +66,8 @@ export const reducer = (state: EditorFrameState, action: Action): EditorFrameSta
         },
         visualization: {
           ...state.visualization,
-          // purge all visualizations on datasource switch
-          stateMap: {},
+          // purge visualization on datasource switch
+          state: null,
           activeId: null,
         },
       };
@@ -83,10 +77,7 @@ export const reducer = (state: EditorFrameState, action: Action): EditorFrameSta
         visualization: {
           ...state.visualization,
           activeId: action.newVisualizationId,
-          stateMap: {
-            ...state.visualization.stateMap,
-            [action.newVisualizationId]: action.initialState,
-          },
+          state: action.initialState,
         },
       };
     case 'UPDATE_DATASOURCE_STATE':
@@ -107,10 +98,7 @@ export const reducer = (state: EditorFrameState, action: Action): EditorFrameSta
         ...state,
         visualization: {
           ...state.visualization,
-          stateMap: {
-            ...state.visualization.stateMap,
-            [state.visualization.activeId]: action.newState,
-          },
+          state: action.newState,
         },
       };
     default:

From fc47428b21838f65f32eb1ceb18991878b469f35 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 17:07:09 +0200
Subject: [PATCH 34/41] add comment

---
 x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx
index eb5326e21ebe9..1d81f315bf525 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.test.tsx
@@ -7,6 +7,8 @@
 import { EditorFramePlugin } from './plugin';
 import { createMockDatasource, createMockVisualization } from './mock_extensions';
 
+// calling this function will wait for all pending Promises from mock
+// datasources to be processed by its callers.
 const waitForPromises = () => new Promise(resolve => setTimeout(resolve));
 
 describe('editor_frame plugin', () => {

From eb8f2e8b2384ea96a9ea0092e8865e77f8113f11 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Tue, 21 May 2019 18:35:02 +0200
Subject: [PATCH 35/41] start adding tests for suggestion helpers and improve
 mock typing

---
 .../editor_frame/suggestion_helpers.test.ts   | 151 ++++++++++++++++++
 .../editor_frame/suggestion_helpers.ts        |   1 +
 .../editor_frame/suggestion_panel.tsx         |   2 +-
 .../editor_frame/workspace_panel.tsx          |   2 +-
 .../editor_frame_plugin/mock_extensions.ts    |  33 ++--
 5 files changed, 176 insertions(+), 13 deletions(-)
 create mode 100644 x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts
new file mode 100644
index 0000000000000..622cba5e995a4
--- /dev/null
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts
@@ -0,0 +1,151 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { getInitialState, reducer } from './state_management';
+import { EditorFrameProps } from '.';
+import { Datasource, Visualization } from '../../types';
+import { getSuggestions } from './suggestion_helpers';
+import { createMockVisualization, createMockDatasource } from '../mock_extensions';
+
+describe('suggestion helpers', () => {
+  it('should return suggestions array', () => {
+    const mockVisualization = createMockVisualization();
+    const mockDatasource = createMockDatasource();
+    const suggestedState = {};
+    const suggestions = getSuggestions(
+      [{ state: {}, tableColumns: [] }],
+      {
+        vis1: {
+          ...mockVisualization,
+          getSuggestions: () => [
+            { datasourceSuggestionId: 0, score: 0.5, title: 'Test', state: suggestedState },
+          ],
+        },
+      },
+      'vis1',
+      {},
+      mockDatasource.publicAPIMock
+    );
+    expect(suggestions.length).toBe(1);
+    expect(suggestions[0].state).toBe(suggestedState);
+  });
+
+  it('should concatenate suggestions from all visualizations', () => {
+    const mockVisualization1 = createMockVisualization();
+    const mockVisualization2 = createMockVisualization();
+    const mockDatasource = createMockDatasource();
+    const suggestions = getSuggestions(
+      [{ state: {}, tableColumns: [] }],
+      {
+        vis1: {
+          ...mockVisualization1,
+          getSuggestions: () => [
+            { datasourceSuggestionId: 0, score: 0.5, title: 'Test', state: {} },
+            { datasourceSuggestionId: 0, score: 0.5, title: 'Test2', state: {} },
+          ],
+        },
+        vis2: {
+          ...mockVisualization2,
+          getSuggestions: () => [
+            { datasourceSuggestionId: 0, score: 0.5, title: 'Test3', state: {} },
+          ],
+        },
+      },
+      'vis1',
+      {},
+      mockDatasource.publicAPIMock
+    );
+    expect(suggestions.length).toBe(3);
+  });
+
+  it('should rank the visualizations by score', () => {
+    const mockVisualization1 = createMockVisualization();
+    const mockVisualization2 = createMockVisualization();
+    const mockDatasource = createMockDatasource();
+    const suggestions = getSuggestions(
+      [{ state: {}, tableColumns: [] }],
+      {
+        vis1: {
+          ...mockVisualization1,
+          getSuggestions: () => [
+            { datasourceSuggestionId: 0, score: 0.2, title: 'Test', state: {} },
+            { datasourceSuggestionId: 0, score: 0.8, title: 'Test2', state: {} },
+          ],
+        },
+        vis2: {
+          ...mockVisualization2,
+          getSuggestions: () => [
+            { datasourceSuggestionId: 0, score: 0.6, title: 'Test3', state: {} },
+          ],
+        },
+      },
+      'vis1',
+      {},
+      mockDatasource.publicAPIMock
+    );
+    expect(suggestions[0].score).toBe(0.8);
+    expect(suggestions[1].score).toBe(0.6);
+    expect(suggestions[2].score).toBe(0.2);
+  });
+
+  it('should call all suggestion getters with all available data tables', () => {
+    const mockVisualization1 = createMockVisualization();
+    const mockVisualization2 = createMockVisualization();
+    const mockDatasource = createMockDatasource();
+    const tableState1 = {};
+    const table1 = [];
+    const tableState2 = {};
+    const table2 = [];
+    getSuggestions(
+      [{ state: tableState1, tableColumns: table1 }, { state: tableState2, tableColumns: table2 }],
+      {
+        vis1: mockVisualization1,
+        vis2: mockVisualization2,
+      },
+      'vis1',
+      {},
+      mockDatasource.publicAPIMock
+    );
+    expect(mockVisualization1.getSuggestions.mock.calls[0][0].tableColumns[0]).toBe(table1);
+    expect(mockVisualization1.getSuggestions.mock.calls[0][0].tableColumns[1]).toBe(table2);
+    expect(mockVisualization2.getSuggestions.mock.calls[0][0].tableColumns[0]).toBe(table1);
+    expect(mockVisualization2.getSuggestions.mock.calls[0][0].tableColumns[1]).toBe(table2);
+  });
+
+  it('should map the suggestion ids back to the correct datasource states', () => {
+    const mockVisualization1 = createMockVisualization();
+    const mockVisualization2 = createMockVisualization();
+    const mockDatasource = createMockDatasource();
+    const tableState1 = {};
+    const table1 = [];
+    const tableState2 = {};
+    const table2 = [];
+    const suggestions = getSuggestions(
+      [{ state: tableState1, tableColumns: table1 }, { state: tableState2, tableColumns: table2 }],
+      {
+        vis1: {
+          ...mockVisualization1,
+          getSuggestions: () => [
+            { datasourceSuggestionId: 0, score: 0.3, title: 'Test', state: {} },
+            { datasourceSuggestionId: 1, score: 0.2, title: 'Test2', state: {} },
+          ],
+        },
+        vis2: {
+          ...mockVisualization2,
+          getSuggestions: () => [
+            { datasourceSuggestionId: 1, score: 0.1, title: 'Test3', state: {} },
+          ],
+        },
+      },
+      'vis1',
+      {},
+      mockDatasource.publicAPIMock
+    );
+    expect(suggestions[0].datasourceState).toBe(tableState1);
+    expect(suggestions[1].datasourceState).toBe(tableState2);
+    expect(suggestions[1].datasourceState).toBe(tableState2);
+  });
+});
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
index 645eef8b40973..97b2b75494015 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
@@ -48,6 +48,7 @@ export function getSuggestions(
           .getSuggestions({
             tableColumns: datasourceTableMetas,
             roles: roleMapping,
+            state: visualizationId === activeVisualizationId ? visualizationState : undefined,
           })
           .map(({ datasourceSuggestionId, ...suggestion }) => ({
             ...suggestion,
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
index 47f4d279008dd..a979a16f2c4c2 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
@@ -9,7 +9,7 @@ import { Action } from './state_management';
 import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
 import { getSuggestions, toSwitchAction } from './suggestion_helpers';
 
-interface SuggestionPanelProps {
+export interface SuggestionPanelProps {
   activeDatasource: Datasource;
   datasourceState: unknown;
   activeVisualizationId: string | null;
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
index 729ed9af44e46..887bb91c0664e 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
@@ -10,7 +10,7 @@ import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
 import { DragDrop } from '../../drag_drop';
 import { getSuggestions, toSwitchAction } from './suggestion_helpers';
 
-interface WorkspacePanelProps {
+export interface WorkspacePanelProps {
   activeDatasource: Datasource;
   datasourceState: unknown;
   activeVisualizationId: string | null;
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
index a61b531b4494a..2d07d2af33f45 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
@@ -4,19 +4,30 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import { DatasourcePublicAPI } from '../types';
+import { DatasourcePublicAPI, Visualization, Datasource } from '../types';
 
-export const createMockVisualization = () => ({
-  getMappingOfTableToRoles: jest.fn(() => []),
-  getPersistableState: jest.fn(() => ({})),
-  getSuggestions: jest.fn(() => []),
-  initialize: jest.fn(() => {}),
-  renderConfigPanel: jest.fn(),
-  toExpression: jest.fn(() => ''),
-});
+type ArgumentType<T> = T extends (...args: infer A) => unknown ? A : unknown[];
+
+// This type returns a mocked version of type T to allow the caller to access
+// the mock APIs in a type safe way
+type Mocked<T> = {
+  [K in keyof T]: T[K] extends (...args: unknown[]) => unknown
+    ? jest.Mock<ReturnType<T[K]>, ArgumentType<T[K]> | []>
+    : T[K]
+};
+
+export const createMockVisualization = () =>
+  ({
+    getMappingOfTableToRoles: jest.fn(() => []),
+    getPersistableState: jest.fn(() => ({})),
+    getSuggestions: jest.fn(() => []),
+    initialize: jest.fn(() => ({})),
+    renderConfigPanel: jest.fn(),
+    toExpression: jest.fn(() => ''),
+  } as Mocked<Visualization>);
 
 export const createMockDatasource = () => {
-  const publicAPIMock: DatasourcePublicAPI = {
+  const publicAPIMock: Mocked<DatasourcePublicAPI> = {
     getTableSpec: jest.fn(() => []),
     getOperationForColumnId: jest.fn(() => null),
     renderDimensionPanel: jest.fn(),
@@ -37,5 +48,5 @@ export const createMockDatasource = () => {
     // this is an additional property which doesn't exist on real datasources
     // but can be used to validate whether specific API mock functions are called
     publicAPIMock,
-  };
+  } as Mocked<Datasource> & { publicAPIMock: Mocked<DatasourcePublicAPI> };
 };

From e0be75331e96e7d9d04cccf6e7d4d1f2fac2b052 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Wed, 22 May 2019 10:51:12 +0200
Subject: [PATCH 36/41] improve getSuggestions tests

---
 .../editor_frame/suggestion_helpers.test.ts   | 43 ++++++++++++-----
 .../editor_frame_plugin/mock_extensions.ts    | 46 +++++++++++--------
 2 files changed, 59 insertions(+), 30 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts
index 622cba5e995a4..096b48eefda0d 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts
@@ -4,11 +4,9 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import { getInitialState, reducer } from './state_management';
-import { EditorFrameProps } from '.';
-import { Datasource, Visualization } from '../../types';
 import { getSuggestions } from './suggestion_helpers';
 import { createMockVisualization, createMockDatasource } from '../mock_extensions';
+import { TableColumn } from '../../types';
 
 describe('suggestion helpers', () => {
   it('should return suggestions array', () => {
@@ -95,12 +93,10 @@ describe('suggestion helpers', () => {
     const mockVisualization1 = createMockVisualization();
     const mockVisualization2 = createMockVisualization();
     const mockDatasource = createMockDatasource();
-    const tableState1 = {};
-    const table1 = [];
-    const tableState2 = {};
-    const table2 = [];
+    const table1: TableColumn[] = [];
+    const table2: TableColumn[] = [];
     getSuggestions(
-      [{ state: tableState1, tableColumns: table1 }, { state: tableState2, tableColumns: table2 }],
+      [{ state: {}, tableColumns: table1 }, { state: {}, tableColumns: table2 }],
       {
         vis1: mockVisualization1,
         vis2: mockVisualization2,
@@ -120,11 +116,9 @@ describe('suggestion helpers', () => {
     const mockVisualization2 = createMockVisualization();
     const mockDatasource = createMockDatasource();
     const tableState1 = {};
-    const table1 = [];
     const tableState2 = {};
-    const table2 = [];
     const suggestions = getSuggestions(
-      [{ state: tableState1, tableColumns: table1 }, { state: tableState2, tableColumns: table2 }],
+      [{ state: tableState1, tableColumns: [] }, { state: tableState2, tableColumns: [] }],
       {
         vis1: {
           ...mockVisualization1,
@@ -148,4 +142,31 @@ describe('suggestion helpers', () => {
     expect(suggestions[1].datasourceState).toBe(tableState2);
     expect(suggestions[1].datasourceState).toBe(tableState2);
   });
+
+  it('should pass the state of the currently active visualization to getSuggestions', () => {
+    const mockVisualization1 = createMockVisualization();
+    const mockVisualization2 = createMockVisualization();
+    const mockDatasource = createMockDatasource();
+    const currentState = {};
+    getSuggestions(
+      [{ state: {}, tableColumns: [] }, { state: {}, tableColumns: [] }],
+      {
+        vis1: mockVisualization1,
+        vis2: mockVisualization2,
+      },
+      'vis1',
+      currentState,
+      mockDatasource.publicAPIMock
+    );
+    expect(mockVisualization1.getSuggestions).toHaveBeenCalledWith(
+      expect.objectContaining({
+        state: currentState,
+      })
+    );
+    expect(mockVisualization2.getSuggestions).not.toHaveBeenCalledWith(
+      expect.objectContaining({
+        state: currentState,
+      })
+    );
+  });
 });
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
index 2d07d2af33f45..277cc6ba423be 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
@@ -11,25 +11,33 @@ type ArgumentType<T> = T extends (...args: infer A) => unknown ? A : unknown[];
 // This type returns a mocked version of type T to allow the caller to access
 // the mock APIs in a type safe way
 type Mocked<T> = {
-  [K in keyof T]: T[K] extends (...args: unknown[]) => unknown
-    ? jest.Mock<ReturnType<T[K]>, ArgumentType<T[K]> | []>
+  // due to a bug in TS 3.3.3333 `(...args: unknown[]) => unknown` doesn't
+  // work here. After an upgrade to 3.4.3 this can be removed.
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  [K in keyof T]: T[K] extends (...args: any[]) => unknown
+    ? jest.Mock<ReturnType<T[K]>, ArgumentType<T[K]>>
     : T[K]
 };
 
-export const createMockVisualization = () =>
-  ({
-    getMappingOfTableToRoles: jest.fn(() => []),
-    getPersistableState: jest.fn(() => ({})),
-    getSuggestions: jest.fn(() => []),
-    initialize: jest.fn(() => ({})),
+export function createMockVisualization(): Mocked<Visualization> {
+  return {
+    getMappingOfTableToRoles: jest.fn((_state, _datasource) => []),
+    getPersistableState: jest.fn(_state => ({})),
+    getSuggestions: jest.fn(_options => []),
+    initialize: jest.fn(_state => ({})),
     renderConfigPanel: jest.fn(),
-    toExpression: jest.fn(() => ''),
-  } as Mocked<Visualization>);
+    toExpression: jest.fn((_state, _datasource) => ''),
+  };
+}
+
+export type DatasourceMock = Mocked<Datasource> & {
+  publicAPIMock: Mocked<DatasourcePublicAPI>;
+};
 
-export const createMockDatasource = () => {
+export function createMockDatasource(): DatasourceMock {
   const publicAPIMock: Mocked<DatasourcePublicAPI> = {
     getTableSpec: jest.fn(() => []),
-    getOperationForColumnId: jest.fn(() => null),
+    getOperationForColumnId: jest.fn(),
     renderDimensionPanel: jest.fn(),
     removeColumnInTableSpec: jest.fn(),
     moveColumnTo: jest.fn(),
@@ -37,16 +45,16 @@ export const createMockDatasource = () => {
   };
 
   return {
-    getDatasourceSuggestionsForField: jest.fn(() => []),
-    getDatasourceSuggestionsFromCurrentState: jest.fn(() => []),
+    getDatasourceSuggestionsForField: jest.fn(_state => []),
+    getDatasourceSuggestionsFromCurrentState: jest.fn(_state => []),
     getPersistableState: jest.fn(),
-    getPublicAPI: jest.fn(() => publicAPIMock),
-    initialize: jest.fn(() => Promise.resolve()),
+    getPublicAPI: jest.fn((_state, _setState) => publicAPIMock),
+    initialize: jest.fn(_state => Promise.resolve()),
     renderDataPanel: jest.fn(),
-    toExpression: jest.fn(() => ''),
+    toExpression: jest.fn(_state => ''),
 
     // this is an additional property which doesn't exist on real datasources
     // but can be used to validate whether specific API mock functions are called
     publicAPIMock,
-  } as Mocked<Datasource> & { publicAPIMock: Mocked<DatasourcePublicAPI> };
-};
+  };
+}

From e464ce26b6b13cfd75645007b1f9025f162c64e9 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Wed, 22 May 2019 11:26:43 +0200
Subject: [PATCH 37/41] add i18n

---
 .../editor_frame/suggestion_panel.tsx          | 10 ++++++++--
 .../editor_frame/workspace_panel.tsx           | 18 +++++++++++++++---
 .../lens/public/editor_frame_plugin/plugin.tsx | 17 ++++++++++-------
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
index a979a16f2c4c2..1efcc084d495f 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
@@ -5,6 +5,8 @@
  */
 
 import React from 'react';
+import { FormattedMessage } from '@kbn/i18n/react';
+
 import { Action } from './state_management';
 import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
 import { getSuggestions, toSwitchAction } from './suggestion_helpers';
@@ -42,8 +44,12 @@ export function SuggestionPanel({
 
   return (
     <>
-      {/* TODO: I18N */}
-      <h2>Suggestions</h2>
+      <h2>
+        <FormattedMessage
+          id="xpack.lens.editorFrame.suggestionPanelTitle"
+          defaultMessage="Suggestions"
+        />
+      </h2>
       {suggestions.map((suggestion, index) => {
         return (
           <button
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
index 887bb91c0664e..83d09fc5e4f05 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
@@ -5,6 +5,8 @@
  */
 
 import React from 'react';
+import { FormattedMessage } from '@kbn/i18n/react';
+
 import { Action } from './state_management';
 import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
 import { DragDrop } from '../../drag_drop';
@@ -53,8 +55,7 @@ export function WorkspacePanel({
     );
 
     if (suggestions.length === 0) {
-      // TODO specify and implement behavior in case
-      // of no valid suggestions
+      // TODO specify and implement behavior in case of no valid suggestions
       return;
     }
 
@@ -64,9 +65,20 @@ export function WorkspacePanel({
     dispatch(toSwitchAction(suggestion));
   }
 
+  function renderEmptyWorkspace() {
+    return (
+      <p>
+        <FormattedMessage
+          id="xpack.lens.editorFrame.emptyWorkspace"
+          defaultMessage="This is the workspace panel. Drop fields here"
+        />
+      </p>
+    );
+  }
+
   function renderVisualization() {
     if (activeVisualizationId === null) {
-      return <p>{/* TODO: I18N */}This is the workspace panel. Drop fields here</p>;
+      return renderEmptyWorkspace();
     }
 
     const activeVisualization = visualizationMap[activeVisualizationId];
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
index 910df2ab0666e..7c12d95cb69c1 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/plugin.tsx
@@ -6,8 +6,9 @@
 
 import React from 'react';
 import { render, unmountComponentAtNode } from 'react-dom';
-import { Datasource, Visualization, EditorFrameSetup, EditorFrameInstance } from '../types';
+import { I18nProvider } from '@kbn/i18n/react';
 
+import { Datasource, Visualization, EditorFrameSetup, EditorFrameInstance } from '../types';
 import { EditorFrame } from './editor_frame';
 
 export class EditorFramePlugin {
@@ -34,12 +35,14 @@ export class EditorFramePlugin {
         const firstVisualizationId = Object.keys(this.visualizations)[0];
 
         render(
-          <EditorFrame
-            datasourceMap={this.datasources}
-            visualizationMap={this.visualizations}
-            initialDatasourceId={firstDatasourceId || null}
-            initialVisualizationId={firstVisualizationId || null}
-          />,
+          <I18nProvider>
+            <EditorFrame
+              datasourceMap={this.datasources}
+              visualizationMap={this.visualizations}
+              initialDatasourceId={firstDatasourceId || null}
+              initialVisualizationId={firstVisualizationId || null}
+            />
+          </I18nProvider>,
           domElement
         );
       },

From 7f1d66d80a017e9d9c38deced173cb00c7e42576 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Wed, 22 May 2019 11:39:55 +0200
Subject: [PATCH 38/41] use jest types

---
 .../editor_frame_plugin/mock_extensions.ts    | 21 ++++---------------
 1 file changed, 4 insertions(+), 17 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
index 277cc6ba423be..cabffb1af6eb8 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
@@ -6,20 +6,7 @@
 
 import { DatasourcePublicAPI, Visualization, Datasource } from '../types';
 
-type ArgumentType<T> = T extends (...args: infer A) => unknown ? A : unknown[];
-
-// This type returns a mocked version of type T to allow the caller to access
-// the mock APIs in a type safe way
-type Mocked<T> = {
-  // due to a bug in TS 3.3.3333 `(...args: unknown[]) => unknown` doesn't
-  // work here. After an upgrade to 3.4.3 this can be removed.
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  [K in keyof T]: T[K] extends (...args: any[]) => unknown
-    ? jest.Mock<ReturnType<T[K]>, ArgumentType<T[K]>>
-    : T[K]
-};
-
-export function createMockVisualization(): Mocked<Visualization> {
+export function createMockVisualization(): jest.Mocked<Visualization> {
   return {
     getMappingOfTableToRoles: jest.fn((_state, _datasource) => []),
     getPersistableState: jest.fn(_state => ({})),
@@ -30,12 +17,12 @@ export function createMockVisualization(): Mocked<Visualization> {
   };
 }
 
-export type DatasourceMock = Mocked<Datasource> & {
-  publicAPIMock: Mocked<DatasourcePublicAPI>;
+export type DatasourceMock = jest.Mocked<Datasource> & {
+  publicAPIMock: jest.Mocked<DatasourcePublicAPI>;
 };
 
 export function createMockDatasource(): DatasourceMock {
-  const publicAPIMock: Mocked<DatasourcePublicAPI> = {
+  const publicAPIMock: jest.Mocked<DatasourcePublicAPI> = {
     getTableSpec: jest.fn(() => []),
     getOperationForColumnId: jest.fn(),
     renderDimensionPanel: jest.fn(),

From 1b1c9b782c3ef52a7a83c6fa81a1b767114d2c7e Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Thu, 23 May 2019 15:34:44 +0200
Subject: [PATCH 39/41] review fixes

---
 .../editor_frame/editor_frame.tsx             | 28 +++++++++----------
 .../editor_frame/suggestion_helpers.ts        |  2 +-
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 68015aa4c2f46..9ffb7efc89299 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -114,19 +114,19 @@ export function EditorFrame(props: EditorFrameProps) {
         }
       />
     );
-  } else {
-    return (
-      <FrameLayout
-        dataPanel={
-          <DataPanelWrapper
-            activeDatasource={state.datasource.activeId}
-            datasourceIsLoading={state.datasource.isLoading}
-            datasourceState={state.datasource.state}
-            datasourceMap={props.datasourceMap}
-            dispatch={dispatch}
-          />
-        }
-      />
-    );
   }
+
+  return (
+    <FrameLayout
+      dataPanel={
+        <DataPanelWrapper
+          activeDatasource={state.datasource.activeId}
+          datasourceIsLoading={state.datasource.isLoading}
+          datasourceState={state.datasource.state}
+          datasourceMap={props.datasourceMap}
+          dispatch={dispatch}
+        />
+      }
+    />
+  );
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
index 97b2b75494015..a1b6d90a38e78 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
@@ -16,7 +16,7 @@ export interface Suggestion {
 }
 
 /**
- * This function takes a list of available data tables and a list of viusalization
+ * This function takes a list of available data tables and a list of visualization
  * extensions and creates a ranked list of suggestions which contain a pair of a data table
  * and a visualization.
  *

From 8652dd28735b3f8afe7499f5aaed738fc2c67384 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Thu, 23 May 2019 16:57:27 +0200
Subject: [PATCH 40/41] remove roles from API

---
 .../editor_frame/suggestion_helpers.ts            |  7 -------
 .../public/editor_frame_plugin/mock_extensions.ts |  1 -
 x-pack/plugins/lens/public/types.ts               | 15 ---------------
 .../xy_visualization_plugin/xy_visualization.tsx  |  6 ++----
 4 files changed, 2 insertions(+), 27 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
index a1b6d90a38e78..01727d0345b25 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
@@ -30,12 +30,6 @@ export function getSuggestions(
   visualizationState: unknown,
   datasourcePublicAPI: DatasourcePublicAPI
 ): Suggestion[] {
-  const roleMapping = activeVisualizationId
-    ? visualizationMap[activeVisualizationId].getMappingOfTableToRoles(
-        visualizationState,
-        datasourcePublicAPI
-      )
-    : [];
   const datasourceTableMetas: Record<string, TableColumn[]> = {};
   datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
     datasourceTableMetas[datasourceSuggestionId] = tableColumns;
@@ -47,7 +41,6 @@ export function getSuggestions(
         return visualization
           .getSuggestions({
             tableColumns: datasourceTableMetas,
-            roles: roleMapping,
             state: visualizationId === activeVisualizationId ? visualizationState : undefined,
           })
           .map(({ datasourceSuggestionId, ...suggestion }) => ({
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
index cabffb1af6eb8..45162d7c07960 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/mock_extensions.ts
@@ -8,7 +8,6 @@ import { DatasourcePublicAPI, Visualization, Datasource } from '../types';
 
 export function createMockVisualization(): jest.Mocked<Visualization> {
   return {
-    getMappingOfTableToRoles: jest.fn((_state, _datasource) => []),
     getPersistableState: jest.fn(_state => ({})),
     getSuggestions: jest.fn(_options => []),
     initialize: jest.fn(_state => ({})),
diff --git a/x-pack/plugins/lens/public/types.ts b/x-pack/plugins/lens/public/types.ts
index 27fe9542987c0..b1f231227d26f 100644
--- a/x-pack/plugins/lens/public/types.ts
+++ b/x-pack/plugins/lens/public/types.ts
@@ -21,16 +21,6 @@ export interface EditorFrameSetup {
 // Hints the default nesting to the data source. 0 is the highest priority
 export type DimensionPriority = 0 | 1 | 2;
 
-// For switching between visualizations and correctly matching columns
-export type DimensionRole =
-  | 'splitChart'
-  | 'series'
-  | 'primary'
-  | 'secondary'
-  | 'color'
-  | 'size'
-  | string; // Some visualizations will use custom names that have other meaning
-
 export interface TableColumn {
   columnId: string;
   operation: Operation;
@@ -131,8 +121,6 @@ export interface VisualizationProps<T = unknown> {
 }
 
 export interface SuggestionRequest<T = unknown> {
-  // Roles currently being used
-  roles?: DimensionRole[];
   // It is up to the Visualization to rank these tables
   tableColumns: Record<number, TableColumn[]>;
   state?: T; // State is only passed if the visualization is active
@@ -155,9 +143,6 @@ export interface Visualization<T = unknown, P = unknown> {
 
   toExpression: (state: T, datasource: DatasourcePublicAPI) => string;
 
-  // Frame will request the list of roles currently being used when calling `getInitialStateFromOtherVisualization`
-  getMappingOfTableToRoles: (state: T, datasource: DatasourcePublicAPI) => DimensionRole[];
-
   // The frame will call this function on all visualizations when the table changes, or when
   // rendering additional ways of using the data
   getSuggestions: (options: SuggestionRequest<T>) => Array<VisualizationSuggestion<T>>;
diff --git a/x-pack/plugins/lens/public/xy_visualization_plugin/xy_visualization.tsx b/x-pack/plugins/lens/public/xy_visualization_plugin/xy_visualization.tsx
index 7ec2994161216..c5d69175e90cd 100644
--- a/x-pack/plugins/lens/public/xy_visualization_plugin/xy_visualization.tsx
+++ b/x-pack/plugins/lens/public/xy_visualization_plugin/xy_visualization.tsx
@@ -6,10 +6,10 @@
 
 import React from 'react';
 import { render } from 'react-dom';
-import { Visualization, DimensionRole } from '../types';
+import { Visualization } from '../types';
 
 export interface XyVisualizationState {
-  roles: DimensionRole[];
+  roles: string[];
 }
 
 export type XyVisualizationPersistedState = XyVisualizationState;
@@ -31,7 +31,5 @@ export const xyVisualization: Visualization<XyVisualizationState, XyVisualizatio
 
   getSuggestions: options => [],
 
-  getMappingOfTableToRoles: (state, datasource) => [],
-
   toExpression: state => '',
 };

From 082a2d3c57d18155c460106bfb829f1c5aec4520 Mon Sep 17 00:00:00 2001
From: Joe Reuter <johannes.reuter@elastic.co>
Date: Fri, 24 May 2019 10:11:08 +0200
Subject: [PATCH 41/41] simplify API and fix i18n context in tests

---
 .../editor_frame/editor_frame.test.tsx        | 23 +++----
 .../editor_frame/editor_frame.tsx             |  1 -
 .../editor_frame/suggestion_helpers.test.ts   | 60 +++++++------------
 .../editor_frame/suggestion_helpers.ts        | 14 ++---
 .../editor_frame/suggestion_panel.tsx         |  7 +--
 .../editor_frame/workspace_panel.tsx          |  3 +-
 x-pack/plugins/lens/public/types.ts           |  4 +-
 7 files changed, 43 insertions(+), 69 deletions(-)

diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
index c92f29eb2684c..904c11d2f29de 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.test.tsx
@@ -5,7 +5,8 @@
  */
 
 import React from 'react';
-import { mount, ReactWrapper } from 'enzyme';
+import { ReactWrapper } from 'enzyme';
+import { mountWithIntl as mount } from 'test_utils/enzyme_helpers';
 import { EditorFrame } from './editor_frame';
 import { Visualization, Datasource, DatasourcePublicAPI } from '../../types';
 import { act } from 'react-dom/test-utils';
@@ -474,13 +475,13 @@ describe('editor_frame', () => {
               ...mockVisualization,
               getSuggestions: () => [
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.5,
                   state: {},
                   title: 'Suggestion2',
                 },
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.8,
                   state: {},
                   title: 'Suggestion1',
@@ -491,13 +492,13 @@ describe('editor_frame', () => {
               ...mockVisualization,
               getSuggestions: () => [
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.4,
                   state: {},
                   title: 'Suggestion4',
                 },
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.45,
                   state: {},
                   title: 'Suggestion3',
@@ -539,7 +540,7 @@ describe('editor_frame', () => {
               ...mockVisualization,
               getSuggestions: () => [
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.8,
                   state: suggestionVisState,
                   title: 'Suggestion1',
@@ -594,13 +595,13 @@ describe('editor_frame', () => {
               ...mockVisualization,
               getSuggestions: () => [
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.2,
                   state: {},
                   title: 'Suggestion1',
                 },
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.8,
                   state: suggestionVisState,
                   title: 'Suggestion2',
@@ -647,13 +648,13 @@ describe('editor_frame', () => {
               ...mockVisualization,
               getSuggestions: () => [
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.2,
                   state: {},
                   title: 'Suggestion1',
                 },
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.6,
                   state: {},
                   title: 'Suggestion2',
@@ -664,7 +665,7 @@ describe('editor_frame', () => {
               ...mockVisualization2,
               getSuggestions: () => [
                 {
-                  datasourceSuggestionId: 0,
+                  tableIndex: 0,
                   score: 0.8,
                   state: suggestionVisState,
                   title: 'Suggestion3',
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index 9ffb7efc89299..fb0bb662915ba 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -105,7 +105,6 @@ export function EditorFrame(props: EditorFrameProps) {
           <SuggestionPanel
             activeDatasource={props.datasourceMap[state.datasource.activeId]}
             activeVisualizationId={state.visualization.activeId}
-            datasourcePublicAPI={datasourcePublicAPI!}
             datasourceState={state.datasource.state}
             visualizationState={state.visualization.state}
             visualizationMap={props.visualizationMap}
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts
index 096b48eefda0d..f79a4b1000991 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts
@@ -5,13 +5,12 @@
  */
 
 import { getSuggestions } from './suggestion_helpers';
-import { createMockVisualization, createMockDatasource } from '../mock_extensions';
+import { createMockVisualization } from '../mock_extensions';
 import { TableColumn } from '../../types';
 
 describe('suggestion helpers', () => {
   it('should return suggestions array', () => {
     const mockVisualization = createMockVisualization();
-    const mockDatasource = createMockDatasource();
     const suggestedState = {};
     const suggestions = getSuggestions(
       [{ state: {}, tableColumns: [] }],
@@ -19,13 +18,12 @@ describe('suggestion helpers', () => {
         vis1: {
           ...mockVisualization,
           getSuggestions: () => [
-            { datasourceSuggestionId: 0, score: 0.5, title: 'Test', state: suggestedState },
+            { tableIndex: 0, score: 0.5, title: 'Test', state: suggestedState },
           ],
         },
       },
       'vis1',
-      {},
-      mockDatasource.publicAPIMock
+      {}
     );
     expect(suggestions.length).toBe(1);
     expect(suggestions[0].state).toBe(suggestedState);
@@ -34,27 +32,23 @@ describe('suggestion helpers', () => {
   it('should concatenate suggestions from all visualizations', () => {
     const mockVisualization1 = createMockVisualization();
     const mockVisualization2 = createMockVisualization();
-    const mockDatasource = createMockDatasource();
     const suggestions = getSuggestions(
       [{ state: {}, tableColumns: [] }],
       {
         vis1: {
           ...mockVisualization1,
           getSuggestions: () => [
-            { datasourceSuggestionId: 0, score: 0.5, title: 'Test', state: {} },
-            { datasourceSuggestionId: 0, score: 0.5, title: 'Test2', state: {} },
+            { tableIndex: 0, score: 0.5, title: 'Test', state: {} },
+            { tableIndex: 0, score: 0.5, title: 'Test2', state: {} },
           ],
         },
         vis2: {
           ...mockVisualization2,
-          getSuggestions: () => [
-            { datasourceSuggestionId: 0, score: 0.5, title: 'Test3', state: {} },
-          ],
+          getSuggestions: () => [{ tableIndex: 0, score: 0.5, title: 'Test3', state: {} }],
         },
       },
       'vis1',
-      {},
-      mockDatasource.publicAPIMock
+      {}
     );
     expect(suggestions.length).toBe(3);
   });
@@ -62,27 +56,23 @@ describe('suggestion helpers', () => {
   it('should rank the visualizations by score', () => {
     const mockVisualization1 = createMockVisualization();
     const mockVisualization2 = createMockVisualization();
-    const mockDatasource = createMockDatasource();
     const suggestions = getSuggestions(
       [{ state: {}, tableColumns: [] }],
       {
         vis1: {
           ...mockVisualization1,
           getSuggestions: () => [
-            { datasourceSuggestionId: 0, score: 0.2, title: 'Test', state: {} },
-            { datasourceSuggestionId: 0, score: 0.8, title: 'Test2', state: {} },
+            { tableIndex: 0, score: 0.2, title: 'Test', state: {} },
+            { tableIndex: 0, score: 0.8, title: 'Test2', state: {} },
           ],
         },
         vis2: {
           ...mockVisualization2,
-          getSuggestions: () => [
-            { datasourceSuggestionId: 0, score: 0.6, title: 'Test3', state: {} },
-          ],
+          getSuggestions: () => [{ tableIndex: 0, score: 0.6, title: 'Test3', state: {} }],
         },
       },
       'vis1',
-      {},
-      mockDatasource.publicAPIMock
+      {}
     );
     expect(suggestions[0].score).toBe(0.8);
     expect(suggestions[1].score).toBe(0.6);
@@ -92,7 +82,6 @@ describe('suggestion helpers', () => {
   it('should call all suggestion getters with all available data tables', () => {
     const mockVisualization1 = createMockVisualization();
     const mockVisualization2 = createMockVisualization();
-    const mockDatasource = createMockDatasource();
     const table1: TableColumn[] = [];
     const table2: TableColumn[] = [];
     getSuggestions(
@@ -102,19 +91,17 @@ describe('suggestion helpers', () => {
         vis2: mockVisualization2,
       },
       'vis1',
-      {},
-      mockDatasource.publicAPIMock
+      {}
     );
-    expect(mockVisualization1.getSuggestions.mock.calls[0][0].tableColumns[0]).toBe(table1);
-    expect(mockVisualization1.getSuggestions.mock.calls[0][0].tableColumns[1]).toBe(table2);
-    expect(mockVisualization2.getSuggestions.mock.calls[0][0].tableColumns[0]).toBe(table1);
-    expect(mockVisualization2.getSuggestions.mock.calls[0][0].tableColumns[1]).toBe(table2);
+    expect(mockVisualization1.getSuggestions.mock.calls[0][0].tables[0]).toBe(table1);
+    expect(mockVisualization1.getSuggestions.mock.calls[0][0].tables[1]).toBe(table2);
+    expect(mockVisualization2.getSuggestions.mock.calls[0][0].tables[0]).toBe(table1);
+    expect(mockVisualization2.getSuggestions.mock.calls[0][0].tables[1]).toBe(table2);
   });
 
   it('should map the suggestion ids back to the correct datasource states', () => {
     const mockVisualization1 = createMockVisualization();
     const mockVisualization2 = createMockVisualization();
-    const mockDatasource = createMockDatasource();
     const tableState1 = {};
     const tableState2 = {};
     const suggestions = getSuggestions(
@@ -123,20 +110,17 @@ describe('suggestion helpers', () => {
         vis1: {
           ...mockVisualization1,
           getSuggestions: () => [
-            { datasourceSuggestionId: 0, score: 0.3, title: 'Test', state: {} },
-            { datasourceSuggestionId: 1, score: 0.2, title: 'Test2', state: {} },
+            { tableIndex: 0, score: 0.3, title: 'Test', state: {} },
+            { tableIndex: 1, score: 0.2, title: 'Test2', state: {} },
           ],
         },
         vis2: {
           ...mockVisualization2,
-          getSuggestions: () => [
-            { datasourceSuggestionId: 1, score: 0.1, title: 'Test3', state: {} },
-          ],
+          getSuggestions: () => [{ tableIndex: 1, score: 0.1, title: 'Test3', state: {} }],
         },
       },
       'vis1',
-      {},
-      mockDatasource.publicAPIMock
+      {}
     );
     expect(suggestions[0].datasourceState).toBe(tableState1);
     expect(suggestions[1].datasourceState).toBe(tableState2);
@@ -146,7 +130,6 @@ describe('suggestion helpers', () => {
   it('should pass the state of the currently active visualization to getSuggestions', () => {
     const mockVisualization1 = createMockVisualization();
     const mockVisualization2 = createMockVisualization();
-    const mockDatasource = createMockDatasource();
     const currentState = {};
     getSuggestions(
       [{ state: {}, tableColumns: [] }, { state: {}, tableColumns: [] }],
@@ -155,8 +138,7 @@ describe('suggestion helpers', () => {
         vis2: mockVisualization2,
       },
       'vis1',
-      currentState,
-      mockDatasource.publicAPIMock
+      currentState
     );
     expect(mockVisualization1.getSuggestions).toHaveBeenCalledWith(
       expect.objectContaining({
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
index 01727d0345b25..c6f5d4d4538f3 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.ts
@@ -4,7 +4,7 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import { DatasourcePublicAPI, TableColumn, Visualization, DatasourceSuggestion } from '../../types';
+import { Visualization, DatasourceSuggestion } from '../../types';
 import { Action } from './state_management';
 
 export interface Suggestion {
@@ -27,23 +27,19 @@ export function getSuggestions(
   datasourceTableSuggestions: DatasourceSuggestion[],
   visualizationMap: Record<string, Visualization>,
   activeVisualizationId: string | null,
-  visualizationState: unknown,
-  datasourcePublicAPI: DatasourcePublicAPI
+  visualizationState: unknown
 ): Suggestion[] {
-  const datasourceTableMetas: Record<string, TableColumn[]> = {};
-  datasourceTableSuggestions.map(({ tableColumns }, datasourceSuggestionId) => {
-    datasourceTableMetas[datasourceSuggestionId] = tableColumns;
-  });
+  const datasourceTables = datasourceTableSuggestions.map(({ tableColumns }) => tableColumns);
 
   return (
     Object.entries(visualizationMap)
       .map(([visualizationId, visualization]) => {
         return visualization
           .getSuggestions({
-            tableColumns: datasourceTableMetas,
+            tables: datasourceTables,
             state: visualizationId === activeVisualizationId ? visualizationState : undefined,
           })
-          .map(({ datasourceSuggestionId, ...suggestion }) => ({
+          .map(({ tableIndex: datasourceSuggestionId, ...suggestion }) => ({
             ...suggestion,
             visualizationId,
             datasourceState: datasourceTableSuggestions[datasourceSuggestionId].state,
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
index 1efcc084d495f..9d9730db37651 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_panel.tsx
@@ -8,7 +8,7 @@ import React from 'react';
 import { FormattedMessage } from '@kbn/i18n/react';
 
 import { Action } from './state_management';
-import { Datasource, Visualization, DatasourcePublicAPI } from '../../types';
+import { Datasource, Visualization } from '../../types';
 import { getSuggestions, toSwitchAction } from './suggestion_helpers';
 
 export interface SuggestionPanelProps {
@@ -17,7 +17,6 @@ export interface SuggestionPanelProps {
   activeVisualizationId: string | null;
   visualizationMap: Record<string, Visualization>;
   visualizationState: unknown;
-  datasourcePublicAPI: DatasourcePublicAPI;
   dispatch: (action: Action) => void;
 }
 
@@ -27,7 +26,6 @@ export function SuggestionPanel({
   activeVisualizationId,
   visualizationMap,
   visualizationState,
-  datasourcePublicAPI,
   dispatch,
 }: SuggestionPanelProps) {
   const datasourceSuggestions = activeDatasource.getDatasourceSuggestionsFromCurrentState(
@@ -38,8 +36,7 @@ export function SuggestionPanel({
     datasourceSuggestions,
     visualizationMap,
     activeVisualizationId,
-    visualizationState,
-    datasourcePublicAPI
+    visualizationState
   );
 
   return (
diff --git a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
index 83d09fc5e4f05..761b77757df62 100644
--- a/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/workspace_panel.tsx
@@ -50,8 +50,7 @@ export function WorkspacePanel({
       datasourceSuggestions,
       visualizationMap,
       activeVisualizationId,
-      visualizationState,
-      datasourcePublicAPI
+      visualizationState
     );
 
     if (suggestions.length === 0) {
diff --git a/x-pack/plugins/lens/public/types.ts b/x-pack/plugins/lens/public/types.ts
index b1f231227d26f..4e2c623416c9a 100644
--- a/x-pack/plugins/lens/public/types.ts
+++ b/x-pack/plugins/lens/public/types.ts
@@ -122,7 +122,7 @@ export interface VisualizationProps<T = unknown> {
 
 export interface SuggestionRequest<T = unknown> {
   // It is up to the Visualization to rank these tables
-  tableColumns: Record<number, TableColumn[]>;
+  tables: TableColumn[][];
   state?: T; // State is only passed if the visualization is active
 }
 
@@ -130,7 +130,7 @@ export interface VisualizationSuggestion<T = unknown> {
   score: number;
   title: string;
   state: T;
-  datasourceSuggestionId: number;
+  tableIndex: number;
 }
 
 export interface Visualization<T = unknown, P = unknown> {