Skip to content

Commit 18d92fc

Browse files
aws-ajanggaws-ajangg
committed
add more comprehensive testing
1 parent 9df0b42 commit 18d92fc

File tree

3 files changed

+324
-0
lines changed

3 files changed

+324
-0
lines changed

‎packages/core/src/test/awsService/sagemaker/commands.test.ts‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -400,4 +400,154 @@ describe('SageMaker Commands', () => {
400400
})
401401
})
402402
})
403+
404+
describe('HyperPod connection with clusterArn', function () {
405+
let mockDeeplinkConnect: sinon.SinonStub
406+
let mockIsRemoteWorkspace: sinon.SinonStub
407+
let deeplinkConnect: any
408+
409+
beforeEach(function () {
410+
mockDeeplinkConnect = sandbox.stub().resolves()
411+
mockIsRemoteWorkspace = sandbox.stub().returns(false)
412+
413+
sandbox.replace(require('../../../shared/vscode/env'), 'isRemoteWorkspace', mockIsRemoteWorkspace)
414+
415+
const freshModule = require('../../../awsService/sagemaker/commands')
416+
deeplinkConnect = freshModule.deeplinkConnect
417+
sandbox.replace(freshModule, 'deeplinkConnect', mockDeeplinkConnect)
418+
})
419+
420+
it('should create session with underscores from HyperPod clusterArn', async function () {
421+
const ctx = {
422+
extensionContext: {},
423+
} as any
424+
425+
await deeplinkConnect(
426+
ctx,
427+
'',
428+
'session-id',
429+
'wss://example.com',
430+
'token',
431+
'',
432+
undefined,
433+
'demo0',
434+
'default',
435+
'arn:aws:sagemaker:us-east-2:123456789012:cluster/n4nkkc5fbwg5'
436+
)
437+
438+
// Verify the session format uses underscores
439+
const sessionArg = mockDeeplinkConnect.firstCall?.args[10] // session parameter
440+
if (sessionArg) {
441+
assert.ok(sessionArg.includes('_'), 'Session should use underscores as separators')
442+
assert.ok(sessionArg.includes('demo0'), 'Session should include workspace name')
443+
assert.ok(sessionArg.includes('default'), 'Session should include namespace')
444+
assert.ok(sessionArg.includes('n4nkkc5fbwg5'), 'Session should include cluster name')
445+
assert.ok(sessionArg.includes('us-east-2'), 'Session should include region')
446+
assert.ok(sessionArg.includes('123456789012'), 'Session should include account ID')
447+
}
448+
})
449+
450+
it('should handle EKS clusterArn format', async function () {
451+
const ctx = {
452+
extensionContext: {},
453+
} as any
454+
455+
await deeplinkConnect(
456+
ctx,
457+
'',
458+
'session-id',
459+
'wss://example.com',
460+
'token',
461+
'',
462+
undefined,
463+
'workspace',
464+
'namespace',
465+
'arn:aws:eks:us-west-2:987654321098:cluster/eks-cluster-name'
466+
)
467+
468+
const sessionArg = mockDeeplinkConnect.firstCall?.args[10]
469+
if (sessionArg) {
470+
assert.ok(sessionArg.includes('eks-cluster-name'), 'Session should include EKS cluster name')
471+
assert.ok(sessionArg.includes('us-west-2'), 'Session should include region')
472+
assert.ok(sessionArg.includes('987654321098'), 'Session should include account ID')
473+
}
474+
})
475+
476+
it('should sanitize invalid characters in session components', async function () {
477+
const ctx = {
478+
extensionContext: {},
479+
} as any
480+
481+
await deeplinkConnect(
482+
ctx,
483+
'',
484+
'session-id',
485+
'wss://example.com',
486+
'token',
487+
'',
488+
undefined,
489+
'My@Workspace!',
490+
'my_namespace',
491+
'arn:aws:sagemaker:us-east-2:123456789012:cluster/test-cluster'
492+
)
493+
494+
const sessionArg = mockDeeplinkConnect.firstCall?.args[10]
495+
if (sessionArg) {
496+
assert.ok(!sessionArg.includes('@'), 'Session should not contain @ symbol')
497+
assert.ok(!sessionArg.includes('!'), 'Session should not contain ! symbol')
498+
assert.strictEqual(sessionArg, sessionArg.toLowerCase(), 'Session should be lowercase')
499+
}
500+
})
501+
502+
it('should handle long component names by truncating', async function () {
503+
const ctx = {
504+
extensionContext: {},
505+
} as any
506+
507+
const longWorkspace = 'a'.repeat(100)
508+
const longNamespace = 'b'.repeat(100)
509+
const longCluster = 'c'.repeat(100)
510+
511+
await deeplinkConnect(
512+
ctx,
513+
'',
514+
'session-id',
515+
'wss://example.com',
516+
'token',
517+
'',
518+
undefined,
519+
longWorkspace,
520+
longNamespace,
521+
`arn:aws:sagemaker:us-east-2:123456789012:cluster/${longCluster}`
522+
)
523+
524+
const sessionArg = mockDeeplinkConnect.firstCall?.args[10]
525+
if (sessionArg) {
526+
assert.ok(sessionArg.length <= 224, 'Session should not exceed max length')
527+
}
528+
})
529+
530+
it('should not create HyperPod session when domain is provided', async function () {
531+
const ctx = {
532+
extensionContext: {},
533+
} as any
534+
535+
await deeplinkConnect(
536+
ctx,
537+
'connection-id',
538+
'session-id',
539+
'wss://example.com',
540+
'token',
541+
'my-domain', // Domain provided - should use SageMaker Studio flow
542+
undefined,
543+
'workspace',
544+
'namespace',
545+
'arn:aws:sagemaker:us-east-2:123456789012:cluster/cluster'
546+
)
547+
548+
// Should not create HyperPod session when domain is present
549+
const sessionArg = mockDeeplinkConnect.firstCall?.args[10]
550+
assert.strictEqual(sessionArg, 'session-id', 'Should use original session when domain is provided')
551+
})
552+
})
403553
})

