4 files changed, 40 insertions, 37 deletions
diff --git a/packages/gitbook-html/package.json b/packages/gitbook-html/package.json
index e0a4214..1552786 100644
--- a/packages/gitbook-html/package.json
+++ b/packages/gitbook-html/package.json
@@ -8,7 +8,7 @@
     "cheerio": "^0.20.0 && >=0.20.0"
   },
   "scripts": {
-    "test": "mocha --reporter spec --bail",
+    "test": "mocha --reporter spec --bail --compilers js:babel-register",
     "prepublish": "rm -rf lib/ && babel -d lib/ src/ --source-maps"
   },
   "repository": {
diff --git a/packages/gitbook-markdown/package.json b/packages/gitbook-markdown/package.json
index 9cb7916..f8a528c 100644
--- a/packages/gitbook-markdown/package.json
+++ b/packages/gitbook-markdown/package.json
@@ -11,7 +11,7 @@
     "slate": "^0.16.8"
   },
   "scripts": {
-    "test": "mocha --reporter list --bail"
+    "test": "mocha --reporter list --bail --compilers js:babel-register"
   },
   "repository": {
     "type": "git",
diff --git a/packages/gitbook-markdown/src/page.js b/packages/gitbook-markdown/src/page.js
index 5bd20e2..b899d5b 100644
--- a/packages/gitbook-markdown/src/page.js
+++ b/packages/gitbook-markdown/src/page.js
@@ -69,28 +69,52 @@ function escapeTextNode(node) {
     return nodes;
 }
 
-
 /**
  * Annotate a block container.
  * @param  {Node} parent
- * @return {Node}
+ * @param  {Number} levelRaw
+ * @return {Node} node
+ * @return {Number} levelRaw
  */
 function annotateNode(parent, levelRaw) {
     let { nodes } = parent;
 
     nodes = nodes.reduce((out, node) => {
-        if (node.type === BLOCKS.CODE) {
-            return out.concat(escapeCodeBlock(node));
+        if (node.type === INLINES.TEMPLATE) {
+            const { type, text } = node.data.toJS();
+
+            if (type === 'expr') {
+                if (text === 'raw') {
+                    levelRaw = levelRaw + 1;
+                } else if (text == 'endraw') {
+                    levelRaw = 0;
+                }
+            }
+
+            return out.concat([ node ]);
+        }
+
+        else if (node.type === BLOCKS.CODE) {
+            return out.concat(
+                levelRaw == 0 ? escapeCodeBlock(node) : [ node ]
+            );
         }
 
         else if (node.kind == 'text') {
-            return out.concat(escapeTextNode(node));
+            return out.concat(
+                levelRaw == 0 ? escapeTextNode(node) : [ node ]
+            );
         }
 
-        return out.concat([node]);
+        const result = annotateNode(node, levelRaw);
+        levelRaw = result.levelRaw;
+        return out.concat([result.node]);
     }, []);
 
-    return parent.merge({ nodes });
+    return {
+        levelRaw,
+        node: parent.merge({ nodes })
+    };
 }
 
 /**
@@ -101,30 +125,9 @@ function annotateNode(parent, levelRaw) {
  * @return {String}
  */
 function preparePage(src) {
-    let levelRaw = 0;
-
     const state = State.create(markdown);
     let document = state.deserializeToDocument(src);
-
-    document = document.mapDescendants((node) => {
-        if (node.type === INLINES.TEMPLATE) {
-            const { type, text } = node.data;
-
-            if (type === 'expr') {
-                if (text === 'raw') {
-                    levelRaw = levelRaw + 1;
-                } else if (text == 'endraw') {
-                    levelRaw = 0;
-                }
-            }
-
-            return node;
-        } else if (node.kind !== 'text') {
-            return annotateNode(node, levelRaw);
-        } else {
-            return node;
-        }
-    });
+    document = annotateNode(document, 0).node;
 
     return state.serializeDocument(document);
 }
diff --git a/packages/gitbook-markdown/test/page.js b/packages/gitbook-markdown/test/page.js
index 9ea48f4..649827e 100644
--- a/packages/gitbook-markdown/test/page.js
+++ b/packages/gitbook-markdown/test/page.js
@@ -25,26 +25,26 @@ describe('Page', () => {
         expect(page.prepare('Hello `world`')).toBe('Hello {% raw %}`world`{% endraw %}\n\n');
         expect(page.prepare('Hello `world test`')).toBe('Hello {% raw %}`world test`{% endraw %}\n\n');
         expect(page.prepare('Hello ```world test```')).toBe('Hello {% raw %}`world test`{% endraw %}\n\n');
-        expect(page.prepare('Hello\n```js\nworld test\n```\n')).toBe('Hello\n\n{% raw %}```js\nworld test\n```\n\n{% endraw %}');
-        expect(page.prepare('Hello\n```\ntest\n\tworld\n\ttest\n```')).toBe('Hello\n\n{% raw %}```\ntest\n    world\n    test\n```\n\n{% endraw %}');
+        expect(page.prepare('Hello\n```js\nworld test\n```\n')).toBe('Hello\n\n{% raw %}\n```js\nworld test\n```\n\n{% endraw %}\n');
+        expect(page.prepare('Hello\n```\ntest\n\tworld\n\ttest\n```')).toBe('Hello\n\n{% raw %}\n```\ntest\n    world\n    test\n```\n\n{% endraw %}\n');
     });
 
     it('should escape codeblocks in preparation (2)', () => {
         expect(
             page.prepare('Hello\n\n\n\tworld\n\thello\n\n\ntest')
         ).toBe(
-            'Hello\n\n{% raw %}```\nworld\nhello```\n\n{% endraw %}test\n\n'
+            'Hello\n\n{% raw %}\n```\nworld\nhello\n```\n\n{% endraw %}\ntest\n\n'
         );
         expect(
             page.prepare('Hello\n\n\n\tworld\n\thello\n\n\n')
         ).toBe(
-            'Hello\n\n{% raw %}```\nworld\nhello```\n\n{% endraw %}'
+            'Hello\n\n{% raw %}\n```\nworld\nhello\n```\n\n{% endraw %}\n'
         );
     });
 
     it('should escape codeblocks with nunjucks tags', () => {
         expect(
-            page.prepare('Hello {{ "Bonjour" }} ```test```')
+            page.prepare('Hello {{ "Bonjour" }} {% raw %}```test```{% endraw %}')
         ).toBe(
             'Hello {{ "Bonjour" }} {% raw %}`test`{% endraw %}\n\n'
         );
@@ -70,7 +70,7 @@ describe('Page', () => {
         expect(
             page.prepare('```\ntest\n```\n\n\n### Test')
         ).toBe(
-            '{% raw %}```\ntest\n```\n\n{% endraw %}### Test\n\n'
+            '{% raw %}\n```\ntest\n```\n\n{% endraw %}\n### Test\n\n'
         );
     });