From 3f791694da2bbc1fc0b42b05a8be261fe9bbebf9 Mon Sep 17 00:00:00 2001
From: Daniel Imms <2193314+Tyriar@users.noreply.github.com>
Date: Tue, 5 Sep 2023 10:28:45 -0700
Subject: [PATCH] Return an optional DisposableStore from leak fn

Part of #190503
---
 src/vs/base/test/common/utils.ts | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/vs/base/test/common/utils.ts b/src/vs/base/test/common/utils.ts
index 62508dcf7f1..03641b0c23a 100644
--- a/src/vs/base/test/common/utils.ts
+++ b/src/vs/base/test/common/utils.ts
@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { IDisposable, IDisposableTracker, setDisposableTracker } from 'vs/base/common/lifecycle';
+import { DisposableStore, IDisposable, IDisposableTracker, setDisposableTracker } from 'vs/base/common/lifecycle';
 import { join } from 'vs/base/common/path';
 import { isWindows } from 'vs/base/common/platform';
 import { URI } from 'vs/base/common/uri';
@@ -128,21 +128,34 @@ export class DisposableTracker implements IDisposableTracker {
  *
  * Use `markAsSingleton` if disposable singletons are created lazily that are allowed to outlive the test.
  * Make sure that the singleton properly registers all child disposables so that they are excluded too.
+ *
+ * @returns A {@link DisposableStore} that can optionally be used to track disposables in the test.
+ * This will be automatically disposed on test teardown.
 */
-export function ensureNoDisposablesAreLeakedInTestSuite() {
+export function ensureNoDisposablesAreLeakedInTestSuite(): Pick<DisposableStore, 'add'> {
 	let tracker: DisposableTracker | undefined;
+	let store: DisposableStore;
 	setup(() => {
+		store = new DisposableStore();
 		tracker = new DisposableTracker();
 		setDisposableTracker(tracker);
 	});
 
 	teardown(function (this: import('mocha').Context) {
+		store.dispose();
 		setDisposableTracker(null);
-
 		if (this.currentTest?.state !== 'failed') {
 			tracker!.ensureNoLeakingDisposables();
 		}
 	});
+
+	// Wrap store as the suite function is called before it's initialized
+	const testContext = {
+		add<T extends IDisposable>(o: T): T {
+			return store.add(o);
+		}
+	};
+	return testContext;
 }
 
 export function throwIfDisposablesAreLeaked(body: () => void): void {