Use a TestAuthUsageService for integration tests and clean up disposables a bit (microsoft#251936)

* Use a `TestAuthUsageService` for integration tests and clean up disposables a bit

The test service is needed because otherwise it creates an event that create a disposable... and the test itself doesn't clean that up properly.

That service isn't being tested here, so it's replaced with a mocked one.

* add a similar class for the mainthread tests

Author: TylerLeonhardt

src/vs/workbench/api/test/browser/extHostAuthentication.integrationTest.ts

@@ -28,7 +28,7 @@ import type { AuthenticationProvider, AuthenticationSession } from 'vscode';
 import { IBrowserWorkbenchEnvironmentService } from '../../../services/environment/browser/environmentService.js';
 import { IProductService } from '../../../../platform/product/common/productService.js';
 import { AuthenticationAccessService, IAuthenticationAccessService } from '../../../services/authentication/browser/authenticationAccessService.js';
-import { AuthenticationUsageService, IAuthenticationUsageService } from '../../../services/authentication/browser/authenticationUsageService.js';
+import { IAccountUsage, IAuthenticationUsageService } from '../../../services/authentication/browser/authenticationUsageService.js';
 import { AuthenticationExtensionsService } from '../../../services/authentication/browser/authenticationExtensionsService.js';
 import { ILogService, NullLogService } from '../../../../platform/log/common/log.js';
 import { IExtHostInitDataService } from '../../common/extHostInitDataService.js';
@@ -76,6 +76,15 @@ class AuthTestQuickInputService extends TestQuickInputService {
 	}
 }
 
