diff --git a/src/components/ChartProvider.js b/src/components/ChartProvider.js
index cfca443246..b3c440c319 100644
--- a/src/components/ChartProvider.js
+++ b/src/components/ChartProvider.js
@@ -6,8 +6,8 @@ const throwIfNotInitialized = () => {
 }
 
 export const ChartContext = createContext({
-    getChartInstance: throwIfNotInitialized,
-    setChartInstance: throwIfNotInitialized,
+    getChart: throwIfNotInitialized,
+    setChart: throwIfNotInitialized,
     isHighchartsChartInstance: throwIfNotInitialized,
 })
 
diff --git a/src/components/Visualization/Visualization.js b/src/components/Visualization/Visualization.js
index d61fe63373..31ac541f8a 100644
--- a/src/components/Visualization/Visualization.js
+++ b/src/components/Visualization/Visualization.js
@@ -269,6 +269,12 @@ UnconnectedVisualization.propTypes = {
 }
 
 UnconnectedVisualization.contextType = ChartContext
+// Needed for Jest/Enzyme context mocking to work
+UnconnectedVisualization.contextTypes = {
+    getChart: PropTypes.func,
+    setChart: PropTypes.func,
+    isHighchartsChartInstance: PropTypes.func,
+}
 
 const mapStateToProps = (state) => ({
     visualization: sGetCurrent(state),
diff --git a/src/components/Visualization/__tests__/Visualization.spec.js b/src/components/Visualization/__tests__/Visualization.spec.js
index 98f17d26ef..d18268091c 100644
--- a/src/components/Visualization/__tests__/Visualization.spec.js
+++ b/src/components/Visualization/__tests__/Visualization.spec.js
@@ -10,9 +10,14 @@ describe('Visualization', () => {
     describe('component', () => {
         let props
         let shallowVisualization
+        const setChart = jest.fn()
         const vis = () => {
             if (!shallowVisualization) {
-                shallowVisualization = shallow(<Visualization {...props} />)
+                shallowVisualization = shallow(<Visualization {...props} />, {
+                    context: {
+                        setChart,
+                    },
+                })
             }
             return shallowVisualization
         }
@@ -26,7 +31,6 @@ describe('Visualization', () => {
                 error: null,
                 rightSidebarOpen: false,
                 addMetadata: jest.fn(),
-                setChart: jest.fn(),
                 clearLoadError: jest.fn(),
                 setLoadError: jest.fn(),
                 onLoadingComplete: jest.fn(),
@@ -71,12 +75,12 @@ describe('Visualization', () => {
         })
 
         it('triggers setChart action when chart has been generated', () => {
-            const svg = 'coolChart'
+            const highChartChartInstanceMock = {}
 
-            vis().instance().onChartGenerated(svg)
+            vis().instance().onChartGenerated(highChartChartInstanceMock)
 
-            expect(props.setChart).toHaveBeenCalled()
-            expect(props.setChart).toHaveBeenCalledWith(svg)
+            expect(setChart).toHaveBeenCalled()
+            expect(setChart).toHaveBeenCalledWith(highChartChartInstanceMock)
         })
 
         it('renders visualization with new id when rightSidebarOpen prop changes', () => {
diff --git a/src/components/VisualizationPlugin/__tests__/ChartPlugin.spec.js b/src/components/VisualizationPlugin/__tests__/ChartPlugin.spec.js
index 4d5aff6482..13043ae667 100644
--- a/src/components/VisualizationPlugin/__tests__/ChartPlugin.spec.js
+++ b/src/components/VisualizationPlugin/__tests__/ChartPlugin.spec.js
@@ -72,7 +72,7 @@ class MockAnalyticsResponse {
 
 const createVisualizationMock = {
     visualization: {
-        getSVGForExport: () => '<svg />',
+        exportChartLocal: jest.fn(),
     },
     config: {
         getConfig: () => {},
@@ -148,7 +148,7 @@ describe('ChartPlugin', () => {
             setTimeout(() => {
                 expect(props.onChartGenerated).toHaveBeenCalled()
                 expect(props.onChartGenerated).toHaveBeenCalledWith(
-                    createVisualizationMock.visualization.getSVGForExport()
+                    createVisualizationMock.visualization
                 )
                 done()
             })