Fix #2 and modify test

index.js

@@ -41,16 +41,17 @@ CssShortener.prototype.cssStream = CssShortener.prototype.stream = function(call
 }
 CssShortener.prototype.htmlStream = function(callback) {
   const t = this;
-  return replaceStream(HTML_CLASS_REGEX, function(match) {
-    const classes = match.trim().split(' ');
+  return replaceStream(HTML_CLASS_REGEX, function(match, capturingGroup) {
+    if (!capturingGroup) return match;
+    const classes = capturingGroup.trim().split(' ');
     const classCount = classes.length;
     var result = '';
     for (var i = 0; i < classCount; i++) {
       // Check if class is mapped and add it to the result
       result += (t._classNameMap[classes[i]] != null ? t._classNameMap[classes[i]] : classes[i]);
       if (i < classCount - 1) result += ' ';
     }
-    return result;
+    return `class="${result}"`;
   });
 }
 

test/index.js

@@ -173,10 +173,19 @@ describe('CssShortener', function() {
         'test-class': 'a'
       };
       c.importMap(map);
+
       const stream = str('<div class="abc test-class 15a"></div>').pipe(c.htmlStream());
       toString(stream).then(function(msg) {
         msg.should.equal('<div class="abc a 15a"></div>');
-        done();
+        const stream = str('<div class="test-class"></div>').pipe(c.htmlStream());
+        toString(stream).then(function(msg) {
+          msg.should.equal('<div class="a"></div>');
+          const stream = str('<div class="myclass"></div>').pipe(c.htmlStream());
+          toString(stream).then(function(msg) {
+            msg.should.equal('<div class="myclass"></div>');
+            done();
+          });
+        });
       });
     });
   });