3 files changed, 10 insertions, 6 deletions

diff --git a/src/test/CppTests/Config.test.ts b/src/test/CppTests/Config.test.ts
index 8a8e939..0459dcb 100644
--- a/src/test/CppTests/Config.test.ts
+++ b/src/test/CppTests/Config.test.ts
@@ -83,12 +83,13 @@ suite("C++ - Configuration Tests", () => {
 
     test("Lines to get test", () => {
         testSetup.cfg = new Config();
-        testSetup.cfg.linesToGet = 1;
-        const negativeResult = testSetup.SetLines(["template<typename T> bool \n", "foo(T a);"]).GetResult();
-        assert.equal("/**\n * @brief \n * \n */", negativeResult);
         testSetup.cfg.linesToGet = 2;
         const positiveResult = testSetup.SetLines(["template<typename T> bool \n", "foo(T a);"]).GetResult();
         assert.equal("/**\n * @brief \n * \n * @tparam T \n * @param a \n * "
             + "@return true \n * @return false \n */", positiveResult);
+        testSetup.cfg.linesToGet = 0;
+        testSetup.firstLine = 1;
+        const negativeResult = testSetup.SetLines(["template<typename T> bool \n", "foo(T a);"]).GetResult();
+        assert.equal("/**\n * @brief \n * \n */", negativeResult);
     });
 });
diff --git a/src/test/CppTests/TestSetup.ts b/src/test/CppTests/TestSetup.ts
index ee7256d..2e6c7f1 100644
--- a/src/test/CppTests/TestSetup.ts
+++ b/src/test/CppTests/TestSetup.ts
@@ -13,11 +13,12 @@ import MockSelection from "../tools/MockSelection";
 
 export default class TestSetup {
     public cfg: Config;
-
+    public firstLine: number;
     private editor: MockEditor;
 
     constructor(method: string) {
         this.cfg = new Config();
+        this.firstLine = 0;
         this.SetLine(method);
     }
 
@@ -38,7 +39,7 @@ export default class TestSetup {
                 .map((l) => new MockLine(l));
         }
 
-        const selection: MockSelection =  new MockSelection(new MockPosition(0, 0));
+        const selection: MockSelection =  new MockSelection(new MockPosition(this.firstLine, 0));
         const doc: MockDocument = new MockDocument(mockLines);
 
         this.editor = new MockEditor(selection, doc);
@@ -51,7 +52,8 @@ export default class TestSetup {
         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)));
+        // tslint:disable-next-line:max-line-length
+        gen.GenerateDoc(new vscode.Range(new vscode.Position(this.firstLine, 0), new vscode.Position(this.firstLine, 0)));
 
         return this.editor.editBuilder.text;
     }
diff --git a/src/test/tools/MockDocument.ts b/src/test/tools/MockDocument.ts
index 781dbfe..1dca76e 100644
--- a/src/test/tools/MockDocument.ts
+++ b/src/test/tools/MockDocument.ts
@@ -14,6 +14,7 @@ export default class MockDocument implements vscode.TextDocument {
     private lines: vscode.TextLine[];
     public constructor(lines: vscode.TextLine[]) {
         this.lines = lines;
+        this.fileName = "MockDocument.h";
     }
     public save(): Thenable<boolean> {
         throw new Error("Method not implemented.");