diff --git a/src/components/fields/TraceSelector.js b/src/components/fields/TraceSelector.js
index 48c02e958..00270429c 100644
--- a/src/components/fields/TraceSelector.js
+++ b/src/components/fields/TraceSelector.js
@@ -118,6 +118,7 @@ class TraceSelector extends Component {
                 handleClick={() =>
                   this.context.openModal(TraceTypeSelector, {
                     ...props,
+                    traceTypesConfig: this.context.traceTypesConfig,
                     glByDefault: this.context.glByDefault,
                   })
                 }
diff --git a/src/components/widgets/TraceTypeSelector.js b/src/components/widgets/TraceTypeSelector.js
index 7a5883065..d249d19b2 100644
--- a/src/components/widgets/TraceTypeSelector.js
+++ b/src/components/widgets/TraceTypeSelector.js
@@ -138,12 +138,10 @@ class TraceTypeSelector extends Component {
 
   renderCategories() {
     const {fullValue} = this.props;
+    const {mapBoxAccess, localize: _, chartHelp} = this.context;
     const {
       traceTypesConfig: {traces, categories, complex},
-      mapBoxAccess,
-      localize: _,
-      chartHelp,
-    } = this.context;
+    } = this.props;
 
     return categories(_).map((category, i) => {
       let items = traces(_)
@@ -156,14 +154,11 @@ class TraceTypeSelector extends Component {
 
       const MAX_ITEMS = 4;
 
-      let columnClasses = 'trace-grid__column';
-
-      if (
+      const columnClasses =
         (items.length > MAX_ITEMS && !category.maxColumns) ||
         (category.maxColumns && category.maxColumns > 1)
-      ) {
-        columnClasses += ' trace-grid__column--double';
-      }
+          ? 'trace-grid__column trace-grid__column--double'
+          : 'trace-grid__column';
 
       return (
         <div className={columnClasses} key={i}>
@@ -192,43 +187,34 @@ class TraceTypeSelector extends Component {
 
   renderSingleBlock() {
     const {fullValue} = this.props;
+    const {localize: _} = this.context;
     const {
       traceTypesConfig: {traces, complex},
-      localize: _,
-    } = this.context;
-
-    const items = traces(_).map(item => (
-      <Item
-        key={item.value}
-        complex={complex}
-        active={fullValue === item.value}
-        item={item}
-        actions={this.actions}
-        showActions={false}
-        handleClick={() => this.selectAndClose(item.value)}
-        style={{display: 'inline-block'}}
-      />
-    ));
+    } = this.props;
 
     return (
-      <div
-        style={{
-          maxWidth: '460px',
-          display: 'flex',
-          flexFlow: 'wrap',
-          padding: '5px',
-        }}
-      >
-        {items}
+      <div className="trace-grid-single-block">
+        {traces(_).map(item => (
+          <Item
+            key={item.value}
+            complex={complex}
+            active={fullValue === item.value}
+            item={item}
+            actions={this.actions}
+            showActions={false}
+            handleClick={() => this.selectAndClose(item.value)}
+            style={{display: 'inline-block'}}
+          />
+        ))}
       </div>
     );
   }
 
   render() {
+    const {localize: _} = this.context;
     const {
       traceTypesConfig: {categories},
-      localize: _,
-    } = this.context;
+    } = this.props;
 
     return (
       <Modal title={_('Select Trace Type')}>
@@ -243,9 +229,9 @@ TraceTypeSelector.propTypes = {
   fullValue: PropTypes.string,
   fullContainer: PropTypes.object,
   glByDefault: PropTypes.bool,
+  traceTypesConfig: PropTypes.object,
 };
 TraceTypeSelector.contextTypes = {
-  traceTypesConfig: PropTypes.object,
   handleClose: PropTypes.func,
   localize: PropTypes.func,
   mapBoxAccess: PropTypes.bool,
diff --git a/src/styles/components/widgets/_trace-type-selector.scss b/src/styles/components/widgets/_trace-type-selector.scss
index ab2de3427..6a77f6288 100644
--- a/src/styles/components/widgets/_trace-type-selector.scss
+++ b/src/styles/components/widgets/_trace-type-selector.scss
@@ -4,6 +4,13 @@ $label-height: 34px;
 $row-height: $image-size + $label-height + $default-half-spacing-unit +
   $default-quarter-spacing-unit;
 
+.trace-grid-single-block {
+  max-width: 460px;
+  display: flex;
+  flex-flow: wrap;
+  padding: var(--spacing-quarter-unit);
+}
+
 .trace-grid {
   display: grid;
   grid-template-columns: repeat(4, 1fr);