Package suggestion for _test.go

src/goSuggest.ts

@@ -12,6 +12,7 @@ import { getBinPath, parameters, parseFilePrelude, isPositionInString, goKeyword
 import { promptForMissingTool } from './goInstallTools';
 import { getTextEditForAddImport } from './goImport';
 import { getImportablePackages } from './goPackages';
+import { readDir } from './util';
 
 function vscodeKindFromGoCodeClass(kind: string): vscode.CompletionItemKind {
 	switch (kind) {
@@ -29,6 +30,69 @@ function vscodeKindFromGoCodeClass(kind: string): vscode.CompletionItemKind {
 	return vscode.CompletionItemKind.Property; // TODO@EG additional mappings needed?
 }
 
+function defaultPackageName(filename: string): string {
+	let goFilename = basename(filename);
+	if (goFilename === 'main.go') {
+		return 'main';
+	}
+
+	if (goFilename.endsWith('internal_test.go')) {
+		return basename(dirname(filename));
+	}
+
+	if (goFilename.endsWith('_test.go')) {
+		return basename(dirname(filename)) + '_test';
+	}
+
+	return basename(dirname(filename));
+}
+
+function packageNameSuggestion(filename: string): Promise<string> {
+	return new Promise((resolve, reject) => {
+		const goFilename = basename(filename);
+		if (goFilename === 'main.go') {
+			return resolve('main');
+		}
+
+		if (goFilename.endsWith('internal_test.go')) {
+			return resolve(guessPackageName(basename(dirname(filename))));
+		}
+
+		if (goFilename.endsWith('_test.go')) {
+			return resolve(guessPackageName(basename(dirname(filename))) + '_test');
+		}
+
+		readDir(dirname(filename)).then(files => {
+			if (files.indexOf('main.go') > -1) {
+				return resolve('main');
+			}
+
+			resolve(guessPackageName(basename(dirname(filename))));
+		}, err => reject(err));
+	});
+}
+
+/**
+ * Guess the package name based on directory name.
+ *
+ * Cases:
+ * - dir 'go-i18n' -> 'i18n'
+ * - dir 'go-spew' -> 'spew'
+ * - dir 'kingpin' -> 'kingpin'
+ * - dir 'go-expand-tilde' -> 'tilde'
+ * - dir 'gax-go' -> 'gax'
+ * - dir 'go-difflib' -> 'difflib'
+ * - dir 'jwt-go' -> 'jwt'
+ * - dir 'go-radix' -> 'radix'
+ *
+ * @param {string} dirName where the go file located.
+ */
+function guessPackageName(dirName) {
+	let segments = dirName.split(/[\.-]/);
+	segments = segments.filter(val => val !== 'go');
+	return segments[segments.length - 1];
+}
+
 interface GoCodeSuggestion {
 	class: string;
 	name: string;
@@ -148,21 +212,6 @@ export class GoCompletionItemProvider implements vscode.CompletionItemProvider {
 					let suggestions = [];
 					let suggestionSet = new Set<string>();
 
-					// 'Smart Snippet' for package clause
-					// TODO: Factor this out into a general mechanism
-					if (!inputText.match(/package\s+(\w+)/)) {
-						let defaultPackageName =
-							basename(filename) === 'main.go'
-								? 'main'
-								: basename(dirname(filename));
-						if (defaultPackageName.match(/[a-zA-Z_]\w*/)) {
-							let packageItem = new vscode.CompletionItem('package ' + defaultPackageName);
-							packageItem.kind = vscode.CompletionItemKind.Snippet;
-							packageItem.insertText = 'package ' + defaultPackageName + '\r\n\r\n';
-							suggestions.push(packageItem);
-						}
-
-					}
 					if (results[1]) {
 						for (let suggest of results[1]) {
 							if (inString && suggest.class !== 'import') continue;
@@ -204,6 +253,20 @@ export class GoCompletionItemProvider implements vscode.CompletionItemProvider {
 					let importablePkgs = includeUnimportedPkgs ? this.getMatchingPackages(currentWord, suggestionSet) : [];
 					suggestions = suggestions.concat(importablePkgs);
 
+					// 'Smart Snippet' for package clause
+					// TODO: Factor this out into a general mechanism
+					if (!inputText.match(/package\s+(\w+)/)) {
+						return packageNameSuggestion(filename).then(pkgName => {
+							if (pkgName.match(/[a-zA-Z_]\w*/)) {
+								let packageItem = new vscode.CompletionItem('package ' + pkgName);
+								packageItem.kind = vscode.CompletionItemKind.Snippet;
+								packageItem.insertText = 'package ' + pkgName + '\r\n\r\n';
+								suggestions.push(packageItem);
+							}
+							resolve(suggestions);
+						});
+					}
+
 					resolve(suggestions);
 				} catch (e) {
 					reject(e);

src/util.ts

@@ -428,3 +428,14 @@ export function getImportPath(text: string): string {
 
 	return '';
 }
+
+export function readDir(path): Promise<string[]> {
+	return new Promise<string[]>((resolve, reject) => {
+		fs.readdir(path, (err, files) => {
+			if (err) {
+				return reject(err);
+			}
+			resolve(files);
+		});
+	});
+}