diff --git a/src/components/test/code-snippet.jsx b/src/components/test/code-snippet.jsx
index 7cda7d8e..46cc118e 100644
--- a/src/components/test/code-snippet.jsx
+++ b/src/components/test/code-snippet.jsx
@@ -54,8 +54,10 @@ class CodeSnippet extends Component {
   render() {
     const { className, code, lang, label, showLabel, useInlineDiffs } = this.props;
     const isDiff = lang === 'diff';
-    const cxName = cx(className, lang, {
-      hljs: !this.shouldHighlight(),
+    const shouldHighlight = this.shouldHighlight();
+    const cxName = cx(className, {
+      [lang]: shouldHighlight,
+      hljs: !shouldHighlight,
       'code-diff': isDiff,
       'inline-diff': isDiff && useInlineDiffs
     });
diff --git a/src/components/test/context.jsx b/src/components/test/context.jsx
index c2452c3d..748e3598 100644
--- a/src/components/test/context.jsx
+++ b/src/components/test/context.jsx
@@ -63,8 +63,13 @@ class TestContext extends Component {
       return this.renderLink(content, title);
     }
 
-    // Default
-    const code = isString(content) ? content : JSON.stringify(content, null, 2);
+    // Simple string
+    if (isString(content)) {
+      return <CodeSnippet className={ cx('code-snippet') } code={ content } highlight={ false } />;
+    }
+
+    // All other types (primitives, objects, arrays...)
+    const code = JSON.stringify(content, null, 2);
     return (
       <CodeSnippet className={ cx('code-snippet') } code={ code } highlight={ highlight } />
     );
diff --git a/test/spec/components/test/code-snippet.test.jsx b/test/spec/components/test/code-snippet.test.jsx
index a94523c9..60fe2b1d 100644
--- a/test/spec/components/test/code-snippet.test.jsx
+++ b/test/spec/components/test/code-snippet.test.jsx
@@ -39,7 +39,8 @@ describe('<CodeSnippet />', () => {
     const props = {
       code: 'function(){console.log(\'sample code\');}'
     };
-    getInstance(props);
+    const wrapper = getInstance(props);
+    expect(wrapper.hasClass('javascript')).to.equal(true);
     expect(document.querySelectorAll('.hljs-keyword').length).to.equal(1);
   });
 
@@ -49,6 +50,7 @@ describe('<CodeSnippet />', () => {
       lang: 'diff'
     };
     const wrapper = getInstance(props);
+    expect(wrapper.hasClass('diff')).to.equal(true);
     expect(wrapper.hasClass('inline-diff')).to.equal(false);
     expect(document.querySelectorAll('.test-code-diff-expected').length).to.equal(1);
     expect(document.querySelectorAll('.test-code-diff-actual').length).to.equal(1);
@@ -67,6 +69,7 @@ describe('<CodeSnippet />', () => {
       lang: 'diff'
     };
     const wrapper = getInstance(props, { useInlineDiffs: true });
+    expect(wrapper.hasClass('diff')).to.equal(false);
     expect(wrapper.hasClass('inline-diff')).to.equal(true);
     expect(document.querySelectorAll('.test-code-diff-expected').length).to.equal(2);
     expect(document.querySelectorAll('.test-code-diff-actual').length).to.equal(2);
@@ -75,9 +78,11 @@ describe('<CodeSnippet />', () => {
   it('does not highlight when prop is false', () => {
     const props = {
       code: 'function(){console.log(\'sample code\');}',
-      highlight: false
+      highlight: false,
+      lang: 'piglatin'
     };
-    getInstance(props);
+    const wrapper = getInstance(props);
+    expect(wrapper.hasClass('piglatin')).to.equal(false);
     expect(document.querySelectorAll('.hljs-keyword').length).to.equal(0);
   });
 
diff --git a/test/spec/components/test/context.test.jsx b/test/spec/components/test/context.test.jsx
index 074289a2..c75daf7d 100644
--- a/test/spec/components/test/context.test.jsx
+++ b/test/spec/components/test/context.test.jsx
@@ -21,143 +21,152 @@ describe('<TestContext />', () => {
     };
   };
 
-  it('renders context when value is undefined', () => {
-    const context = {
-      title: 'sample context',
-      value: 'undefined'
-    };
-    const { wrapper, ctxItems } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+  describe('when context is a string', () => {
+    it('renders simple string', () => {
+      const context = 'sample context';
+      const { wrapper, snippet } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(1);
     });
-    expect(wrapper).to.have.className('test');
-    expect(ctxItems).to.have.lengthOf(1);
-  });
 
-  it('renders context with string', () => {
-    const context = 'sample context';
-    const { wrapper, snippet } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+    it('renders url with protocol', () => {
+      const context = 'http://test.url.com/somepath';
+      const { wrapper, snippet, link } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(0);
+      expect(link).to.have.lengthOf(1);
+      link.simulate('click', { stopPropagation: noop });
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(1);
-  });
 
-  it('renders context with string, url with protocol', () => {
-    const context = 'http://test.url.com/somepath';
-    const { wrapper, snippet, link } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+    it('renders url without protocol', () => {
+      const context = 'test.url.com/somepath';
+      const { wrapper, snippet, link } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(0);
+      expect(link).to.have.lengthOf(1);
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(0);
-    expect(link).to.have.lengthOf(1);
-    link.simulate('click', { stopPropagation: noop });
-  });
 
-  it('renders context with string, url without protocol', () => {
-    const context = 'test.url.com/somepath';
-    const { wrapper, snippet, link } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+    it('renders image url with protocol', () => {
+      const context = 'http://test.url.com/testimage.png';
+      const { wrapper, snippet, img } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(0);
+      expect(img).to.have.lengthOf(1);
+      img.simulate('click', { stopPropagation: noop });
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(0);
-    expect(link).to.have.lengthOf(1);
-  });
 
-  it('renders context with string, image url with protocol', () => {
-    const context = 'http://test.url.com/testimage.png';
-    const { wrapper, snippet, img } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+    it('renders image url without protocol', () => {
+      const context = 'test.url.com/testimage.png';
+      const { wrapper, snippet, img } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(0);
+      expect(img).to.have.lengthOf(1);
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(0);
-    expect(img).to.have.lengthOf(1);
-    img.simulate('click', { stopPropagation: noop });
-  });
 
-  it('renders context with string, image url without protocol', () => {
-    const context = 'test.url.com/testimage.png';
-    const { wrapper, snippet, img } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+    it('renders local image', () => {
+      const context = '/testimage.png';
+      const { wrapper, snippet, img } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(0);
+      expect(img).to.have.lengthOf(1);
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(0);
-    expect(img).to.have.lengthOf(1);
   });
 
-  it('renders context with string, local image', () => {
-    const context = '/testimage.png';
-    const { wrapper, snippet, img } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+  describe('when context is an object', () => {
+    it('renders undefined value', () => {
+      const context = {
+        title: 'sample context',
+        value: 'undefined'
+      };
+      const { wrapper, ctxItems } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(ctxItems).to.have.lengthOf(1);
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(0);
-    expect(img).to.have.lengthOf(1);
-  });
 
-  it('renders context with object, string value', () => {
-    const context = {
-      title: 'sample context',
-      value: 'context string'
-    };
-    const { wrapper, snippet } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+
+    it('renders string value', () => {
+      const context = {
+        title: 'sample context',
+        value: 'context string'
+      };
+      const { wrapper, snippet } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(1);
+      expect(snippet.prop('highlight')).to.equal(false);
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(1);
-  });
 
-  it('renders context with object, image url value', () => {
-    const context = {
-      title: 'sample context',
-      value: 'http://test.url.com/testimage.png'
-    };
-    const { wrapper, snippet, img } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+    it('renders image url value', () => {
+      const context = {
+        title: 'sample context',
+        value: 'http://test.url.com/testimage.png'
+      };
+      const { wrapper, snippet, img } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(0);
+      expect(img).to.have.lengthOf(1);
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(0);
-    expect(img).to.have.lengthOf(1);
-  });
 
-  it('renders context with object, object value', () => {
-    const context = {
-      title: 'sample context',
-      value: { testing: true }
-    };
-    const { wrapper, snippet } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+    it('renders object value', () => {
+      const context = {
+        title: 'sample context',
+        value: { testing: true }
+      };
+      const { wrapper, snippet } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(1);
+      expect(snippet.prop('highlight')).to.equal(true);
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(1);
   });
 
-  it('renders context with array', () => {
-    const context = [
-      {
-        title: 'sample context a',
-        value: 'http://test.url.com/testimage.png'
-      },
-      {
-        title: 'sample context b',
-        value: { testing: true }
-      }
-    ];
-    const { wrapper, snippet, img } = getInstance({
-      context: JSON.stringify(context),
-      className: 'test'
+  describe('when context is an array', () => {
+    it('renders', () => {
+      const context = [
+        {
+          title: 'sample context a',
+          value: 'http://test.url.com/testimage.png'
+        },
+        {
+          title: 'sample context b',
+          value: { testing: true }
+        }
+      ];
+      const { wrapper, snippet, img } = getInstance({
+        context: JSON.stringify(context),
+        className: 'test'
+      });
+      expect(wrapper).to.have.className('test');
+      expect(snippet).to.have.lengthOf(1);
+      expect(img).to.have.lengthOf(1);
     });
-    expect(wrapper).to.have.className('test');
-    expect(snippet).to.have.lengthOf(1);
-    expect(img).to.have.lengthOf(1);
   });
 });