1 files changed, 5 insertions, 10 deletions
diff --git a/src/test/tools/MockDocument.ts b/src/test/tools/MockDocument.ts
index cb6bbfe..781dbfe 100644
--- a/src/test/tools/MockDocument.ts
+++ b/src/test/tools/MockDocument.ts
@@ -11,24 +11,19 @@ export default class MockDocument implements vscode.TextDocument {
     public isClosed: boolean;
     public eol: vscode.EndOfLine;
     public lineCount: number;
-    private line: vscode.TextLine;
-    private callCount: number;
-    public constructor(line: vscode.TextLine) {
-        this.line = line;
-        this.callCount = 0;
+    private lines: vscode.TextLine[];
+    public constructor(lines: vscode.TextLine[]) {
+        this.lines = lines;
     }
     public save(): Thenable<boolean> {
         throw new Error("Method not implemented.");
     }
     public lineAt(line: number | vscode.Position): vscode.TextLine;
     public lineAt(position: any): any {
-        if (++this.callCount === 1) {
-            return new MockLine("");
-        } else if (this.callCount === 2) {
-            return this.line;
-        } else {
+        if (position >= this.lines.length) {
             return new MockLine(";");
         }
+        return this.lines[position];
     }
     public offsetAt(position: vscode.Position): number {
         throw new Error("Method not implemented.");