1 files changed, 16 insertions, 10 deletions

diff --git a/src/test/CppTests/TestSetup.ts b/src/test/CppTests/TestSetup.ts
index 068cff2..ee7256d 100644
--- a/src/test/CppTests/TestSetup.ts
+++ b/src/test/CppTests/TestSetup.ts
@@ -17,23 +17,29 @@ export default class TestSetup {
     private editor: MockEditor;
 
     constructor(method: string) {
+        this.cfg = new Config();
         this.SetLine(method);
     }
 
     public SetLine(method: string): TestSetup {
-        this.cfg = new Config();
+        return this.SetLines([method]);
+    }
 
-        let position: MockPosition;
-        position = new MockPosition(0, 0);
+    public SetLines(lines: string[]): TestSetup {
+        let mockLines: MockLine[];
 
-        let selection: MockSelection;
-        selection = new MockSelection(position);
+        // Add two empty lines with just the indentation
+        // since one is for the triggersequence line.
+        // And the other is if the next empty line.
+        if (lines.length > 0) {
+            const indent: string = lines[0].match("^\\s*")[0];
 
-        let line: MockLine;
-        line = new MockLine(method);
+            mockLines = [indent, indent].concat(lines)
+                .map((l) => new MockLine(l));
+        }
 
-        let doc: MockDocument;
-        doc = new MockDocument(line);
+        const selection: MockSelection =  new MockSelection(new MockPosition(0, 0));
+        const doc: MockDocument = new MockDocument(mockLines);
 
         this.editor = new MockEditor(selection, doc);
 
@@ -42,7 +48,7 @@ export default class TestSetup {
 
     public GetResult(): string {
         let parser: CodeParser;
-        parser = new CppParser(new Config());
+        parser = new CppParser(this.cfg);
 
         const gen: IDocGen = parser.Parse(this.editor);
         gen.GenerateDoc(new vscode.Range(new vscode.Position(0, 0), new vscode.Position(0, 0)));