From 2ca72964f27717de3afc6d022d977ed1418ddc04 Mon Sep 17 00:00:00 2001
From: David Gardiner <dmgardiner25@github.com>
Date: Fri, 7 Jul 2023 14:44:47 -0700
Subject: [PATCH] Allow optional `axios` adapter in
 `TunnelManagementHttpClient` (#271)

---
 .../management/tunnelManagementHttpClient.ts  |  6 ++-
 ts/test/tunnels-test/tunnelManagementTests.ts | 53 +++++++++++++++++--
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/ts/src/management/tunnelManagementHttpClient.ts b/ts/src/management/tunnelManagementHttpClient.ts
index 93577680..26b2f401 100644
--- a/ts/src/management/tunnelManagementHttpClient.ts
+++ b/ts/src/management/tunnelManagementHttpClient.ts
@@ -21,7 +21,7 @@ import {
 import { TunnelRequestOptions } from './tunnelRequestOptions';
 import { TunnelAccessTokenProperties } from './tunnelAccessTokenProperties';
 import { tunnelSdkUserAgent } from './version';
-import axios, { AxiosError, AxiosRequestConfig, AxiosResponse, Method } from 'axios';
+import axios, { AxiosAdapter, AxiosError, AxiosRequestConfig, AxiosResponse, Method } from 'axios';
 import * as https from 'https';
 
 type NullableIfNotBoolean<T> = T extends boolean ? T : T | null;
@@ -101,12 +101,14 @@ export class TunnelManagementHttpClient implements TunnelManagementClient {
      * the global tunnel service URI.
      * @param httpsAgent Optional agent that will be invoked for HTTPS requests to the tunnel
      * service.
+     * @param adapter Optional axios adapter to use for HTTP requests.
      */
     public constructor(
         userAgents: (ProductHeaderValue | string)[] | ProductHeaderValue | string,
         userTokenCallback?: () => Promise<string | null>,
         tunnelServiceUri?: string,
         public readonly httpsAgent?: https.Agent,
+        private readonly adapter?: AxiosAdapter
     ) {
         if (!userAgents) {
             throw new TypeError('User agent must be provided.');
@@ -535,6 +537,7 @@ export class TunnelManagementHttpClient implements TunnelManagementClient {
         );
         const config: AxiosRequestConfig = {
             httpsAgent: this.httpsAgent,
+            adapter: this.adapter,
         };
         return await this.request<boolean>('GET', uri, undefined, config);
     }
@@ -697,6 +700,7 @@ export class TunnelManagementHttpClient implements TunnelManagementClient {
         const config: AxiosRequestConfig = {
             headers,
             ...(this.httpsAgent && { httpsAgent: this.httpsAgent }),
+            ...(this.adapter && { adapter: this.adapter }),
         };
 
         if (options?.followRedirects === false) {
diff --git a/ts/test/tunnels-test/tunnelManagementTests.ts b/ts/test/tunnels-test/tunnelManagementTests.ts
index e642fe8b..beb547be 100644
--- a/ts/test/tunnels-test/tunnelManagementTests.ts
+++ b/ts/test/tunnels-test/tunnelManagementTests.ts
@@ -2,16 +2,17 @@
 // Licensed under the MIT license.
 
 import * as assert from 'assert';
+import axios, { AxiosPromise, AxiosRequestConfig, Method } from 'axios';
+import * as https from 'https';
 import { suite, test, slow, timeout } from '@testdeck/mocha';
 import { TunnelManagementHttpClient } from '@microsoft/dev-tunnels-management';
-import { AxiosRequestConfig, Method } from 'axios';
 
 @suite
 @slow(3000)
 @timeout(10000)
 export class TunnelManagementTests {
 
-    private readonly managementClient : TunnelManagementHttpClient;
+    private readonly managementClient: TunnelManagementHttpClient;
 
     public constructor() {
         this.managementClient = new TunnelManagementHttpClient(
@@ -47,7 +48,6 @@ export class TunnelManagementTests {
         assert(!this.lastRequest.uri.includes('global=true'));
     }
 
-
     @test
     public async listTunnelsGlobal() {
         this.nextResponse = [];
@@ -74,4 +74,51 @@ export class TunnelManagementTests {
         assert.equal(this.lastRequest.method, 'GET');
         assert(this.lastRequest.uri.endsWith('/api/v1/userlimits'));
     }
+
+    @test
+    public async configDoesNotContainHttpsAgentAndAdapter() {
+        this.nextResponse = [];
+        await this.managementClient.listUserLimits();
+        assert(this.lastRequest);
+        assert(this.lastRequest.config.httpsAgent === undefined);
+        assert(this.lastRequest.config.adapter === undefined);
+    }
+
+    @test
+    public async configContainsHttpsAgentAndAdapter() {
+        // Create a mock https agent
+        const httpsAgent = new https.Agent({
+            rejectUnauthorized: true,
+            keepAlive: true,
+        });
+
+        // Create a mock axios adapter
+        interface AxiosAdapter {
+            (config: AxiosRequestConfig): AxiosPromise<any>;
+        }
+
+        class AxiosAdapter implements AxiosAdapter {
+            constructor(private client: any, private auth: any) { }
+        }
+
+        const axiosAdapter = new AxiosAdapter(axios, { auth: { username: 'test', password: 'test' } });
+
+        // Create a management client with a mock https agent and adapter
+        const managementClient = new TunnelManagementHttpClient(
+            'test/0.0.0', undefined, 'http://global.tunnels.test.api.visualstudio.com', httpsAgent, axiosAdapter);
+        (<any>managementClient).request = this.mockRequest.bind(this);
+
+        this.nextResponse = [];
+        await managementClient.listUserLimits();
+        assert(this.lastRequest);
+
+        // Assert that the https agent and adapter are the same as the ones we passed into the constructor
+        assert(this.lastRequest.config.httpsAgent === httpsAgent);
+        assert(this.lastRequest.config.httpsAgent !== new https.Agent({
+            rejectUnauthorized: true,
+            keepAlive: true,
+        }));
+        assert(this.lastRequest.config.adapter === axiosAdapter);
+        assert(this.lastRequest.config.adapter !== new AxiosAdapter(axios, { auth: { username: 'test', password: 'test' } }))
+    }
 }