From d745df16ab442d40cce1e64463f6e881972a7b61 Mon Sep 17 00:00:00 2001 From: Austin Jang Date: Thu, 20 Nov 2025 19:13:10 -0800 Subject: [PATCH 1/5] Add unit tests --- .../explorer/sagemakerDevSpaceNode.test.ts | 219 +++++++++++++ .../explorer/sagemakerHyperpodNode.test.ts | 224 +++++++++++++ .../test/shared/clients/kubectlClient.test.ts | 304 ++++++++++++++++++ 3 files changed, 747 insertions(+) create mode 100644 packages/core/src/test/awsService/sagemaker/explorer/sagemakerDevSpaceNode.test.ts create mode 100644 packages/core/src/test/awsService/sagemaker/explorer/sagemakerHyperpodNode.test.ts create mode 100644 packages/core/src/test/shared/clients/kubectlClient.test.ts diff --git a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerDevSpaceNode.test.ts b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerDevSpaceNode.test.ts new file mode 100644 index 00000000000..44bd8537bb9 --- /dev/null +++ b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerDevSpaceNode.test.ts @@ -0,0 +1,219 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as sinon from 'sinon' +import * as assert from 'assert' +import { SagemakerDevSpaceNode } from '../../../../awsService/sagemaker/explorer/sagemakerDevSpaceNode' +import { SagemakerHyperpodNode } from '../../../../awsService/sagemaker/explorer/sagemakerHyperpodNode' +import { HyperpodDevSpace, HyperpodCluster, KubectlClient } from '../../../../shared/clients/kubectlClient' +import { SagemakerClient } from '../../../../shared/clients/sagemaker' + +describe('SagemakerDevSpaceNode', function () { + let testNode: SagemakerDevSpaceNode + let mockParent: SagemakerHyperpodNode + let mockKubectlClient: sinon.SinonStubbedInstance + let mockDevSpace: HyperpodDevSpace + let mockHyperpodCluster: HyperpodCluster + let mockSagemakerClient: sinon.SinonStubbedInstance + const testRegion = 'us-east-1' + + beforeEach(function () { + mockSagemakerClient = sinon.createStubInstance(SagemakerClient) + mockParent = new SagemakerHyperpodNode(testRegion, mockSagemakerClient as any) + mockKubectlClient = sinon.createStubInstance(KubectlClient) + mockDevSpace = { + name: 'test-space', + namespace: 'test-namespace', + cluster: 'test-cluster', + group: 'sagemaker.aws.amazon.com', + version: 'v1', + plural: 'devspaces', + status: 'Stopped', + appType: 'jupyterlab', + creator: 'test-user', + accessType: 'Public', + } + mockHyperpodCluster = { + clusterName: 'test-cluster', + clusterArn: 'arn:aws:sagemaker:us-east-1:123456789012:cluster/test-cluster', + status: 'InService', + regionCode: testRegion, + } + + sinon.stub(mockParent, 'getKubectlClient').returns(mockKubectlClient as any) + sinon.stub(mockParent, 'trackPendingNode').returns() + + testNode = new SagemakerDevSpaceNode(mockParent, mockDevSpace, mockHyperpodCluster, testRegion) + }) + + afterEach(function () { + sinon.restore() + }) + + describe('buildLabel', function () { + it('should return formatted label with name and status', function () { + const label = testNode.buildLabel() + assert.strictEqual(label, 'test-space (Stopped)') + }) + }) + + describe('buildDescription', function () { + it('should return access type description', function () { + const description = testNode.buildDescription() + assert.strictEqual(description, 'Public space') + }) + + it('should default to Public when accessType is undefined', function () { + const newDevSpace = { ...mockDevSpace, accessType: 'Public' } + const newNode = new SagemakerDevSpaceNode(mockParent, newDevSpace, mockHyperpodCluster, testRegion) + const description = newNode.buildDescription() + assert.strictEqual(description, 'Public space') + }) + }) + + describe('getContext', function () { + it('should return transitional context for Starting status', function () { + const getStatusStub = sinon.stub(testNode, 'status').get(() => 'Starting') + const context = (testNode as any).getContext() + assert.strictEqual(context, 'awsSagemakerHyperpodDevSpaceTransitionalNode') + getStatusStub.restore() + }) + + it('should return stopped context for Stopped status', function () { + const getStatusStub = sinon.stub(testNode, 'status').get(() => 'Stopped') + const context = (testNode as any).getContext() + assert.strictEqual(context, 'awsSagemakerHyperpodDevSpaceStoppedNode') + getStatusStub.restore() + }) + + it('should return running context for Running status', function () { + const getStatusStub = sinon.stub(testNode, 'status').get(() => 'Running') + const context = (testNode as any).getContext() + assert.strictEqual(context, 'awsSagemakerHyperpodDevSpaceRunningNode') + getStatusStub.restore() + }) + + it('should return error context for unknown status', function () { + const getStatusStub = sinon.stub(testNode, 'status').get(() => 'Unknown') + const context = (testNode as any).getContext() + assert.strictEqual(context, 'awsSagemakerHyperpodDevSpaceErrorNode') + getStatusStub.restore() + }) + }) + + describe('isPending', function () { + it('should return false for Running status', function () { + const getStatusStub = sinon.stub(testNode, 'status').get(() => 'Running') + assert.strictEqual(testNode.isPending(), false) + getStatusStub.restore() + }) + + it('should return false for Stopped status', function () { + const getStatusStub = sinon.stub(testNode, 'status').get(() => 'Stopped') + assert.strictEqual(testNode.isPending(), false) + getStatusStub.restore() + }) + + it('should return true for Starting status', function () { + const getStatusStub = sinon.stub(testNode, 'status').get(() => 'Starting') + assert.strictEqual(testNode.isPending(), true) + getStatusStub.restore() + }) + }) + + describe('getDevSpaceKey', function () { + it('should return formatted devspace key', function () { + const key = testNode.getDevSpaceKey() + assert.strictEqual(key, 'test-cluster-test-namespace-test-space') + }) + }) + + describe('updateWorkspaceStatus', function () { + it('should update status from kubectl client', async function () { + mockKubectlClient.getHyperpodSpaceStatus.resolves('Running') + + await testNode.updateWorkspaceStatus() + + assert.strictEqual(testNode.status, 'Running') + sinon.assert.calledOnce(mockKubectlClient.getHyperpodSpaceStatus) + }) + + it('should handle errors gracefully', async function () { + mockKubectlClient.getHyperpodSpaceStatus.rejects(new Error('API Error')) + + await testNode.updateWorkspaceStatus() + + // Should not throw, just log warning + sinon.assert.calledOnce(mockKubectlClient.getHyperpodSpaceStatus) + }) + }) + + describe('buildTooltip', function () { + it('should format tooltip with all devspace details', function () { + const tooltip = testNode.buildTooltip() + + assert.ok(tooltip.includes('test-space')) + assert.ok(tooltip.includes('test-namespace')) + assert.ok(tooltip.includes('test-cluster')) + assert.ok(tooltip.includes('test-user')) + assert.ok(tooltip.includes('Hyperpod')) + }) + }) + + describe('buildIconPath', function () { + it('should return jupyter icon for jupyterlab app type', function () { + testNode.devSpace.appType = 'jupyterlab' + + const iconPath = testNode.buildIconPath() + + assert.ok(iconPath !== undefined) + }) + + it('should return code editor icon for code-editor app type', function () { + testNode.devSpace.appType = 'code-editor' + + const iconPath = testNode.buildIconPath() + + assert.ok(iconPath !== undefined) + }) + + it('should return undefined for unknown app types', function () { + testNode.devSpace.appType = 'unknown-type' + + const iconPath = testNode.buildIconPath() + + assert.strictEqual(iconPath, undefined) + }) + }) + + describe('updateWorkspace', function () { + it('should update all node properties and track pending if needed', function () { + const isPendingStub = sinon.stub(testNode, 'isPending').returns(true) + + testNode.updateWorkspace() + + // Should update properties + assert.ok(testNode.label) + assert.ok(testNode.description) + assert.ok(testNode.tooltip) + assert.ok(testNode.contextValue) + + isPendingStub.restore() + }) + }) + + describe('refreshNode', function () { + it('should update status and refresh VS Code explorer', async function () { + const updateStatusStub = sinon.stub(testNode, 'updateWorkspaceStatus').resolves() + + await testNode.refreshNode() + + sinon.assert.calledOnce(updateStatusStub) + // Note: VS Code commands.executeCommand is mocked by the test framework + + updateStatusStub.restore() + }) + }) +}) diff --git a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerHyperpodNode.test.ts b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerHyperpodNode.test.ts new file mode 100644 index 00000000000..0014db93f2b --- /dev/null +++ b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerHyperpodNode.test.ts @@ -0,0 +1,224 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as sinon from 'sinon' +import * as assert from 'assert' +import { SagemakerHyperpodNode } from '../../../../awsService/sagemaker/explorer/sagemakerHyperpodNode' +import { SagemakerDevSpaceNode } from '../../../../awsService/sagemaker/explorer/sagemakerDevSpaceNode' +import { KubectlClient, HyperpodCluster } from '../../../../shared/clients/kubectlClient' +import { SagemakerClient } from '../../../../shared/clients/sagemaker' + +describe('SagemakerHyperpodNode', function () { + let testNode: SagemakerHyperpodNode + let mockDevSpaceNode: sinon.SinonStubbedInstance + let mockKubectlClient: sinon.SinonStubbedInstance + let mockSagemakerClient: sinon.SinonStubbedInstance + let mockHyperpodCluster: HyperpodCluster + const testRegion = 'us-east-1' + + beforeEach(function () { + mockSagemakerClient = sinon.createStubInstance(SagemakerClient) + // Mock the EKS client that will be returned + const mockEksClient = { send: sinon.stub() } + mockSagemakerClient.getEKSClient.returns(mockEksClient as any) + + testNode = new SagemakerHyperpodNode(testRegion, mockSagemakerClient as any) + mockDevSpaceNode = sinon.createStubInstance(SagemakerDevSpaceNode) + mockKubectlClient = sinon.createStubInstance(KubectlClient) + mockHyperpodCluster = { + clusterName: 'test-cluster', + clusterArn: 'arn:aws:sagemaker:us-east-1:123456789012:cluster/test-cluster', + status: 'InService', + regionCode: testRegion, + } + + mockDevSpaceNode.getDevSpaceKey.returns('test-cluster-test-namespace-test-space') + mockDevSpaceNode.isPending.returns(false) + mockDevSpaceNode.updateWorkspaceStatus.resolves() + mockDevSpaceNode.refreshNode.resolves() + }) + + afterEach(function () { + testNode.pollingSet.clear() + testNode.pollingSet.clearTimer() + sinon.restore() + }) + + describe('constructor', function () { + it('should initialize with correct properties', function () { + assert.strictEqual(testNode.regionCode, testRegion) + assert.strictEqual(testNode.label, 'HyperPod') + assert.ok(testNode.hyperpodDevSpaceNodes instanceof Map) + assert.ok(testNode.kubectlClients instanceof Map) + assert.ok(testNode.pollingSet) + }) + }) + + describe('getKubectlClient', function () { + it('should return kubectl client for cluster', function () { + const clusterName = 'test-cluster' + testNode.kubectlClients.set(clusterName, mockKubectlClient as any) + + const client = testNode.getKubectlClient(clusterName) + assert.strictEqual(client, mockKubectlClient) + }) + }) + + describe('trackPendingNode', function () { + it('should add devspace key to polling set', function () { + const devSpaceKey = 'test-cluster-test-namespace-test-space' + + testNode.trackPendingNode(devSpaceKey) + + assert.ok(testNode.pollingSet.has(devSpaceKey)) + }) + }) + + describe('updatePendingNodes', function () { + it('should update pending nodes and remove from polling when not pending', async function () { + const devSpaceKey = 'test-cluster-test-namespace-test-space' + testNode.hyperpodDevSpaceNodes.set(devSpaceKey, mockDevSpaceNode as any) + testNode.pollingSet.add(devSpaceKey) + + mockDevSpaceNode.isPending.returns(false) + + await (testNode as any).updatePendingNodes() + + sinon.assert.calledOnce(mockDevSpaceNode.updateWorkspaceStatus) + sinon.assert.calledOnce(mockDevSpaceNode.refreshNode) + assert.ok(!testNode.pollingSet.has(devSpaceKey)) + }) + + it('should keep pending nodes in polling set', async function () { + const devSpaceKey = 'test-cluster-test-namespace-test-space' + testNode.hyperpodDevSpaceNodes.set(devSpaceKey, mockDevSpaceNode as any) + testNode.pollingSet.add(devSpaceKey) + + mockDevSpaceNode.isPending.returns(true) + + await (testNode as any).updatePendingNodes() + + sinon.assert.calledOnce(mockDevSpaceNode.updateWorkspaceStatus) + sinon.assert.notCalled(mockDevSpaceNode.refreshNode) + assert.ok(testNode.pollingSet.has(devSpaceKey)) + }) + + it('should throw error when devspace not found in map', async function () { + const devSpaceKey = 'missing-key' + testNode.pollingSet.add(devSpaceKey) + + await assert.rejects( + (testNode as any).updatePendingNodes(), + /Devspace missing-key from polling set not found/ + ) + }) + }) + + describe('listSpaces', function () { + it('should discover spaces across multiple clusters', async function () { + const mockClusters = [ + { + clusterName: 'cluster1', + clusterArn: 'arn:aws:sagemaker:us-east-1:123:cluster/cluster1', + status: 'InService', + eksClusterName: 'eks1', + regionCode: testRegion, + }, + ] + mockSagemakerClient.listHyperpodClusters.resolves(mockClusters) + + const mockEksResponse = { cluster: { name: 'eks1', endpoint: 'https://test.com' } } + ;(testNode.eksClient as any).send.resolves(mockEksResponse) + + const mockKubectl = { getSpacesForCluster: sinon.stub().resolves([]) } + testNode.kubectlClients.set('cluster1', mockKubectl as any) + + const result = await testNode.listSpaces() + + assert.ok(result instanceof Map) + sinon.assert.calledOnce(mockSagemakerClient.listHyperpodClusters) + }) + + it('should handle clusters without EKS integration', async function () { + const mockClusters = [ + { + clusterName: 'cluster1', + clusterArn: 'arn:aws:sagemaker:us-east-1:123:cluster/cluster1', + status: 'InService', + regionCode: testRegion, + }, + ] // No eksClusterName + mockSagemakerClient.listHyperpodClusters.resolves(mockClusters) + + const result = await testNode.listSpaces() + + assert.strictEqual(result.size, 0) + }) + + it('should handle kubectl client creation errors', async function () { + mockSagemakerClient.listHyperpodClusters.rejects(new Error('API Error')) + + await assert.rejects(testNode.listSpaces(), /No workspaces listed/) + }) + }) + + describe('updateChildren', function () { + it('should filter spaces based on selected cluster namespaces', async function () { + const mockDevSpace = { + name: 'test-space', + namespace: 'test-namespace', + cluster: 'test-cluster', // This is the key field needed + environment: 'test-env', + application: 'test-app', + group: 'test-group', + version: 'v1', + plural: 'spaces', + status: 'Running', + appType: 'jupyterlab', + creator: 'test-user', + accessType: 'Public', + } + const mockSpaces = new Map([ + [ + 'key1', + { + cluster: mockHyperpodCluster, + devSpace: mockDevSpace, + }, + ], + ]) + sinon.stub(testNode, 'listSpaces').resolves(mockSpaces) + sinon.stub(testNode, 'getSelectedClusterNamespaces').resolves(new Set(['test-cluster-test-namespace'])) + const stsStub = sinon.stub((testNode as any).stsClient, 'getCallerIdentity').resolves({ Arn: 'test-arn' }) + + await testNode.updateChildren() + + assert.ok(testNode.hyperpodDevSpaceNodes instanceof Map) + stsStub.restore() + }) + + it('should handle caller identity retrieval', async function () { + sinon.stub(testNode, 'listSpaces').resolves(new Map()) + sinon.stub(testNode, 'getSelectedClusterNamespaces').resolves(new Set()) + const stsStub = sinon.stub((testNode as any).stsClient, 'getCallerIdentity').resolves({ Arn: 'test-arn' }) + + await testNode.updateChildren() + + sinon.assert.calledOnce(stsStub) + stsStub.restore() + }) + }) + + describe('getSelectedClusterNamespaces', function () { + it('should return defaults when no cache exists', async function () { + sinon.stub(testNode, 'getDefaultSelectedClusterNamespaces').resolves(['default-selection']) + ;(testNode as any).callerIdentity = { Arn: 'test-arn' } + + const result = await testNode.getSelectedClusterNamespaces() + + assert.ok(result.has('default-selection')) + }) + }) +}) diff --git a/packages/core/src/test/shared/clients/kubectlClient.test.ts b/packages/core/src/test/shared/clients/kubectlClient.test.ts new file mode 100644 index 00000000000..5a7bac8c770 --- /dev/null +++ b/packages/core/src/test/shared/clients/kubectlClient.test.ts @@ -0,0 +1,304 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as sinon from 'sinon' +import * as assert from 'assert' +import * as k8s from '@kubernetes/client-node' +import { KubectlClient, HyperpodDevSpace, HyperpodCluster } from '../../../shared/clients/kubectlClient' +import { SagemakerDevSpaceNode } from '../../../awsService/sagemaker/explorer/sagemakerDevSpaceNode' +import { Cluster } from '@aws-sdk/client-eks' +import { IncomingMessage } from 'http' + +describe('KubectlClient', function () { + let client: KubectlClient + let mockK8sApi: sinon.SinonStubbedInstance + let mockDevSpace: HyperpodDevSpace + let mockHyperpodCluster: HyperpodCluster + let mockDevSpaceNode: sinon.SinonStubbedInstance + let mockEksCluster: Cluster + + beforeEach(function () { + mockK8sApi = sinon.createStubInstance(k8s.CustomObjectsApi) + mockDevSpace = { + name: 'test-space', + namespace: 'test-namespace', + cluster: 'test-cluster', + group: 'sagemaker.aws.amazon.com', + version: 'v1', + plural: 'devspaces', + status: 'Stopped', + appType: 'jupyterlab', + creator: 'test-user', + accessType: 'Public', + } + mockHyperpodCluster = { + clusterName: 'test-cluster', + clusterArn: 'arn:aws:sagemaker:us-east-1:123456789012:cluster/test-cluster', + status: 'InService', + regionCode: 'us-east-1', + } + mockEksCluster = { + name: 'test-cluster', + endpoint: 'https://test-endpoint.com', + certificateAuthority: { data: 'test-cert-data' }, + } + mockDevSpaceNode = sinon.createStubInstance(SagemakerDevSpaceNode) + Object.defineProperty(mockDevSpaceNode, 'devSpace', { + value: mockDevSpace, + writable: false, + }) + + client = new KubectlClient(mockEksCluster, mockHyperpodCluster) + ;(client as any).k8sApi = mockK8sApi + }) + + afterEach(function () { + sinon.restore() + }) + + describe('getHyperpodSpaceStatus', function () { + it('should return Running status when available and not progressing', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + conditions: [ + { type: 'Available', status: 'True' }, + { type: 'Progressing', status: 'False' }, + { type: 'Stopped', status: 'False' }, + ], + }, + spec: { desiredStatus: 'Running' }, + }, + } + mockK8sApi.getNamespacedCustomObject.resolves(mockResponse) + + const status = await client.getHyperpodSpaceStatus(mockDevSpace) + assert.strictEqual(status, 'Running') + }) + + it('should return Starting status when progressing with Running desired status', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + conditions: [ + { type: 'Available', status: 'False' }, + { type: 'Progressing', status: 'True' }, + ], + }, + spec: { desiredStatus: 'Running' }, + }, + } + mockK8sApi.getNamespacedCustomObject.resolves(mockResponse) + + const status = await client.getHyperpodSpaceStatus(mockDevSpace) + assert.strictEqual(status, 'Starting') + }) + + it('should return Error status when degraded', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + conditions: [{ type: 'Degraded', status: 'True' }], + }, + }, + } + mockK8sApi.getNamespacedCustomObject.resolves(mockResponse) + + const status = await client.getHyperpodSpaceStatus(mockDevSpace) + assert.strictEqual(status, 'Error') + }) + + it('should throw error when API call fails', async function () { + mockK8sApi.getNamespacedCustomObject.rejects(new Error('API Error')) + + await assert.rejects( + client.getHyperpodSpaceStatus(mockDevSpace), + /Failed to get status for devSpace: test-space/ + ) + }) + }) + + describe('patchDevSpaceStatus', function () { + it('should patch devspace with Running status', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: {}, + } + mockK8sApi.patchNamespacedCustomObject.resolves(mockResponse) + + await client.patchDevSpaceStatus(mockDevSpace, 'Running') + + sinon.assert.calledOnceWithExactly( + mockK8sApi.patchNamespacedCustomObject, + 'sagemaker.aws.amazon.com', + 'v1', + 'test-namespace', + 'devspaces', + 'test-space', + { spec: { desiredStatus: 'Running' } }, + undefined, + undefined, + undefined, + { headers: { 'Content-Type': 'application/merge-patch+json' } } + ) + }) + + it('should throw error when patch fails', async function () { + mockK8sApi.patchNamespacedCustomObject.rejects(new Error('Patch failed')) + + await assert.rejects( + client.patchDevSpaceStatus(mockDevSpace, 'Stopped'), + /Failed to update transitional status for devSpace test-space/ + ) + }) + }) + + describe('createWorkspaceConnection', function () { + it('should create workspace connection and return connection details', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + workspaceConnectionUrl: 'https://test-url.com', + workspaceConnectionType: 'vscode-remote', + }, + }, + } + mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) + + const result = await client.createWorkspaceConnection(mockDevSpace) + + assert.strictEqual(result.type, 'vscode-remote') + assert.strictEqual(result.url, 'https://test-url.com') + }) + + it('should throw error when workspace connection creation fails', async function () { + mockK8sApi.createNamespacedCustomObject.rejects(new Error('Creation failed')) + + await assert.rejects( + client.createWorkspaceConnection(mockDevSpace), + /Failed to create workspace connection/ + ) + }) + }) + + describe('getSpacesForCluster', function () { + it('should return mapped workspaces from Kubernetes API', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { + items: [ + { + metadata: { + name: 'test-workspace', + namespace: 'test-namespace', + annotations: { 'workspace.jupyter.org/created-by': 'test-user' }, + }, + spec: { appType: 'jupyterlab', accessType: 'Public', desiredStatus: 'Running' }, + status: { conditions: [{ type: 'Available', status: 'True' }] }, + }, + ], + }, + } + mockK8sApi.listClusterCustomObject.resolves(mockResponse) + + const result = await client.getSpacesForCluster(mockEksCluster) + + assert.strictEqual(result.length, 1) + assert.strictEqual(result[0].name, 'test-workspace') + assert.strictEqual(result[0].namespace, 'test-namespace') + assert.strictEqual(result[0].creator, 'test-user') + }) + + it('should handle 403 permission errors with user message', async function () { + const error = new Error('Forbidden') + ;(error as any).statusCode = 403 + mockK8sApi.listClusterCustomObject.rejects(error) + + const result = await client.getSpacesForCluster(mockEksCluster) + + assert.strictEqual(result.length, 0) + }) + + it('should return empty array when API returns no items', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: {}, + } + mockK8sApi.listClusterCustomObject.resolves(mockResponse) + + const result = await client.getSpacesForCluster(mockEksCluster) + + assert.strictEqual(result.length, 0) + }) + + it('should return empty array when no spaces found', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { items: [] }, + } + mockK8sApi.listClusterCustomObject.resolves(mockResponse) + + const result = await client.getSpacesForCluster(mockEksCluster) + + assert.strictEqual(result.length, 0) + }) + }) + + describe('startHyperpodDevSpace', function () { + it('should patch status to Running and track pending node', async function () { + const mockParent = { trackPendingNode: sinon.stub() } + mockDevSpaceNode.getParent.returns(mockParent as any) + mockDevSpaceNode.getDevSpaceKey.returns('test-key') + mockK8sApi.patchNamespacedCustomObject.resolves({} as any) + mockK8sApi.getNamespacedCustomObject.resolves({ + response: {} as IncomingMessage, + body: { status: { conditions: [] } }, + }) + + await client.startHyperpodDevSpace(mockDevSpaceNode as any) + + sinon.assert.calledWith( + mockK8sApi.patchNamespacedCustomObject, + mockDevSpace.group, + mockDevSpace.version, + mockDevSpace.namespace, + mockDevSpace.plural, + mockDevSpace.name, + { spec: { desiredStatus: 'Running' } } + ) + sinon.assert.calledWith(mockParent.trackPendingNode, 'test-key') + }) + }) + + describe('stopHyperpodDevSpace', function () { + it('should patch status to Stopped and track pending node', async function () { + const mockParent = { trackPendingNode: sinon.stub() } + mockDevSpaceNode.getParent.returns(mockParent as any) + mockDevSpaceNode.getDevSpaceKey.returns('test-key') + mockK8sApi.patchNamespacedCustomObject.resolves({} as any) + mockK8sApi.getNamespacedCustomObject.resolves({ + response: {} as IncomingMessage, + body: { status: { conditions: [] } }, + }) + + await client.stopHyperpodDevSpace(mockDevSpaceNode as any) + + sinon.assert.calledWith( + mockK8sApi.patchNamespacedCustomObject, + mockDevSpace.group, + mockDevSpace.version, + mockDevSpace.namespace, + mockDevSpace.plural, + mockDevSpace.name, + { spec: { desiredStatus: 'Stopped' } } + ) + sinon.assert.calledWith(mockParent.trackPendingNode, 'test-key') + }) + }) +}) From 9df0b4292904c5fb5637274bb48eec497f08681e Mon Sep 17 00:00:00 2001 From: aws-ajangg Date: Mon, 12 Jan 2026 13:27:45 -0800 Subject: [PATCH 2/5] lint fix for duplicate code --- .../explorer/sagemakerDevSpaceNode.test.ts | 23 +++-------- .../test/shared/clients/kubectlClient.test.ts | 33 ++++------------ .../test/shared/clients/kubectlTestHelpers.ts | 38 +++++++++++++++++++ 3 files changed, 51 insertions(+), 43 deletions(-) create mode 100644 packages/core/src/test/shared/clients/kubectlTestHelpers.ts diff --git a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerDevSpaceNode.test.ts b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerDevSpaceNode.test.ts index 44bd8537bb9..a7f95b21929 100644 --- a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerDevSpaceNode.test.ts +++ b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerDevSpaceNode.test.ts @@ -9,6 +9,7 @@ import { SagemakerDevSpaceNode } from '../../../../awsService/sagemaker/explorer import { SagemakerHyperpodNode } from '../../../../awsService/sagemaker/explorer/sagemakerHyperpodNode' import { HyperpodDevSpace, HyperpodCluster, KubectlClient } from '../../../../shared/clients/kubectlClient' import { SagemakerClient } from '../../../../shared/clients/sagemaker' +import { createMockK8sSetup } from '../../../shared/clients/kubectlTestHelpers' describe('SagemakerDevSpaceNode', function () { let testNode: SagemakerDevSpaceNode @@ -23,24 +24,10 @@ describe('SagemakerDevSpaceNode', function () { mockSagemakerClient = sinon.createStubInstance(SagemakerClient) mockParent = new SagemakerHyperpodNode(testRegion, mockSagemakerClient as any) mockKubectlClient = sinon.createStubInstance(KubectlClient) - mockDevSpace = { - name: 'test-space', - namespace: 'test-namespace', - cluster: 'test-cluster', - group: 'sagemaker.aws.amazon.com', - version: 'v1', - plural: 'devspaces', - status: 'Stopped', - appType: 'jupyterlab', - creator: 'test-user', - accessType: 'Public', - } - mockHyperpodCluster = { - clusterName: 'test-cluster', - clusterArn: 'arn:aws:sagemaker:us-east-1:123456789012:cluster/test-cluster', - status: 'InService', - regionCode: testRegion, - } + + const mockSetup = createMockK8sSetup() + mockDevSpace = mockSetup.mockDevSpace as HyperpodDevSpace + mockHyperpodCluster = mockSetup.mockHyperpodCluster sinon.stub(mockParent, 'getKubectlClient').returns(mockKubectlClient as any) sinon.stub(mockParent, 'trackPendingNode').returns() diff --git a/packages/core/src/test/shared/clients/kubectlClient.test.ts b/packages/core/src/test/shared/clients/kubectlClient.test.ts index 5a7bac8c770..1333cb85902 100644 --- a/packages/core/src/test/shared/clients/kubectlClient.test.ts +++ b/packages/core/src/test/shared/clients/kubectlClient.test.ts @@ -9,6 +9,7 @@ import * as k8s from '@kubernetes/client-node' import { KubectlClient, HyperpodDevSpace, HyperpodCluster } from '../../../shared/clients/kubectlClient' import { SagemakerDevSpaceNode } from '../../../awsService/sagemaker/explorer/sagemakerDevSpaceNode' import { Cluster } from '@aws-sdk/client-eks' +import { createMockK8sSetup, setupMockDevSpaceNode } from './kubectlTestHelpers' import { IncomingMessage } from 'http' describe('KubectlClient', function () { @@ -20,25 +21,11 @@ describe('KubectlClient', function () { let mockEksCluster: Cluster beforeEach(function () { - mockK8sApi = sinon.createStubInstance(k8s.CustomObjectsApi) - mockDevSpace = { - name: 'test-space', - namespace: 'test-namespace', - cluster: 'test-cluster', - group: 'sagemaker.aws.amazon.com', - version: 'v1', - plural: 'devspaces', - status: 'Stopped', - appType: 'jupyterlab', - creator: 'test-user', - accessType: 'Public', - } - mockHyperpodCluster = { - clusterName: 'test-cluster', - clusterArn: 'arn:aws:sagemaker:us-east-1:123456789012:cluster/test-cluster', - status: 'InService', - regionCode: 'us-east-1', - } + const mockSetup = createMockK8sSetup() + mockK8sApi = mockSetup.mockK8sApi + mockDevSpace = mockSetup.mockDevSpace + mockHyperpodCluster = mockSetup.mockHyperpodCluster + mockEksCluster = { name: 'test-cluster', endpoint: 'https://test-endpoint.com', @@ -252,9 +239,7 @@ describe('KubectlClient', function () { describe('startHyperpodDevSpace', function () { it('should patch status to Running and track pending node', async function () { - const mockParent = { trackPendingNode: sinon.stub() } - mockDevSpaceNode.getParent.returns(mockParent as any) - mockDevSpaceNode.getDevSpaceKey.returns('test-key') + const mockParent = setupMockDevSpaceNode(mockDevSpaceNode) mockK8sApi.patchNamespacedCustomObject.resolves({} as any) mockK8sApi.getNamespacedCustomObject.resolves({ response: {} as IncomingMessage, @@ -278,9 +263,7 @@ describe('KubectlClient', function () { describe('stopHyperpodDevSpace', function () { it('should patch status to Stopped and track pending node', async function () { - const mockParent = { trackPendingNode: sinon.stub() } - mockDevSpaceNode.getParent.returns(mockParent as any) - mockDevSpaceNode.getDevSpaceKey.returns('test-key') + const mockParent = setupMockDevSpaceNode(mockDevSpaceNode) mockK8sApi.patchNamespacedCustomObject.resolves({} as any) mockK8sApi.getNamespacedCustomObject.resolves({ response: {} as IncomingMessage, diff --git a/packages/core/src/test/shared/clients/kubectlTestHelpers.ts b/packages/core/src/test/shared/clients/kubectlTestHelpers.ts new file mode 100644 index 00000000000..88b731064d3 --- /dev/null +++ b/packages/core/src/test/shared/clients/kubectlTestHelpers.ts @@ -0,0 +1,38 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as sinon from 'sinon' +import * as k8s from '@kubernetes/client-node' + +export function createMockK8sSetup() { + const mockK8sApi = sinon.createStubInstance(k8s.CustomObjectsApi) + const mockDevSpace = { + name: 'test-space', + namespace: 'test-namespace', + cluster: 'test-cluster', + group: 'sagemaker.aws.amazon.com', + version: 'v1', + plural: 'devspaces', + status: 'Stopped', + appType: 'jupyterlab', + creator: 'test-user', + accessType: 'Public', + } + const mockHyperpodCluster = { + clusterName: 'test-cluster', + clusterArn: 'arn:aws:sagemaker:us-east-1:123456789012:cluster/test-cluster', + status: 'InService', + regionCode: 'us-east-1', + } + + return { mockK8sApi, mockDevSpace, mockHyperpodCluster } +} + +export function setupMockDevSpaceNode(mockDevSpaceNode: any) { + const mockParent = { trackPendingNode: sinon.stub() } + mockDevSpaceNode.getParent.returns(mockParent as any) + mockDevSpaceNode.getDevSpaceKey.returns('test-key') + return mockParent +} From 18d92fc92795e5208774659e1ae1c30856649fce Mon Sep 17 00:00:00 2001 From: aws-ajangg Date: Thu, 15 Jan 2026 14:50:25 -0800 Subject: [PATCH 3/5] add more comprehensive testing --- .../awsService/sagemaker/commands.test.ts | 150 ++++++++++++++++++ .../awsService/sagemaker/uriHandlers.test.ts | 65 ++++++++ .../test/shared/clients/kubectlClient.test.ts | 109 +++++++++++++ 3 files changed, 324 insertions(+) diff --git a/packages/core/src/test/awsService/sagemaker/commands.test.ts b/packages/core/src/test/awsService/sagemaker/commands.test.ts index 8b7a445b1b4..35d05b47c86 100644 --- a/packages/core/src/test/awsService/sagemaker/commands.test.ts +++ b/packages/core/src/test/awsService/sagemaker/commands.test.ts @@ -400,4 +400,154 @@ describe('SageMaker Commands', () => { }) }) }) + + describe('HyperPod connection with clusterArn', function () { + let mockDeeplinkConnect: sinon.SinonStub + let mockIsRemoteWorkspace: sinon.SinonStub + let deeplinkConnect: any + + beforeEach(function () { + mockDeeplinkConnect = sandbox.stub().resolves() + mockIsRemoteWorkspace = sandbox.stub().returns(false) + + sandbox.replace(require('../../../shared/vscode/env'), 'isRemoteWorkspace', mockIsRemoteWorkspace) + + const freshModule = require('../../../awsService/sagemaker/commands') + deeplinkConnect = freshModule.deeplinkConnect + sandbox.replace(freshModule, 'deeplinkConnect', mockDeeplinkConnect) + }) + + it('should create session with underscores from HyperPod clusterArn', async function () { + const ctx = { + extensionContext: {}, + } as any + + await deeplinkConnect( + ctx, + '', + 'session-id', + 'wss://example.com', + 'token', + '', + undefined, + 'demo0', + 'default', + 'arn:aws:sagemaker:us-east-2:123456789012:cluster/n4nkkc5fbwg5' + ) + + // Verify the session format uses underscores + const sessionArg = mockDeeplinkConnect.firstCall?.args[10] // session parameter + if (sessionArg) { + assert.ok(sessionArg.includes('_'), 'Session should use underscores as separators') + assert.ok(sessionArg.includes('demo0'), 'Session should include workspace name') + assert.ok(sessionArg.includes('default'), 'Session should include namespace') + assert.ok(sessionArg.includes('n4nkkc5fbwg5'), 'Session should include cluster name') + assert.ok(sessionArg.includes('us-east-2'), 'Session should include region') + assert.ok(sessionArg.includes('123456789012'), 'Session should include account ID') + } + }) + + it('should handle EKS clusterArn format', async function () { + const ctx = { + extensionContext: {}, + } as any + + await deeplinkConnect( + ctx, + '', + 'session-id', + 'wss://example.com', + 'token', + '', + undefined, + 'workspace', + 'namespace', + 'arn:aws:eks:us-west-2:987654321098:cluster/eks-cluster-name' + ) + + const sessionArg = mockDeeplinkConnect.firstCall?.args[10] + if (sessionArg) { + assert.ok(sessionArg.includes('eks-cluster-name'), 'Session should include EKS cluster name') + assert.ok(sessionArg.includes('us-west-2'), 'Session should include region') + assert.ok(sessionArg.includes('987654321098'), 'Session should include account ID') + } + }) + + it('should sanitize invalid characters in session components', async function () { + const ctx = { + extensionContext: {}, + } as any + + await deeplinkConnect( + ctx, + '', + 'session-id', + 'wss://example.com', + 'token', + '', + undefined, + 'My@Workspace!', + 'my_namespace', + 'arn:aws:sagemaker:us-east-2:123456789012:cluster/test-cluster' + ) + + const sessionArg = mockDeeplinkConnect.firstCall?.args[10] + if (sessionArg) { + assert.ok(!sessionArg.includes('@'), 'Session should not contain @ symbol') + assert.ok(!sessionArg.includes('!'), 'Session should not contain ! symbol') + assert.strictEqual(sessionArg, sessionArg.toLowerCase(), 'Session should be lowercase') + } + }) + + it('should handle long component names by truncating', async function () { + const ctx = { + extensionContext: {}, + } as any + + const longWorkspace = 'a'.repeat(100) + const longNamespace = 'b'.repeat(100) + const longCluster = 'c'.repeat(100) + + await deeplinkConnect( + ctx, + '', + 'session-id', + 'wss://example.com', + 'token', + '', + undefined, + longWorkspace, + longNamespace, + `arn:aws:sagemaker:us-east-2:123456789012:cluster/${longCluster}` + ) + + const sessionArg = mockDeeplinkConnect.firstCall?.args[10] + if (sessionArg) { + assert.ok(sessionArg.length <= 224, 'Session should not exceed max length') + } + }) + + it('should not create HyperPod session when domain is provided', async function () { + const ctx = { + extensionContext: {}, + } as any + + await deeplinkConnect( + ctx, + 'connection-id', + 'session-id', + 'wss://example.com', + 'token', + 'my-domain', // Domain provided - should use SageMaker Studio flow + undefined, + 'workspace', + 'namespace', + 'arn:aws:sagemaker:us-east-2:123456789012:cluster/cluster' + ) + + // Should not create HyperPod session when domain is present + const sessionArg = mockDeeplinkConnect.firstCall?.args[10] + assert.strictEqual(sessionArg, 'session-id', 'Should use original session when domain is provided') + }) + }) }) diff --git a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts index f27df1fcb11..31007519ace 100644 --- a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts +++ b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts @@ -76,4 +76,69 @@ describe('SageMaker URI handler', function () { assert.ok(deeplinkConnectStub.calledOnce) assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], undefined) }) + + describe('HyperPod workspace connection', function () { + function createHyperPodUri(params: { [key: string]: string }): vscode.Uri { + const query = Object.entries(params) + .map(([k, v]) => `${encodeURIComponent(k)}=${encodeURIComponent(v)}`) + .join('&') + return vscode.Uri.parse(`vscode://${VSCODE_EXTENSION_ID.awstoolkit}/connect/workspace?${query}`) + } + + it('calls deeplinkConnect with clusterArn for HyperPod connections', async function () { + const params = { + sessionId: 'session-123', + streamUrl: 'wss://example.com/stream', + sessionToken: 'token-xyz', + 'cell-number': '5', + workspaceName: 'my-workspace', + namespace: 'default', + clusterArn: 'arn:aws:sagemaker:us-east-2:123456789012:cluster/my-cluster', + } + + const uri = createHyperPodUri(params) + await handler.handleUri(uri) + + assert.ok(deeplinkConnectStub.calledOnce) + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[1], '') + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[2], 'session-123') + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[3], 'wss://example.com/stream&cell-number=5') + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[4], 'token-xyz') + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[5], '') + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[6], undefined) + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[7], 'my-workspace') + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[8], 'default') + assert.deepStrictEqual( + deeplinkConnectStub.firstCall.args[9], + 'arn:aws:sagemaker:us-east-2:123456789012:cluster/my-cluster' + ) + }) + + it('calls deeplinkConnect with undefined optional params when not provided', async function () { + const params = { + sessionId: 'session-123', + streamUrl: 'wss://example.com/stream', + sessionToken: 'token-xyz', + 'cell-number': '5', + } + + const uri = createHyperPodUri(params) + await handler.handleUri(uri) + + assert.ok(deeplinkConnectStub.calledOnce) + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[7], undefined) + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[8], undefined) + assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[9], undefined) + }) + + it('throws error when required params are missing', async function () { + const params = { + sessionId: 'session-123', + // Missing streamUrl, sessionToken, cell-number + } + + const uri = createHyperPodUri(params) + await assert.rejects(handler.handleUri(uri)) + }) + }) }) diff --git a/packages/core/src/test/shared/clients/kubectlClient.test.ts b/packages/core/src/test/shared/clients/kubectlClient.test.ts index 1333cb85902..9f30524c6f6 100644 --- a/packages/core/src/test/shared/clients/kubectlClient.test.ts +++ b/packages/core/src/test/shared/clients/kubectlClient.test.ts @@ -164,6 +164,85 @@ describe('KubectlClient', function () { assert.strictEqual(result.url, 'https://test-url.com') }) + it('should replace eksClusterArn with clusterArn in URL', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + workspaceConnectionUrl: + 'vscode://test/connect?sessionId=123&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster', + workspaceConnectionType: 'vscode-remote', + }, + }, + } + mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) + + const result = await client.createWorkspaceConnection(mockDevSpace) + + assert.strictEqual(result.type, 'vscode-remote') + assert.ok( + result.url.includes( + 'clusterArn=arn%3Aaws%3Asagemaker%3Aus-east-2%3A123456789012%3Acluster%2Ftest-cluster' + ) + ) + assert.ok(!result.url.includes('eksClusterArn')) + }) + + it('should preserve other query parameters when replacing eksClusterArn', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + workspaceConnectionUrl: + 'vscode://test/connect?sessionId=123&workspaceName=demo&namespace=default&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster', + workspaceConnectionType: 'vscode-remote', + }, + }, + } + mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) + + const result = await client.createWorkspaceConnection(mockDevSpace) + + const url = new URL(result.url) + assert.strictEqual(url.searchParams.get('sessionId'), '123') + assert.strictEqual(url.searchParams.get('workspaceName'), 'demo') + assert.strictEqual(url.searchParams.get('namespace'), 'default') + assert.ok(url.searchParams.has('clusterArn')) + assert.ok(!url.searchParams.has('eksClusterArn')) + }) + + it('should not modify URL if eksClusterArn is not present', async function () { + const originalUrl = 'vscode://test/connect?sessionId=123&workspaceName=demo' + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + workspaceConnectionUrl: originalUrl, + workspaceConnectionType: 'vscode-remote', + }, + }, + } + mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) + + const result = await client.createWorkspaceConnection(mockDevSpace) + + assert.strictEqual(result.url, originalUrl) + }) + + it('should throw error when presignedUrl is undefined', async function () { + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + workspaceConnectionType: 'vscode-remote', + }, + }, + } + mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) + + await assert.rejects(client.createWorkspaceConnection(mockDevSpace), /No workspace connection URL returned/) + }) + it('should throw error when workspace connection creation fails', async function () { mockK8sApi.createNamespacedCustomObject.rejects(new Error('Creation failed')) @@ -172,6 +251,36 @@ describe('KubectlClient', function () { /Failed to create workspace connection/ ) }) + + it('should not modify URL when hyperpodCluster.clusterArn is undefined', async function () { + const originalUrl = + 'vscode://test/connect?sessionId=123&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster' + const mockResponse = { + response: {} as IncomingMessage, + body: { + status: { + workspaceConnectionUrl: originalUrl, + workspaceConnectionType: 'vscode-remote', + }, + }, + } + mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) + + // Create a client with undefined clusterArn + const mockHyperpodClusterNoArn = { + clusterName: 'test-cluster', + clusterArn: undefined as any, + status: 'InService', + regionCode: 'us-east-2', + } + const clientNoArn = new KubectlClient(mockEksCluster, mockHyperpodClusterNoArn) + ;(clientNoArn as any).k8sApi = mockK8sApi + + const result = await clientNoArn.createWorkspaceConnection(mockDevSpace) + + // URL should remain unchanged when clusterArn is undefined + assert.strictEqual(result.url, originalUrl) + }) }) describe('getSpacesForCluster', function () { From 2c06f83f13d2bd057e1f6fc398fb0af9b63c667b Mon Sep 17 00:00:00 2001 From: aws-ajangg Date: Fri, 16 Jan 2026 10:42:11 -0800 Subject: [PATCH 4/5] remove presign url modification test --- .../test/shared/clients/kubectlClient.test.ts | 65 ------------------- .../test/shared/clients/kubectlTestHelpers.ts | 4 +- 2 files changed, 2 insertions(+), 67 deletions(-) diff --git a/packages/core/src/test/shared/clients/kubectlClient.test.ts b/packages/core/src/test/shared/clients/kubectlClient.test.ts index 9f30524c6f6..f98661c611a 100644 --- a/packages/core/src/test/shared/clients/kubectlClient.test.ts +++ b/packages/core/src/test/shared/clients/kubectlClient.test.ts @@ -164,71 +164,6 @@ describe('KubectlClient', function () { assert.strictEqual(result.url, 'https://test-url.com') }) - it('should replace eksClusterArn with clusterArn in URL', async function () { - const mockResponse = { - response: {} as IncomingMessage, - body: { - status: { - workspaceConnectionUrl: - 'vscode://test/connect?sessionId=123&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster', - workspaceConnectionType: 'vscode-remote', - }, - }, - } - mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) - - const result = await client.createWorkspaceConnection(mockDevSpace) - - assert.strictEqual(result.type, 'vscode-remote') - assert.ok( - result.url.includes( - 'clusterArn=arn%3Aaws%3Asagemaker%3Aus-east-2%3A123456789012%3Acluster%2Ftest-cluster' - ) - ) - assert.ok(!result.url.includes('eksClusterArn')) - }) - - it('should preserve other query parameters when replacing eksClusterArn', async function () { - const mockResponse = { - response: {} as IncomingMessage, - body: { - status: { - workspaceConnectionUrl: - 'vscode://test/connect?sessionId=123&workspaceName=demo&namespace=default&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster', - workspaceConnectionType: 'vscode-remote', - }, - }, - } - mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) - - const result = await client.createWorkspaceConnection(mockDevSpace) - - const url = new URL(result.url) - assert.strictEqual(url.searchParams.get('sessionId'), '123') - assert.strictEqual(url.searchParams.get('workspaceName'), 'demo') - assert.strictEqual(url.searchParams.get('namespace'), 'default') - assert.ok(url.searchParams.has('clusterArn')) - assert.ok(!url.searchParams.has('eksClusterArn')) - }) - - it('should not modify URL if eksClusterArn is not present', async function () { - const originalUrl = 'vscode://test/connect?sessionId=123&workspaceName=demo' - const mockResponse = { - response: {} as IncomingMessage, - body: { - status: { - workspaceConnectionUrl: originalUrl, - workspaceConnectionType: 'vscode-remote', - }, - }, - } - mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) - - const result = await client.createWorkspaceConnection(mockDevSpace) - - assert.strictEqual(result.url, originalUrl) - }) - it('should throw error when presignedUrl is undefined', async function () { const mockResponse = { response: {} as IncomingMessage, diff --git a/packages/core/src/test/shared/clients/kubectlTestHelpers.ts b/packages/core/src/test/shared/clients/kubectlTestHelpers.ts index 88b731064d3..55236d70db3 100644 --- a/packages/core/src/test/shared/clients/kubectlTestHelpers.ts +++ b/packages/core/src/test/shared/clients/kubectlTestHelpers.ts @@ -22,9 +22,9 @@ export function createMockK8sSetup() { } const mockHyperpodCluster = { clusterName: 'test-cluster', - clusterArn: 'arn:aws:sagemaker:us-east-1:123456789012:cluster/test-cluster', + clusterArn: 'arn:aws:sagemaker:us-east-2:123456789012:cluster/test-cluster', status: 'InService', - regionCode: 'us-east-1', + regionCode: 'us-east-2', } return { mockK8sApi, mockDevSpace, mockHyperpodCluster } From e799173b08f33613e9b17d30f848f7702257f745 Mon Sep 17 00:00:00 2001 From: aws-ajangg Date: Sun, 1 Feb 2026 23:34:44 -0800 Subject: [PATCH 5/5] fix: modify test cases for ekc cluster arn --- .../awsService/sagemaker/commands.test.ts | 10 ++--- .../awsService/sagemaker/uriHandlers.test.ts | 6 +-- .../test/shared/clients/kubectlClient.test.ts | 40 +++++++------------ 3 files changed, 22 insertions(+), 34 deletions(-) diff --git a/packages/core/src/test/awsService/sagemaker/commands.test.ts b/packages/core/src/test/awsService/sagemaker/commands.test.ts index 35d05b47c86..4dab0c73b9b 100644 --- a/packages/core/src/test/awsService/sagemaker/commands.test.ts +++ b/packages/core/src/test/awsService/sagemaker/commands.test.ts @@ -401,7 +401,7 @@ describe('SageMaker Commands', () => { }) }) - describe('HyperPod connection with clusterArn', function () { + describe('HyperPod connection with eksClusterArn', function () { let mockDeeplinkConnect: sinon.SinonStub let mockIsRemoteWorkspace: sinon.SinonStub let deeplinkConnect: any @@ -417,7 +417,7 @@ describe('SageMaker Commands', () => { sandbox.replace(freshModule, 'deeplinkConnect', mockDeeplinkConnect) }) - it('should create session with underscores from HyperPod clusterArn', async function () { + it('should create session with underscores from HyperPod eksClusterArn', async function () { const ctx = { extensionContext: {}, } as any @@ -447,7 +447,7 @@ describe('SageMaker Commands', () => { } }) - it('should handle EKS clusterArn format', async function () { + it('should handle EKS eksClusterArn format', async function () { const ctx = { extensionContext: {}, } as any @@ -499,7 +499,7 @@ describe('SageMaker Commands', () => { } }) - it('should handle long component names by truncating', async function () { + it('should handle long component names by truncating to 253 chars', async function () { const ctx = { extensionContext: {}, } as any @@ -523,7 +523,7 @@ describe('SageMaker Commands', () => { const sessionArg = mockDeeplinkConnect.firstCall?.args[10] if (sessionArg) { - assert.ok(sessionArg.length <= 224, 'Session should not exceed max length') + assert.ok(sessionArg.length <= 253, 'Session should not exceed 253 char max length') } }) diff --git a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts index 31007519ace..8040c45edc4 100644 --- a/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts +++ b/packages/core/src/test/awsService/sagemaker/uriHandlers.test.ts @@ -85,7 +85,7 @@ describe('SageMaker URI handler', function () { return vscode.Uri.parse(`vscode://${VSCODE_EXTENSION_ID.awstoolkit}/connect/workspace?${query}`) } - it('calls deeplinkConnect with clusterArn for HyperPod connections', async function () { + it('calls deeplinkConnect with eksClusterArn for HyperPod connections', async function () { const params = { sessionId: 'session-123', streamUrl: 'wss://example.com/stream', @@ -93,7 +93,7 @@ describe('SageMaker URI handler', function () { 'cell-number': '5', workspaceName: 'my-workspace', namespace: 'default', - clusterArn: 'arn:aws:sagemaker:us-east-2:123456789012:cluster/my-cluster', + eksClusterArn: 'arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster', } const uri = createHyperPodUri(params) @@ -110,7 +110,7 @@ describe('SageMaker URI handler', function () { assert.deepStrictEqual(deeplinkConnectStub.firstCall.args[8], 'default') assert.deepStrictEqual( deeplinkConnectStub.firstCall.args[9], - 'arn:aws:sagemaker:us-east-2:123456789012:cluster/my-cluster' + 'arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster' ) }) diff --git a/packages/core/src/test/shared/clients/kubectlClient.test.ts b/packages/core/src/test/shared/clients/kubectlClient.test.ts index f98661c611a..a43dcbbe24d 100644 --- a/packages/core/src/test/shared/clients/kubectlClient.test.ts +++ b/packages/core/src/test/shared/clients/kubectlClient.test.ts @@ -164,57 +164,45 @@ describe('KubectlClient', function () { assert.strictEqual(result.url, 'https://test-url.com') }) - it('should throw error when presignedUrl is undefined', async function () { + it('should return presigned URL unmodified', async function () { + const presignedUrl = 'https://test-url.com?eksClusterArn=arn:aws:eks:us-east-1:123:cluster/test' const mockResponse = { response: {} as IncomingMessage, body: { status: { + workspaceConnectionUrl: presignedUrl, workspaceConnectionType: 'vscode-remote', }, }, } mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) - await assert.rejects(client.createWorkspaceConnection(mockDevSpace), /No workspace connection URL returned/) - }) - - it('should throw error when workspace connection creation fails', async function () { - mockK8sApi.createNamespacedCustomObject.rejects(new Error('Creation failed')) + const result = await client.createWorkspaceConnection(mockDevSpace) - await assert.rejects( - client.createWorkspaceConnection(mockDevSpace), - /Failed to create workspace connection/ - ) + assert.strictEqual(result.url, presignedUrl, 'Presigned URL should be returned unmodified') }) - it('should not modify URL when hyperpodCluster.clusterArn is undefined', async function () { - const originalUrl = - 'vscode://test/connect?sessionId=123&eksClusterArn=arn:aws:eks:us-east-2:123456789012:cluster/eks-cluster' + it('should throw error when presignedUrl is undefined', async function () { const mockResponse = { response: {} as IncomingMessage, body: { status: { - workspaceConnectionUrl: originalUrl, workspaceConnectionType: 'vscode-remote', }, }, } mockK8sApi.createNamespacedCustomObject.resolves(mockResponse) - // Create a client with undefined clusterArn - const mockHyperpodClusterNoArn = { - clusterName: 'test-cluster', - clusterArn: undefined as any, - status: 'InService', - regionCode: 'us-east-2', - } - const clientNoArn = new KubectlClient(mockEksCluster, mockHyperpodClusterNoArn) - ;(clientNoArn as any).k8sApi = mockK8sApi + await assert.rejects(client.createWorkspaceConnection(mockDevSpace), /No workspace connection URL returned/) + }) - const result = await clientNoArn.createWorkspaceConnection(mockDevSpace) + it('should throw error when workspace connection creation fails', async function () { + mockK8sApi.createNamespacedCustomObject.rejects(new Error('Creation failed')) - // URL should remain unchanged when clusterArn is undefined - assert.strictEqual(result.url, originalUrl) + await assert.rejects( + client.createWorkspaceConnection(mockDevSpace), + /Failed to create workspace connection/ + ) }) })