1 files changed, 34 insertions, 10 deletions
diff --git a/test/codehighlighting.js b/test/codehighlighting.js
index a64bf4b..b0c4110 100644
--- a/test/codehighlighting.js
+++ b/test/codehighlighting.js
@@ -7,27 +7,51 @@ var Plugin = require('../lib/plugin');
 var PLUGINS_ROOT = path.resolve(__dirname, 'plugins');
 
 describe('Code Highlighting', function () {
-    it('should correctly replace highlighting', function() {
+    var book, PAGE;
+
+    before(function() {
         return books.generate('highlight', 'website', {
-            prepare: function(book) {
+            prepare: function(_book) {
+                book = _book;
+
                 var plugin = new Plugin(book, "highlight");
                 plugin.load("./highlight", PLUGINS_ROOT);
 
                 book.plugins.load(plugin);
             }
         })
-        .then(function(book) {
-            var PAGE = fs.readFileSync(
+        .then(function() {
+            PAGE = fs.readFileSync(
                 path.join(book.options.output, "index.html"),
                 { encoding: "utf-8" }
             );
+        });
+    });
 
-            PAGE.should.be.html({
-                "code": {
-                    count: 1,
-                    text: 'code_test\n_code'
-                }
-            });
+    it('should correctly replace highlighting', function() {
+        PAGE.should.be.html({
+            "code": {
+                index: 0,
+                text: 'code_test 1\n_code'
+            }
+        });
+    });
+
+    it('should correctly replace highlighting with language', function() {
+        PAGE.should.be.html({
+            "code": {
+                index: 1,
+                text: 'lang_test 2\n_lang'
+            }
+        });
+    });
+
+    it('should correctly replace highlighting for inline code', function() {
+        PAGE.should.be.html({
+            "code": {
+                index: 2,
+                text: 'code_test 3_code'
+            }
         });
     });
 });