From 56e824bc9e3a27a899513dd8adef92c1574e3ab1 Mon Sep 17 00:00:00 2001
From: Jeff Bean <bean@uber.com>
Date: Mon, 13 Aug 2018 17:26:09 -0700
Subject: [PATCH] Simplify the logic even further

---
 src/goGenerateTests.ts | 28 ++++++++--------------------
 1 file changed, 8 insertions(+), 20 deletions(-)

diff --git a/src/goGenerateTests.ts b/src/goGenerateTests.ts
index 4bdf11140..801f2c988 100644
--- a/src/goGenerateTests.ts
+++ b/src/goGenerateTests.ts
@@ -91,12 +91,13 @@ export function GenerateTests(genType: GenerationType): Thenable<boolean> {
 	}
 
 	let goConfig = getGoConfigObject(editor);
+	let documentPath = editor.document.uri.fsPath;
 
 	switch (genType) {
 		case GenerationType.Package:
-			return generateTestCurrentPackage(editor, goConfig);
+			return generateTests({ dir: path.dirname(documentPath), goConfig: goConfig });
 		case GenerationType.File:
-			return generateTestCurrentFile(editor, goConfig);
+			return generateTests({ dir: documentPath, goConfig: goConfig });
 		case GenerationType.Function:
 			return generateTestCurrentFunction(editor, goConfig);
 		default:
@@ -105,23 +106,9 @@ export function GenerateTests(genType: GenerationType): Thenable<boolean> {
 	}
 }
 
-function generateTestCurrentPackage(editor: vscode.TextEditor, goConfig: vscode.WorkspaceConfiguration): Thenable<boolean> {
-	let dir = path.dirname(editor.document.uri.fsPath);
-	const goGenerateTestsFlags: string[] = goConfig['genTestsFlags'] || [];
-	return generateTests({ dir: dir, genFlags: goGenerateTestsFlags });
-}
-
-function generateTestCurrentFile(editor: vscode.TextEditor, goConfig: vscode.WorkspaceConfiguration): Thenable<boolean> {
-	let file = editor.document.uri.fsPath;
-	const goGenerateTestsFlags: string[] = goConfig['genTestsFlags'] || [];
-	return generateTests({ dir: file, genFlags: goGenerateTestsFlags });
-}
-
 function generateTestCurrentFunction(editor: vscode.TextEditor, goConfig: vscode.WorkspaceConfiguration): Thenable<boolean> {
 	let file = editor.document.uri.fsPath;
 
-	const goGenerateTestsFlags: string[] = goConfig['genTestsFlags'] || [];
-
 	return getFunctions(editor.document).then(functions => {
 		let currentFunction: vscode.SymbolInformation;
 		for (let func of functions) {
@@ -139,7 +126,7 @@ function generateTestCurrentFunction(editor: vscode.TextEditor, goConfig: vscode
 		if (funcName.includes('.')) {
 			funcName = funcName.split('.')[1];
 		}
-		return generateTests({ dir: file, func: funcName , genFlags: goGenerateTestsFlags  });
+		return generateTests({ dir: file, func: funcName , goConfig: goConfig  });
 	});
 }
 
@@ -152,9 +139,9 @@ interface Config {
 	 */
 	dir: string;
 	/**
-	 * Optional Template dir for any custom templates for `gotests`.
+	 * The config object for the go section
 	 */
-	genFlags: string[];
+	goConfig: vscode.WorkspaceConfiguration;
 	/**
 	 * Specific function names to generate tests squeleton.
 	 */
@@ -165,8 +152,9 @@ function generateTests(conf: Config): Thenable<boolean> {
 	return new Promise<boolean>((resolve, reject) => {
 		let cmd = getBinPath('gotests');
 		let args = ['-w'];
+		let goGenerateTestsFlags: string[] = conf.goConfig['genTestsFlags'] || [];
 
-		conf.genFlags.forEach(flag => {
+		goGenerateTestsFlags.forEach(flag => {
 			args.push(flag);
 		});