‎packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,4 +76,69 @@ describe('SageMaker URI handler', function () {
7676
assert.ok(deeplinkConnectStub.calledOnce)
7777
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], undefined)
7878
})
79+
80+
describe('HyperPod workspace connection', function () {
81+
function createHyperPodUri(params: { [key: string]: string }): vscode.Uri {
82+
const query = Object.entries(params)
83+
.map(([k, v]) => `${encodeURIComponent(k)}=${encodeURIComponent(v)}`)
84+
.join('&')
85+
return vscode.Uri.parse(`vscode://${VSCODE_EXTENSION_ID.awstoolkit}/connect/workspace?${query}`)
86+
}
87+
88+
it('calls deeplinkConnect with clusterArn for HyperPod connections', async function () {
89+
const params = {
90+
sessionId: 'session-123',
91+
streamUrl: 'wss://example.com/stream',
92+
sessionToken: 'token-xyz',
93+
'cell-number': '5',
94+
workspaceName: 'my-workspace',
95+
namespace: 'default',
96+
clusterArn: 'arn:aws:sagemaker:us-east-2:123456789012:cluster/my-cluster',
97+
}
98+
99+
const uri = createHyperPodUri(params)
100+
await handler.handleUri(uri)
101+
102+
assert.ok(deeplinkConnectStub.calledOnce)
103+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[1], '')
104+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[2], 'session-123')
105+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[3], 'wss://example.com/stream&cell-number=5')
106+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[4], 'token-xyz')
107+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[5], '')
108+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], undefined)
109+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[7], 'my-workspace')
110+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[8], 'default')
111+
assert.deepStrictEqual(
112+
deeplinkConnectStub.firstCall.args[9],
113+
'arn:aws:sagemaker:us-east-2:123456789012:cluster/my-cluster'
114+
)
115+
})
116+
117+
it('calls deeplinkConnect with undefined optional params when not provided', async function () {
118+
const params = {
119+
sessionId: 'session-123',
120+
streamUrl: 'wss://example.com/stream',
121+
sessionToken: 'token-xyz',
122+
'cell-number': '5',
123+
}
124+
125+
const uri = createHyperPodUri(params)
126+
await handler.handleUri(uri)
127+
128+
assert.ok(deeplinkConnectStub.calledOnce)
129+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[7], undefined)
130+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[8], undefined)
131+
assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[9], undefined)
132+
})
133+
134+
it('throws error when required params are missing', async function () {
135+
const params = {
136+
sessionId: 'session-123',
137+
// Missing streamUrl, sessionToken, cell-number
138+
}
139+
140+
const uri = createHyperPodUri(params)
141+
await assert.rejects(handler.handleUri(uri))
142+
})
143+
})
79144
})

‎packages/core/src/test/shared/clients/kubectlClient.test.ts‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,85 @@ describe('KubectlClient', function () {
164164
assert.strictEqual(result.url, 'https://test-url.com')
165165
})
166166

