From ead068baae553ff2906e9300c23c694338b499dc Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Mon, 18 Nov 2024 09:29:45 -0500
Subject: [PATCH] add tests

---
 packages/core/src/shared/extensions/ssh.ts    |  8 +--
 .../src/test/shared/extensions/ssh.test.ts    | 70 ++++++++++++++++++-
 2 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/packages/core/src/shared/extensions/ssh.ts b/packages/core/src/shared/extensions/ssh.ts
index 6f5d3ad13cf..d785aebf127 100644
--- a/packages/core/src/shared/extensions/ssh.ts
+++ b/packages/core/src/shared/extensions/ssh.ts
@@ -8,7 +8,7 @@ import * as path from 'path'
 import * as nls from 'vscode-nls'
 import fs from '../fs/fs'
 import { getLogger } from '../logger'
-import { ChildProcess } from '../utilities/processUtils'
+import { ChildProcess, ChildProcessResult } from '../utilities/processUtils'
 import { ArrayConstructor, NonNullObject } from '../utilities/typeConstructors'
 import { Settings } from '../settings'
 import { VSCODE_EXTENSION_ID } from '../extensions'
@@ -129,16 +129,16 @@ export async function testSshConnection(
     sshPath: string,
     user: string,
     session: SSM.StartSessionResponse
-): Promise<void> {
+): Promise<ChildProcessResult | never> {
     try {
         const env = { SESSION_ID: session.SessionId, STREAM_URL: session.StreamUrl, TOKEN: session.TokenValue }
-        await new ProcessClass(sshPath, ['-T', `${user}@${hostname}`, 'echo connected && exit']).run({
+        const result = await new ProcessClass(sshPath, ['-T', `${user}@${hostname}`, 'echo connected && exit']).run({
             spawnOptions: {
                 env,
             },
         })
+        return result
     } catch (error) {
-        getLogger().error('SSH connection test failed: %O', error)
         throw new SSHError('SSH connection test failed', { cause: error as Error })
     }
 }
diff --git a/packages/core/src/test/shared/extensions/ssh.test.ts b/packages/core/src/test/shared/extensions/ssh.test.ts
index c7abc7095cd..c300341aacd 100644
--- a/packages/core/src/test/shared/extensions/ssh.test.ts
+++ b/packages/core/src/test/shared/extensions/ssh.test.ts
@@ -4,7 +4,13 @@
  */
 import * as assert from 'assert'
 import { ChildProcess } from '../../../shared/utilities/processUtils'
-import { startSshAgent } from '../../../shared/extensions/ssh'
+import { startSshAgent, testSshConnection } from '../../../shared/extensions/ssh'
+import { createBoundProcess } from '../../../shared/remoteSession'
+import { createExecutableFile, createTestWorkspaceFolder } from '../../testUtil'
+import { WorkspaceFolder } from 'vscode'
+import path from 'path'
+import { SSM } from 'aws-sdk'
+import { fs } from '../../../shared/fs/fs'
 
 describe('SSH Agent', function () {
     it('can start the agent on windows', async function () {
@@ -29,3 +35,65 @@ describe('SSH Agent', function () {
         assert.strictEqual(await getStatus(), 'Running')
     })
 })
+
+describe('testSshConnection', function () {
+    let testWorkspace: WorkspaceFolder
+    let sshPath: string
+
+    before(async function () {
+        testWorkspace = await createTestWorkspaceFolder()
+        sshPath = path.join(testWorkspace.uri.fsPath, 'fakeSSH')
+    })
+
+    after(async function () {
+        await fs.delete(testWorkspace.uri.fsPath, { recursive: true, force: true })
+        await fs.delete(sshPath, { force: true })
+    })
+
+    it('runs in bound process', async function () {
+        const envProvider = async () => ({ MY_VAR: 'yes' })
+        const process = createBoundProcess(envProvider)
+        const session = {
+            SessionId: 'testSession',
+            StreamUrl: 'testUrl',
+            TokenValue: 'testToken',
+        } as SSM.StartSessionResponse
+
+        await createExecutableFile(sshPath, 'echo "$MY_VAR"')
+        const r = await testSshConnection(process, 'localhost', sshPath, 'test-user', session)
+        assert.strictEqual(r.stdout, 'yes')
+        await createExecutableFile(sshPath, 'echo "$UNDEFINED"')
+        const r2 = await testSshConnection(process, 'localhost', sshPath, 'test-user', session)
+        assert.strictEqual(r2.stdout, '')
+    })
+
+    it('injects new session into env', async function () {
+        const oldSession = {
+            SessionId: 'testSession1',
+            StreamUrl: 'testUrl1',
+            TokenValue: 'testToken1',
+        } as SSM.StartSessionResponse
+        const newSession = {
+            SessionId: 'testSession2',
+            StreamUrl: 'testUrl2',
+            TokenValue: 'testToken2',
+        } as SSM.StartSessionResponse
+        const envProvider = async () => ({
+            SESSION_ID: oldSession.SessionId,
+            STREAM_URL: oldSession.StreamUrl,
+            TOKEN: oldSession.TokenValue,
+        })
+        const process = createBoundProcess(envProvider)
+
+        await createExecutableFile(sshPath, 'echo "$SESSION_ID, $STREAM_URL, $TOKEN"')
+        const r = await testSshConnection(process, 'localhost', sshPath, 'test-user', newSession)
+        assert.strictEqual(r.stdout, `${newSession.SessionId}, ${newSession.StreamUrl}, ${newSession.TokenValue}`)
+    })
+
+    it('passes proper args to the ssh invoke', async function () {
+        const process = createBoundProcess(async () => ({}))
+        await createExecutableFile(sshPath, 'echo "$1,$2"')
+        const r = await testSshConnection(process, 'localhost', sshPath, 'test-user', {} as SSM.StartSessionResponse)
+        assert.strictEqual(r.stdout, '-T,test-user@localhost')
+    })
+})