From cc661dafd06770137459b72441e5f7cc877483f0 Mon Sep 17 00:00:00 2001
From: Ramya Achutha Rao <ramyar@microsoft.com>
Date: Tue, 2 May 2017 17:27:29 -0700
Subject: [PATCH] Fixes #601 Toggle test coverage

---
 CHANGELOG.md   |  5 +++--
 package.json   |  2 +-
 src/goCover.ts | 10 +++++++++-
 src/goMain.ts  |  4 ++--
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5cb28a27b..aac7230ff 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,8 +10,9 @@
 ## Test Coverage
 * [Thomas Bradford (@kode4food)](https://github.com/kode4food)
     * New setting `go.coverageOptions` to control whether you want to highlight only covered code or only uncovered code or both when code coverage is run.
-
-
+* [Ramya Rao (@ramya-rao-a)](https://github.com/ramya-rao-a)
+    * The command `Go: Test Coverage In Current Package` is renamed to `Go: Toggle Test Coverage In Current Package` and it does exactly what the name suggests. Toggles test coverage.
+        
 ### Bug Fixes 
 * [Ramya Rao (@ramya-rao-a)](https://github.com/ramya-rao-a)
     * Fix for [Bug 529](https://github.com/Microsoft/vscode-go/issues/529) Code completion for unimported packages now works on unsaved file after deleting imports.
diff --git a/package.json b/package.json
index 19723de4f..4ca16dc90 100644
--- a/package.json
+++ b/package.json
@@ -111,7 +111,7 @@
       },
       {
         "command": "go.test.coverage",
-        "title": "Go: Test Coverage In Current Package",
+        "title": "Go: Toggle Test Coverage In Current Package",
         "description": "Displays test coverage in the current package."
       },
       {
diff --git a/src/goCover.ts b/src/goCover.ts
index 4d3fb785e..b4761eb09 100644
--- a/src/goCover.ts
+++ b/src/goCover.ts
@@ -53,13 +53,21 @@ export function removeCodeCoverage(e: vscode.TextDocumentChangeEvent) {
 	}
 }
 
-export function coverageCurrentPackage() {
+export function toggleCoverageCurrentPackage() {
 	let editor = vscode.window.activeTextEditor;
 	if (!editor) {
 		vscode.window.showInformationMessage('No editor is active.');
 		return;
 	}
 
+	// If current file has highlights, then remove coverage, else add coverage
+	for (let filename in coverageFiles) {
+		if (editor.document.uri.fsPath.endsWith(filename)) {			
+			clearCoverage();
+			return;
+		}
+	}
+
 	// FIXME: is there a better way to get goConfig?
 	let goConfig = vscode.workspace.getConfiguration('go');
 	let cwd = path.dirname(editor.document.uri.fsPath);
diff --git a/src/goMain.ts b/src/goMain.ts
index d9d9c054f..a5de55b45 100644
--- a/src/goMain.ts
+++ b/src/goMain.ts
@@ -21,7 +21,7 @@ import { check, ICheckResult, removeTestStatus } from './goCheck';
 import { updateGoPathGoRootFromConfig, offerToInstallTools } from './goInstallTools';
 import { GO_MODE } from './goMode';
 import { showHideStatus } from './goStatus';
-import { coverageCurrentPackage, getCodeCoverage, removeCodeCoverage } from './goCover';
+import { toggleCoverageCurrentPackage, getCodeCoverage, removeCodeCoverage } from './goCover';
 import { testAtCursor, testCurrentPackage, testCurrentFile, testPrevious, showTestOutput, testWorkspace } from './goTest';
 import * as goGenerateTests from './goGenerateTests';
 import { addImport } from './goImport';
@@ -141,7 +141,7 @@ export function activate(ctx: vscode.ExtensionContext): void {
 	}));
 
 	ctx.subscriptions.push(vscode.commands.registerCommand('go.test.coverage', () => {
-		coverageCurrentPackage();
+		toggleCoverageCurrentPackage();
 	}));
 
 	ctx.subscriptions.push(vscode.commands.registerCommand('go.test.showOutput', () => {