167+
it('should replace eksClusterArn with clusterArn in URL', async function () {
168+
const mockResponse = {
169+
response: {} as IncomingMessage,
170+
body: {
171+
status: {
172+
workspaceConnectionUrl:
173+
'vscode://test/connect?sessionId=123&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster',
174+
workspaceConnectionType: 'vscode-remote',
175+
},
176+
},
177+
}
178+
mockK8sApi.createNamespacedCustomObject.resolves(mockResponse)
179+
180+
const result = await client.createWorkspaceConnection(mockDevSpace)
181+
182+
assert.strictEqual(result.type, 'vscode-remote')
183+
assert.ok(
184+
result.url.includes(
185+
'clusterArn=arn%3Aaws%3Asagemaker%3Aus-east-2%3A123456789012%3Acluster%2Ftest-cluster'
186+
)
187+
)
188+
assert.ok(!result.url.includes('eksClusterArn'))
189+
})
190+
191+
it('should preserve other query parameters when replacing eksClusterArn', async function () {
192+
const mockResponse = {
193+
response: {} as IncomingMessage,
194+
body: {
195+
status: {
196+
workspaceConnectionUrl:
197+
'vscode://test/connect?sessionId=123&workspaceName=demo&namespace=default&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster',
198+
workspaceConnectionType: 'vscode-remote',
199+
},
200+
},
201+
}
202+
mockK8sApi.createNamespacedCustomObject.resolves(mockResponse)
203+
204+
const result = await client.createWorkspaceConnection(mockDevSpace)
205+
206+
const url = new URL(result.url)
207+
assert.strictEqual(url.searchParams.get('sessionId'), '123')
208+
assert.strictEqual(url.searchParams.get('workspaceName'), 'demo')
209+
assert.strictEqual(url.searchParams.get('namespace'), 'default')
210+
assert.ok(url.searchParams.has('clusterArn'))
211+
assert.ok(!url.searchParams.has('eksClusterArn'))
212+
})
213+
214+
it('should not modify URL if eksClusterArn is not present', async function () {
215+
const originalUrl = 'vscode://test/connect?sessionId=123&workspaceName=demo'
216+
const mockResponse = {
217+
response: {} as IncomingMessage,
218+
body: {
219+
status: {
220+
workspaceConnectionUrl: originalUrl,
221+
workspaceConnectionType: 'vscode-remote',
222+
},
223+
},
224+
}
225+
mockK8sApi.createNamespacedCustomObject.resolves(mockResponse)
226+
227+
const result = await client.createWorkspaceConnection(mockDevSpace)
228+
229+
assert.strictEqual(result.url, originalUrl)
230+
})
231+
232+
it('should throw error when presignedUrl is undefined', async function () {
233+
const mockResponse = {
234+
response: {} as IncomingMessage,
235+
body: {
236+
status: {
237+
workspaceConnectionType: 'vscode-remote',
238+
},
239+
},
240+
}
241+
mockK8sApi.createNamespacedCustomObject.resolves(mockResponse)
242+
243+
await assert.rejects(client.createWorkspaceConnection(mockDevSpace), /No workspace connection URL returned/)
244+
})
245+
167246
it('should throw error when workspace connection creation fails', async function () {
168247
mockK8sApi.createNamespacedCustomObject.rejects(new Error('Creation failed'))
169248

@@ -172,6 +251,36 @@ describe('KubectlClient', function () {
172251
/Failed to create workspace connection/
173252
)
174253
})
254+
255+
it('should not modify URL when hyperpodCluster.clusterArn is undefined', async function () {
256+
const originalUrl =
257+
'vscode://test/connect?sessionId=123&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster'
258+
const mockResponse = {
259+
response: {} as IncomingMessage,
260+
body: {
261+
status: {
262+
workspaceConnectionUrl: originalUrl,
263+
workspaceConnectionType: 'vscode-remote',
264+
},
265+
},
266+
}
267+
mockK8sApi.createNamespacedCustomObject.resolves(mockResponse)
268+
269+
// Create a client with undefined clusterArn
270+
const mockHyperpodClusterNoArn = {
271+
clusterName: 'test-cluster',
272+
clusterArn: undefined as any,
273+
status: 'InService',
274+
regionCode: 'us-east-2',
275+
}
276+
const clientNoArn = new KubectlClient(mockEksCluster, mockHyperpodClusterNoArn)
277+
;(clientNoArn as any).k8sApi = mockK8sApi
278+
279+
const result = await clientNoArn.createWorkspaceConnection(mockDevSpace)
280+
281+
// URL should remain unchanged when clusterArn is undefined
282+
assert.strictEqual(result.url, originalUrl)
283+
})
175284
})
176285

177286
describe('getSpacesForCluster', function () {

0 commit comments

Comments
 (0)