diff --git a/packages/grafana-ui/src/components/Badge/Badge.story.tsx b/packages/grafana-ui/src/components/Badge/Badge.story.tsx
index 2ce4585b2fc..3c3c3948983 100644
--- a/packages/grafana-ui/src/components/Badge/Badge.story.tsx
+++ b/packages/grafana-ui/src/components/Badge/Badge.story.tsx
@@ -1,6 +1,7 @@
 import React from 'react';
-import { boolean, text, select } from '@storybook/addon-knobs';
-import { Badge, BadgeColor } from '@grafana/ui';
+import { Story } from '@storybook/react';
+import { Badge, BadgeProps } from '@grafana/ui';
+import { iconOptions } from '../../utils/storybook/knobs';
 
 export default {
   title: 'Data Display/Badge',
@@ -8,28 +9,22 @@ export default {
   decorators: [],
   parameters: {
     docs: {},
+    knobs: {
+      disabled: true,
+    },
+  },
+  argTypes: {
+    icon: { control: { type: 'select', options: iconOptions } },
+    color: { control: 'select' },
   },
 };
 
-export const basic = () => {
-  const badgeColor = select<BadgeColor>(
-    'Badge color',
-    {
-      Red: 'red',
-      Green: 'green',
-      Blue: 'blue',
-      Orange: 'orange',
-    },
-    'blue'
-  );
-  const withIcon = boolean('With icon', true);
-  const tooltipText = text('Tooltip text', '');
-  return (
-    <Badge
-      text={'Badge label'}
-      color={badgeColor}
-      icon={withIcon ? 'rocket' : undefined}
-      tooltip={tooltipText.trim() === '' ? undefined : tooltipText}
-    />
-  );
+const Template: Story<BadgeProps> = args => <Badge {...args} />;
+
+export const Basic = Template.bind({});
+
+Basic.args = {
+  text: 'Badge label',
+  color: 'blue',
+  icon: 'rocket',
 };