From 1b6e0e76979889b8977df7ba3bc3a92cde3dec17 Mon Sep 17 00:00:00 2001
From: Henry Zhu <hi@henryzoo.com>
Date: Mon, 7 Aug 2017 17:06:00 -0400
Subject: [PATCH] fix transform (#154)

---
 packages/rollup-plugin-babel/src/helperPlugin.js  | 4 ++--
 packages/rollup-plugin-babel/test/helperPlugin.js | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/packages/rollup-plugin-babel/src/helperPlugin.js b/packages/rollup-plugin-babel/src/helperPlugin.js
index 844c10848..6c367cdfa 100644
--- a/packages/rollup-plugin-babel/src/helperPlugin.js
+++ b/packages/rollup-plugin-babel/src/helperPlugin.js
@@ -1,8 +1,8 @@
 export default function importHelperPlugin () {
 	return {
 		visitor: {
-			Program (path, state) {
-				state.file.addHelper('classCallCheck');
+			ClassDeclaration (path, state) {
+				path.replaceWith(state.file.addHelper('classCallCheck'));
 			}
 		}
 	};
diff --git a/packages/rollup-plugin-babel/test/helperPlugin.js b/packages/rollup-plugin-babel/test/helperPlugin.js
index 18b9f3442..4fd937bcd 100644
--- a/packages/rollup-plugin-babel/test/helperPlugin.js
+++ b/packages/rollup-plugin-babel/test/helperPlugin.js
@@ -1,8 +1,8 @@
 module.exports = function importHelperPlugin () {
   return {
     visitor: {
-      Program (path, state) {
-        state.file.addHelper('classCallCheck');
+      ClassDeclaration (path, state) {
+        path.replaceWith(state.file.addHelper('classCallCheck'));
       }
     }
   };