diff --git a/test/unit/button.js b/test/unit/button.js
index 714bcc61..6541a6ea 100644
--- a/test/unit/button.js
+++ b/test/unit/button.js
@@ -8,8 +8,7 @@
     it('Should be upgraded to a MaterialButton successfully', function () {
       var el = document.createElement('button');
       componentHandler.upgradeElement(el, 'MaterialButton');
-      var upgraded = el.getAttribute('data-upgraded');
-      expect(upgraded).to.contain('MaterialButton');
+      expect($(el)).to.have.data('upgraded', ',MaterialButton');
     });
 
     it('Should be upgraded to a raised MaterialButton button with ripples successfully', function () {
@@ -17,8 +16,8 @@
       el.innerHTML = '<button class="wsk-button wsk-js-button wsk-button--raised wsk-js-ripple-effect">Raised</button>';
       var btn = el.firstChild;
       componentHandler.upgradeElement(btn, 'MaterialButton');
-      expect(btn.childNodes[1].className).to.contain('wsk-button__ripple-container');
-      expect(btn.childNodes[1].firstChild.className).to.contain('wsk-ripple');
+      expect($(btn.childNodes[1])).to.have.class('wsk-button__ripple-container');
+      expect($(btn.childNodes[1].firstChild)).to.have.class('wsk-ripple');
     });
 
     it('Should be upgraded to a MaterialButton FAB with ripples successfully', function () {
@@ -26,7 +25,7 @@
       el.innerHTML = '<button class="wsk-button wsk-js-button wsk-button--fab wsk-button--colored wsk-js-ripple-effect">♥</button>';
       var btn = el.firstChild;
       componentHandler.upgradeElement(btn, 'MaterialButton');
-      expect(btn.childNodes[1].className).to.contain('wsk-button__ripple-container');
-      expect(btn.childNodes[1].firstChild.className).to.contain('wsk-ripple');
+      expect($(btn.childNodes[1])).to.have.class('wsk-button__ripple-container');
+      expect($(btn.childNodes[1].firstChild)).to.have.class('wsk-ripple');
     });
   });