+class TestAuthUsageService implements IAuthenticationUsageService {
+	_serviceBrand: undefined;
+	initializeExtensionUsageCache(): Promise<void> { return Promise.resolve(); }
+	extensionUsesAuth(extensionId: string): Promise<boolean> { return Promise.resolve(false); }
+	readAccountUsages(providerId: string, accountName: string): IAccountUsage[] { return []; }
+	removeAccountUsage(providerId: string, accountName: string): void { }
+	addAccountUsage(providerId: string, accountName: string, scopes: ReadonlyArray<string>, extensionId: string, extensionName: string): void { }
+}
+
 class TestAuthProvider implements AuthenticationProvider {
 	private id = 1;
 	private sessions = new Map<string, AuthenticationSession>();
@@ -136,7 +145,7 @@ suite('ExtHostAuthentication', () => {
 		services.set(IUserActivityService, new SyncDescriptor(UserActivityService));
 		services.set(IAuthenticationAccessService, new SyncDescriptor(AuthenticationAccessService));
 		services.set(IAuthenticationService, new SyncDescriptor(AuthenticationService));
-		services.set(IAuthenticationUsageService, new SyncDescriptor(AuthenticationUsageService));
+		services.set(IAuthenticationUsageService, new SyncDescriptor(TestAuthUsageService));
 		services.set(IAuthenticationExtensionsService, new SyncDescriptor(AuthenticationExtensionsService));
 		mainInstantiationService = disposables.add(new TestInstantiationService(services, undefined, undefined, true));
 

src/vs/workbench/api/test/browser/mainThreadAuthentication.integrationTest.ts

@@ -26,7 +26,7 @@ import { TestActivityService, TestExtensionService, TestProductService, TestStor
 import { IBrowserWorkbenchEnvironmentService } from '../../../services/environment/browser/environmentService.js';
 import { IProductService } from '../../../../platform/product/common/productService.js';
 import { AuthenticationAccessService, IAuthenticationAccessService } from '../../../services/authentication/browser/authenticationAccessService.js';
-import { AuthenticationUsageService, IAuthenticationUsageService } from '../../../services/authentication/browser/authenticationUsageService.js';
+import { IAccountUsage, IAuthenticationUsageService } from '../../../services/authentication/browser/authenticationUsageService.js';
 import { AuthenticationExtensionsService } from '../../../services/authentication/browser/authenticationExtensionsService.js';
 import { ILogService, NullLogService } from '../../../../platform/log/common/log.js';
 import { IHostService } from '../../../services/host/browser/host.js';
@@ -40,6 +40,15 @@ import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/c
 import { ServiceCollection } from '../../../../platform/instantiation/common/serviceCollection.js';
 import { SyncDescriptor } from '../../../../platform/instantiation/common/descriptors.js';
 
+class TestAuthUsageService implements IAuthenticationUsageService {
+	_serviceBrand: undefined;
+	initializeExtensionUsageCache(): Promise<void> { return Promise.resolve(); }
+	extensionUsesAuth(extensionId: string): Promise<boolean> { return Promise.resolve(false); }
+	readAccountUsages(providerId: string, accountName: string): IAccountUsage[] { return []; }
+	removeAccountUsage(providerId: string, accountName: string): void { }
+	addAccountUsage(providerId: string, accountName: string, scopes: ReadonlyArray<string>, extensionId: string, extensionName: string): void { }
+}
+
 suite('MainThreadAuthentication', () => {
 	const disposables = ensureNoDisposablesAreLeakedInTestSuite();
 
@@ -64,7 +73,7 @@ suite('MainThreadAuthentication', () => {
 		services.set(IUserActivityService, new SyncDescriptor(UserActivityService));
 		services.set(IAuthenticationAccessService, new SyncDescriptor(AuthenticationAccessService));
 		services.set(IAuthenticationService, new SyncDescriptor(AuthenticationService));
-		services.set(IAuthenticationUsageService, new SyncDescriptor(AuthenticationUsageService));
+		services.set(IAuthenticationUsageService, new SyncDescriptor(TestAuthUsageService));
 		services.set(IAuthenticationExtensionsService, new SyncDescriptor(AuthenticationExtensionsService));
 		instantiationService = disposables.add(new TestInstantiationService(services, undefined, undefined, true));
 

src/vs/workbench/services/authentication/browser/authenticationService.ts

@@ -21,7 +21,8 @@ import { ExtensionsRegistry } from '../../extensions/common/extensionsRegistry.j
 import { match } from '../../../../base/common/glob.js';
 import { URI } from '../../../../base/common/uri.js';
 import { IAuthorizationProtectedResourceMetadata, IAuthorizationServerMetadata } from '../../../../base/common/oauth.js';
-import { raceTimeout } from '../../../../base/common/async.js';
+import { raceCancellation, raceTimeout } from '../../../../base/common/async.js';
+import { CancellationTokenSource } from '../../../../base/common/cancellation.js';
 
 export function getAuthenticationProviderActivationEvent(id: string): string { return `onAuthenticationRequest:${id}`; }
 
@@ -110,7 +111,7 @@ export class AuthenticationService extends Disposable implements IAuthentication
 
 	private readonly _delegates: IAuthenticationProviderHostDelegate[] = [];
 
-	private _isDisposable: boolean = false;
+	private _disposedSource = new CancellationTokenSource();
 
 	constructor(
 		@IExtensionService private readonly _extensionService: IExtensionService,
@@ -119,7 +120,7 @@ export class AuthenticationService extends Disposable implements IAuthentication
 		@ILogService private readonly _logService: ILogService
 	) {
 		super();
-		this._register(toDisposable(() => this._isDisposable = true));
+		this._register(toDisposable(() => this._disposedSource.dispose(true)));
 		this._register(authenticationAccessService.onDidChangeExtensionSessionAccess(e => {
 			// The access has changed, not the actual session itself but extensions depend on this event firing
 			// when they have gained access to an account so this fires that event.
@@ -276,7 +277,7 @@ export class AuthenticationService extends Disposable implements IAuthentication
 	}
 
 	async getSessions(id: string, scopes?: string[], options?: IAuthenticationGetSessionsOptions, activateImmediate: boolean = false): Promise<ReadonlyArray<AuthenticationSession>> {
-		if (this._isDisposable) {
+		if (this._disposedSource.token.isCancellationRequested) {
 			return [];
 		}
 
@@ -297,7 +298,7 @@ export class AuthenticationService extends Disposable implements IAuthentication
 	}
 
 	async createSession(id: string, scopes: string[], options?: IAuthenticationCreateSessionOptions): Promise<AuthenticationSession> {
-		if (this._isDisposable) {
+		if (this._disposedSource.token.isCancellationRequested) {
 			throw new Error('Authentication service is disposed.');
 		}
 
@@ -310,7 +311,7 @@ export class AuthenticationService extends Disposable implements IAuthentication
 	}
 
 	async removeSession(id: string, sessionId: string): Promise<void> {
-		if (this._isDisposable) {
+		if (this._disposedSource.token.isCancellationRequested) {
 			throw new Error('Authentication service is disposed.');
 		}
 
@@ -382,17 +383,23 @@ export class AuthenticationService extends Disposable implements IAuthentication
 		if (provider) {
 			return provider;
 		}
+		if (this._disposedSource.token.isCancellationRequested) {
+			throw new Error('Authentication service is disposed.');
+		}
 
-		const store = this._register(new DisposableStore());
+		const store = new DisposableStore();
 		try {
 			const result = await raceTimeout(
-				Event.toPromise(
-					Event.filter(
-						this.onDidRegisterAuthenticationProvider,
-						e => e.id === providerId,
+				raceCancellation(
+					Event.toPromise(
+						Event.filter(
+							this.onDidRegisterAuthenticationProvider,
+							e => e.id === providerId,
+							store
+						),
 						store
 					),
-					store
+					this._disposedSource.token
 				),
 				5000
 			);

src/vs/workbench/services/authentication/browser/authenticationUsageService.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Queue } from '../../../../base/common/async.js';
-import { Disposable } from '../../../../base/common/lifecycle.js';
+import { Disposable, toDisposable } from '../../../../base/common/lifecycle.js';
 import { InstantiationType, registerSingleton } from '../../../../platform/instantiation/common/extensions.js';
 import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js';
 import { ILogService } from '../../../../platform/log/common/log.js';
@@ -56,17 +56,19 @@ export interface IAuthenticationUsageService {
 export class AuthenticationUsageService extends Disposable implements IAuthenticationUsageService {
 	_serviceBrand: undefined;
 
-	private _queue = new Queue();
+	private _queue = this._register(new Queue());
 	private _extensionsUsingAuth = new Set<string>();
 
+	private _disposed = false;
+
 	constructor(
 		@IStorageService private readonly _storageService: IStorageService,
 		@IAuthenticationService private readonly _authenticationService: IAuthenticationService,
 		@ILogService private readonly _logService: ILogService,
 		@IProductService productService: IProductService,
 	) {
 		super();
-
+		this._register(toDisposable(() => this._disposed = true));
 		// If an extension is listed in `trustedExtensionAuthAccess` we should consider it as using auth
 		const trustedExtensionAuthAccess = productService.trustedExtensionAuthAccess;
 		if (Array.isArray(trustedExtensionAuthAccess)) {
@@ -143,6 +145,9 @@ export class AuthenticationUsageService extends Disposable implements IAuthentic
 	}
 
 	private async _addExtensionsToCache(providerId: string) {
+		if (this._disposed) {
+			return;
+		}
 		try {
 			const accounts = await this._authenticationService.getAccounts(providerId);
 			for (const account of accounts) {