{
+ if (!this.sqlClient) {
+ try {
+ // Get the endpoint URL for the region
+ const endpoint = this.getSQLWorkbenchEndpoint(this.region)
+ this.logger.info(`Using SQL Workbench endpoint: ${endpoint}`)
+
+ if (this.connectionCredentialsProvider) {
+ // Create client with provided credentials
+ this.sqlClient = (await globals.sdkClientBuilder.createAwsService(
+ Service,
+ {
+ apiConfig: apiConfig,
+ region: this.region,
+ endpoint: endpoint,
+ credentialProvider: adaptConnectionCredentialsProvider(this.connectionCredentialsProvider),
+ } as ServiceConfigurationOptions,
+ undefined,
+ false
+ )) as SQLWorkbench
+ } else {
+ // Use the SDK client builder for default credentials
+ this.sqlClient = (await globals.sdkClientBuilder.createAwsService(
+ Service,
+ {
+ apiConfig: apiConfig,
+ region: this.region,
+ endpoint: endpoint,
+ } as ServiceConfigurationOptions,
+ undefined,
+ false
+ )) as SQLWorkbench
+ }
+
+ this.logger.debug('SQLWorkbenchClient: Successfully created SQL client')
+ } catch (err) {
+ this.logger.error('SQLWorkbenchClient: Failed to create SQL client: %s', err as Error)
+ throw err
+ }
+ }
+ return this.sqlClient
+ }
+}
diff --git a/packages/core/src/sagemakerunifiedstudio/shared/client/sqlworkbench.json b/packages/core/src/sagemakerunifiedstudio/shared/client/sqlworkbench.json
new file mode 100644
index 00000000000..e403ec34a88
--- /dev/null
+++ b/packages/core/src/sagemakerunifiedstudio/shared/client/sqlworkbench.json
@@ -0,0 +1,2102 @@
+{
+ "version": "2.0",
+ "metadata": {
+ "apiVersion": "2024-02-12",
+ "auth": ["aws.auth#sigv4"],
+ "endpointPrefix": "sqlworkbench",
+ "protocol": "rest-json",
+ "protocols": ["rest-json"],
+ "serviceFullName": "AmazonSQLWorkbench",
+ "serviceId": "SQLWorkbench",
+ "signatureVersion": "v4",
+ "signingName": "sqlworkbench",
+ "uid": "sqlworkbench-2024-02-12"
+ },
+ "operations": {
+ "CancelQueries": {
+ "name": "CancelQueries",
+ "http": {
+ "method": "POST",
+ "requestUri": "/database/cancelQueries",
+ "responseCode": 200
+ },
+ "input": { "shape": "CancelQueriesRequest" },
+ "output": { "shape": "CancelQueriesResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "CreateConnection": {
+ "name": "CreateConnection",
+ "http": {
+ "method": "PUT",
+ "requestUri": "/connections",
+ "responseCode": 200
+ },
+ "input": { "shape": "CreateConnectionRequest" },
+ "output": { "shape": "CreateConnectionResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "DeleteConnection": {
+ "name": "DeleteConnection",
+ "http": {
+ "method": "DELETE",
+ "requestUri": "/connections/{connectionId}",
+ "responseCode": 200
+ },
+ "input": { "shape": "DeleteConnectionRequest" },
+ "output": { "shape": "DeleteConnectionResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "ExecuteQuery": {
+ "name": "ExecuteQuery",
+ "http": {
+ "method": "POST",
+ "requestUri": "/database/executeQuery",
+ "responseCode": 200
+ },
+ "input": { "shape": "ExecuteQueryRequest" },
+ "output": { "shape": "ExecuteQueryResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "ExportQueryResults": {
+ "name": "ExportQueryResults",
+ "http": {
+ "method": "POST",
+ "requestUri": "/database/exportResults",
+ "responseCode": 200
+ },
+ "input": { "shape": "ExportQueryResultsRequest" },
+ "output": { "shape": "ExportQueryResultsResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "GetConnectableResources": {
+ "name": "GetConnectableResources",
+ "http": {
+ "method": "POST",
+ "requestUri": "/database/getConnectableResources",
+ "responseCode": 200
+ },
+ "input": { "shape": "GetConnectableResourcesRequest" },
+ "output": { "shape": "GetConnectableResourcesResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "GetConnection": {
+ "name": "GetConnection",
+ "http": {
+ "method": "GET",
+ "requestUri": "/connections/{connectionId}",
+ "responseCode": 200
+ },
+ "input": { "shape": "GetConnectionRequest" },
+ "output": { "shape": "GetConnectionResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "GetDatabaseConfigurations": {
+ "name": "GetDatabaseConfigurations",
+ "http": {
+ "method": "POST",
+ "requestUri": "/database/configurations",
+ "responseCode": 200
+ },
+ "input": { "shape": "GetDatabaseConfigurationsRequest" },
+ "output": { "shape": "GetDatabaseConfigurationsResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "GetQueryExecutionHistory": {
+ "name": "GetQueryExecutionHistory",
+ "http": {
+ "method": "POST",
+ "requestUri": "/queryExecutionHistory/details",
+ "responseCode": 200
+ },
+ "input": { "shape": "GetQueryExecutionHistoryRequest" },
+ "output": { "shape": "GetQueryExecutionHistoryResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "GetQueryResult": {
+ "name": "GetQueryResult",
+ "http": {
+ "method": "POST",
+ "requestUri": "/database/getQueryResults",
+ "responseCode": 200
+ },
+ "input": { "shape": "GetQueryResultRequest" },
+ "output": { "shape": "GetQueryResultResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "GetResources": {
+ "name": "GetResources",
+ "http": {
+ "method": "POST",
+ "requestUri": "/database/getResources",
+ "responseCode": 200
+ },
+ "input": { "shape": "GetResourcesRequest" },
+ "output": { "shape": "GetResourcesResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "GetTabStates": {
+ "name": "GetTabStates",
+ "http": {
+ "method": "POST",
+ "requestUri": "/tab/state",
+ "responseCode": 200
+ },
+ "input": { "shape": "GetTabStatesRequest" },
+ "output": { "shape": "GetTabStatesResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "ListQueryExecutionHistory": {
+ "name": "ListQueryExecutionHistory",
+ "http": {
+ "method": "POST",
+ "requestUri": "/queryExecutionHistory/list",
+ "responseCode": 200
+ },
+ "input": { "shape": "ListQueryExecutionHistoryRequest" },
+ "output": { "shape": "ListQueryExecutionHistoryResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "ListTagsForResource": {
+ "name": "ListTagsForResource",
+ "http": {
+ "method": "GET",
+ "requestUri": "/tags/{resourceArn}",
+ "responseCode": 200
+ },
+ "input": { "shape": "ListTagsForResourceRequest" },
+ "output": { "shape": "ListTagsForResourceResponse" },
+ "errors": [
+ { "shape": "BadRequestError" },
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "PollQueryExecutionEvents": {
+ "name": "PollQueryExecutionEvents",
+ "http": {
+ "method": "POST",
+ "requestUri": "/database/pollQueryExecutionEvents",
+ "responseCode": 200
+ },
+ "input": { "shape": "PollQueryExecutionEventsRequest" },
+ "output": { "shape": "PollQueryExecutionEventsResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "BadRequestError" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "TagResource": {
+ "name": "TagResource",
+ "http": {
+ "method": "POST",
+ "requestUri": "/tags/{resourceArn}",
+ "responseCode": 204
+ },
+ "input": { "shape": "TagResourceRequest" },
+ "output": { "shape": "TagResourceResponse" },
+ "errors": [
+ { "shape": "BadRequestError" },
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "UntagResource": {
+ "name": "UntagResource",
+ "http": {
+ "method": "DELETE",
+ "requestUri": "/tags/{resourceArn}",
+ "responseCode": 204
+ },
+ "input": { "shape": "UntagResourceRequest" },
+ "output": { "shape": "UntagResourceResponse" },
+ "errors": [
+ { "shape": "BadRequestError" },
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ],
+ "idempotent": true
+ },
+ "UpdateConnection": {
+ "name": "UpdateConnection",
+ "http": {
+ "method": "POST",
+ "requestUri": "/connections",
+ "responseCode": 200
+ },
+ "input": { "shape": "UpdateConnectionRequest" },
+ "output": { "shape": "UpdateConnectionResponse" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "ConflictException" },
+ { "shape": "InternalServerError" },
+ { "shape": "ValidationException" }
+ ]
+ },
+ "VerifyResourcesExistForTagris": {
+ "name": "VerifyResourcesExistForTagris",
+ "http": {
+ "method": "POST",
+ "requestUri": "/verifyResourcesExistForTagris",
+ "responseCode": 200
+ },
+ "input": { "shape": "TagrisVerifyResourcesExistInput" },
+ "output": { "shape": "TagrisVerifyResourcesExistOutput" },
+ "errors": [
+ { "shape": "ThrottlingException" },
+ { "shape": "InternalServerError" },
+ { "shape": "TagrisInvalidParameterException" },
+ { "shape": "TagrisAccessDeniedException" },
+ { "shape": "TagrisInvalidArnException" },
+ { "shape": "ResourceNotFoundException" },
+ { "shape": "TagrisInternalServiceException" },
+ { "shape": "ServiceQuotaExceededException" },
+ { "shape": "AccessDeniedException" },
+ { "shape": "TagrisPartialResourcesExistResultsException" },
+ { "shape": "TagrisThrottledException" },
+ { "shape": "ConflictException" },
+ { "shape": "ValidationException" }
+ ]
+ }
+ },
+ "shapes": {
+ "AccessDeniedException": {
+ "type": "structure",
+ "required": ["message"],
+ "members": {
+ "message": { "shape": "String" },
+ "code": { "shape": "ErrorCode" }
+ },
+ "error": {
+ "httpStatusCode": 403,
+ "senderFault": true
+ },
+ "exception": true
+ },
+ "AckIds": {
+ "type": "list",
+ "member": { "shape": "AckIdsMemberString" }
+ },
+ "AckIdsMemberString": {
+ "type": "string",
+ "max": 100,
+ "min": 0
+ },
+ "Arn": {
+ "type": "string",
+ "max": 1011,
+ "min": 20
+ },
+ "AvailableConnectionConfigurationOptions": {
+ "type": "list",
+ "member": { "shape": "AvailableConnectionConfigurationOptionsMemberString" }
+ },
+ "AvailableConnectionConfigurationOptionsMemberString": {
+ "type": "string",
+ "max": 50,
+ "min": 0
+ },
+ "BadRequestError": {
+ "type": "structure",
+ "required": ["message"],
+ "members": {
+ "message": { "shape": "String" },
+ "code": { "shape": "ErrorCode" }
+ },
+ "error": {
+ "httpStatusCode": 400,
+ "senderFault": true
+ },
+ "exception": true
+ },
+ "Boolean": {
+ "type": "boolean",
+ "box": true
+ },
+ "CancelQueriesRequest": {
+ "type": "structure",
+ "required": ["queryExecutionIds", "databaseType"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "queryExecutionIds": { "shape": "CancelQueriesRequestQueryExecutionIdsList" },
+ "databaseType": {
+ "shape": "DatabaseType",
+ "location": "querystring",
+ "locationName": "databaseType"
+ }
+ }
+ },
+ "CancelQueriesRequestQueryExecutionIdsList": {
+ "type": "list",
+ "member": { "shape": "CancelQueriesRequestQueryExecutionIdsListMemberString" },
+ "max": 100,
+ "min": 1
+ },
+ "CancelQueriesRequestQueryExecutionIdsListMemberString": {
+ "type": "string",
+ "max": 100,
+ "min": 1
+ },
+ "CancelQueriesResponse": {
+ "type": "structure",
+ "required": ["cancelQueryResponses"],
+ "members": {
+ "cancelQueryResponses": { "shape": "CancelQueryResponses" }
+ }
+ },
+ "CancelQueryResponse": {
+ "type": "structure",
+ "required": ["queryExecutionId"],
+ "members": {
+ "queryExecutionId": { "shape": "CancelQueryResponseQueryExecutionIdString" },
+ "queryCancellationStatus": { "shape": "QueryCancellationStatus" }
+ }
+ },
+ "CancelQueryResponseQueryExecutionIdString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "CancelQueryResponses": {
+ "type": "list",
+ "member": { "shape": "CancelQueryResponse" }
+ },
+ "ChildObjectTypes": {
+ "type": "list",
+ "member": { "shape": "ChildObjectTypesMemberString" }
+ },
+ "ChildObjectTypesMemberString": {
+ "type": "string",
+ "max": 50,
+ "min": 0
+ },
+ "Columns": {
+ "type": "list",
+ "member": { "shape": "QueryResultCellValue" }
+ },
+ "ConflictException": {
+ "type": "structure",
+ "required": ["message"],
+ "members": {
+ "message": { "shape": "String" },
+ "code": { "shape": "ErrorCode" }
+ },
+ "error": {
+ "httpStatusCode": 409,
+ "senderFault": true
+ },
+ "exception": true
+ },
+ "ConnectableResource": {
+ "type": "structure",
+ "required": ["displayName", "identifier", "childObjectTypes", "availableConnectionConfigurationOptions"],
+ "members": {
+ "displayName": { "shape": "ResourceDisplayName" },
+ "identifier": { "shape": "ResourceIdentifier" },
+ "type": { "shape": "ConnectableResourceTypeString" },
+ "unavailable": { "shape": "Boolean" },
+ "tooltipTranslationKey": { "shape": "ConnectableResourceTooltipTranslationKeyString" },
+ "childObjectTypes": { "shape": "ChildObjectTypes" },
+ "availableConnectionConfigurationOptions": { "shape": "AvailableConnectionConfigurationOptions" }
+ }
+ },
+ "ConnectableResourceTooltipTranslationKeyString": {
+ "type": "string",
+ "max": 50,
+ "min": 0
+ },
+ "ConnectableResourceTypeString": {
+ "type": "string",
+ "max": 50,
+ "min": 0
+ },
+ "ConnectableResourceTypes": {
+ "type": "list",
+ "member": { "shape": "ConnectableResourceTypesMemberString" }
+ },
+ "ConnectableResourceTypesMemberString": {
+ "type": "string",
+ "max": 50,
+ "min": 0
+ },
+ "ConnectableResources": {
+ "type": "list",
+ "member": { "shape": "ConnectableResource" }
+ },
+ "Connection": {
+ "type": "structure",
+ "members": {
+ "id": {
+ "shape": "String",
+ "documentation": "Id of the connection
"
+ },
+ "name": {
+ "shape": "ConnectionName",
+ "documentation": "Name of the connection
"
+ },
+ "authenticationType": {
+ "shape": "ConnectionAuthenticationTypes",
+ "documentation": "Number representing the type of authentication to use (2 = IAM, 3 = Username and Password). Today we only support the types 2 and 3
"
+ },
+ "secretArn": {
+ "shape": "String",
+ "documentation": "Secret that is linked to this connection
"
+ },
+ "databaseName": {
+ "shape": "DatabaseName",
+ "documentation": "Name of the database where the query is run
"
+ },
+ "clusterId": {
+ "shape": "String",
+ "documentation": "Id of the cluster of the connection
"
+ },
+ "dbUser": {
+ "shape": "DbUser",
+ "documentation": "User of the database
"
+ },
+ "isServerless": { "shape": "Boolean" },
+ "isProd": { "shape": "String" },
+ "isEnabled": { "shape": "String" },
+ "userSettings": { "shape": "UserSettings" },
+ "recordDate": { "shape": "String" },
+ "updatedDate": { "shape": "String" },
+ "tags": { "shape": "Tags" },
+ "databaseType": { "shape": "DatabaseType" },
+ "connectableResourceType": { "shape": "String" },
+ "connectableResourceIdentifier": { "shape": "ResourceIdentifier" }
+ }
+ },
+ "ConnectionAuthenticationTypes": {
+ "type": "string",
+ "enum": ["2", "3", "4", "5", "6", "7", "8"],
+ "sensitive": true
+ },
+ "ConnectionName": {
+ "type": "string",
+ "sensitive": true
+ },
+ "ConnectionProperties": {
+ "type": "map",
+ "key": { "shape": "ConnectionPropertyKey" },
+ "value": { "shape": "ConnectionPropertyValue" },
+ "max": 50,
+ "min": 1
+ },
+ "ConnectionPropertyKey": {
+ "type": "string",
+ "max": 1000,
+ "min": 1
+ },
+ "ConnectionPropertyValue": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "CreateConnectionRequest": {
+ "type": "structure",
+ "required": ["name", "databaseName", "authenticationType"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "name": {
+ "shape": "CreateConnectionRequestNameString",
+ "documentation": "Name of the connection
"
+ },
+ "databaseName": {
+ "shape": "CreateConnectionRequestDatabaseNameString",
+ "documentation": "Name of the database used for this connection
"
+ },
+ "authenticationType": {
+ "shape": "CreateConnectionRequestAuthenticationTypeEnum",
+ "documentation": "Number representing the type of authentication to use (2 = IAM, 3 = Username and Password, 4 = Federated connection)
"
+ },
+ "isProd": { "shape": "CreateConnectionRequestIsProdString" },
+ "userSettings": { "shape": "UserSettings" },
+ "secretArn": {
+ "shape": "CreateConnectionRequestSecretArnString",
+ "documentation": "secretArn for redshift cluster
"
+ },
+ "clusterId": {
+ "shape": "CreateConnectionRequestClusterIdString",
+ "documentation": "Id of the cluster used for this connection
"
+ },
+ "isServerless": {
+ "shape": "Boolean",
+ "documentation": "Is serverless connection
"
+ },
+ "dbUser": {
+ "shape": "DbUser",
+ "documentation": "User of the database used for this connection
"
+ },
+ "isStoreNewSecret": { "shape": "CreateConnectionRequestIsStoreNewSecretString" },
+ "username": {
+ "shape": "DbUser",
+ "documentation": "Username used in the Username_Password connection type
"
+ },
+ "password": {
+ "shape": "CreateConnectionRequestPasswordString",
+ "documentation": "Password of the user used for this connection
"
+ },
+ "tags": { "shape": "Tags" },
+ "host": {
+ "shape": "CreateConnectionRequestHostString",
+ "documentation": "Host address used for creating secret for Username_Password connection type
"
+ },
+ "secretName": { "shape": "CreateConnectionRequestSecretNameString" },
+ "description": { "shape": "CreateConnectionRequestDescriptionString" },
+ "databaseType": { "shape": "DatabaseType" },
+ "connectableResourceIdentifier": {
+ "shape": "CreateConnectionRequestConnectableResourceIdentifierString",
+ "documentation": "Id of the connectable resource used for this connection
"
+ },
+ "connectableResourceType": {
+ "shape": "CreateConnectionRequestConnectableResourceTypeString",
+ "documentation": "Type of the connectable resource used for this connection
"
+ }
+ }
+ },
+ "CreateConnectionRequestAuthenticationTypeEnum": {
+ "type": "string",
+ "enum": ["2", "3", "4", "5", "6", "7", "8"],
+ "max": 1,
+ "min": 1,
+ "sensitive": true
+ },
+ "CreateConnectionRequestClusterIdString": {
+ "type": "string",
+ "max": 63,
+ "min": 1
+ },
+ "CreateConnectionRequestConnectableResourceIdentifierString": {
+ "type": "string",
+ "max": 63,
+ "min": 1,
+ "sensitive": true
+ },
+ "CreateConnectionRequestConnectableResourceTypeString": {
+ "type": "string",
+ "max": 63,
+ "min": 1
+ },
+ "CreateConnectionRequestDatabaseNameString": {
+ "type": "string",
+ "max": 64,
+ "min": 1,
+ "sensitive": true
+ },
+ "CreateConnectionRequestDescriptionString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "CreateConnectionRequestHostString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "CreateConnectionRequestIsProdString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "CreateConnectionRequestIsStoreNewSecretString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "CreateConnectionRequestNameString": {
+ "type": "string",
+ "max": 512,
+ "min": 1,
+ "sensitive": true
+ },
+ "CreateConnectionRequestPasswordString": {
+ "type": "string",
+ "max": 64,
+ "min": 8,
+ "sensitive": true
+ },
+ "CreateConnectionRequestSecretArnString": {
+ "type": "string",
+ "max": 1000,
+ "min": 1
+ },
+ "CreateConnectionRequestSecretNameString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "CreateConnectionResponse": {
+ "type": "structure",
+ "members": {
+ "data": { "shape": "Connection" }
+ }
+ },
+ "DatabaseAuthenticationMethod": {
+ "type": "string",
+ "enum": ["USERNAME_PASSWORD", "TEMPORARY_CREDENTIALS_WITH_IAM"]
+ },
+ "DatabaseAuthenticationMethods": {
+ "type": "list",
+ "member": { "shape": "DatabaseAuthenticationMethod" }
+ },
+ "DatabaseAuthenticationOption": {
+ "type": "structure",
+ "required": ["connectableResourceType", "authenticationMethods"],
+ "members": {
+ "connectableResourceType": { "shape": "String" },
+ "authenticationMethods": { "shape": "DatabaseAuthenticationMethods" }
+ }
+ },
+ "DatabaseAuthenticationOptions": {
+ "type": "list",
+ "member": { "shape": "DatabaseAuthenticationOption" }
+ },
+ "DatabaseConfiguration": {
+ "type": "structure",
+ "required": [
+ "databaseType",
+ "authenticationOptions",
+ "connectableResourceTypes",
+ "sessionSupported",
+ "eventAcknowledgementSupported",
+ "appendingLimitToQuerySupported",
+ "queryStatsSupported"
+ ],
+ "members": {
+ "databaseType": { "shape": "DatabaseType" },
+ "authenticationOptions": { "shape": "DatabaseAuthenticationOptions" },
+ "connectableResourceTypes": { "shape": "ConnectableResourceTypes" },
+ "sessionSupported": { "shape": "Boolean" },
+ "eventAcknowledgementSupported": { "shape": "Boolean" },
+ "appendingLimitToQuerySupported": { "shape": "Boolean" },
+ "queryStatsSupported": { "shape": "Boolean" }
+ }
+ },
+ "DatabaseConfigurations": {
+ "type": "list",
+ "member": { "shape": "DatabaseConfiguration" }
+ },
+ "DatabaseConnectionAccountSettings": {
+ "type": "structure",
+ "members": {
+ "masterKeyArn": { "shape": "KmsKeyArn" }
+ }
+ },
+ "DatabaseConnectionConfiguration": {
+ "type": "structure",
+ "required": ["id", "type", "databaseType", "connectableResourceIdentifier", "connectableResourceType"],
+ "members": {
+ "id": { "shape": "DatabaseConnectionConfigurationIdString" },
+ "type": { "shape": "DatabaseIntegrationConnectionAuthenticationTypes" },
+ "auth": { "shape": "DatabaseConnectionConfigurationAuth" },
+ "databaseType": { "shape": "DatabaseType" },
+ "connectableResourceIdentifier": { "shape": "ResourceIdentifier" },
+ "connectableResourceType": { "shape": "DatabaseConnectionConfigurationConnectableResourceTypeString" },
+ "database": { "shape": "DatabaseName" }
+ }
+ },
+ "DatabaseConnectionConfigurationAuth": {
+ "type": "structure",
+ "members": {
+ "secretArn": { "shape": "SecretKeyArn" },
+ "username": { "shape": "DatabaseConnectionConfigurationAuthUsernameString" },
+ "password": { "shape": "DatabaseConnectionConfigurationAuthPasswordString" }
+ }
+ },
+ "DatabaseConnectionConfigurationAuthPasswordString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0,
+ "sensitive": true
+ },
+ "DatabaseConnectionConfigurationAuthUsernameString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0,
+ "sensitive": true
+ },
+ "DatabaseConnectionConfigurationConnectableResourceTypeString": {
+ "type": "string",
+ "max": 50,
+ "min": 0
+ },
+ "DatabaseConnectionConfigurationIdString": {
+ "type": "string",
+ "max": 2048,
+ "min": 32
+ },
+ "DatabaseIntegrationConnectionAuthenticationTypes": {
+ "type": "string",
+ "enum": ["4", "5", "6", "8"],
+ "sensitive": true
+ },
+ "DatabaseName": {
+ "type": "string",
+ "max": 150,
+ "min": 0,
+ "sensitive": true
+ },
+ "DatabaseType": {
+ "type": "string",
+ "enum": ["REDSHIFT", "ATHENA"]
+ },
+ "DbUser": {
+ "type": "string",
+ "max": 127,
+ "min": 1,
+ "pattern": "[a-zA-Z0-9_][a-zA-Z_0-9+.@$-]*",
+ "sensitive": true
+ },
+ "DeleteConnectionRequest": {
+ "type": "structure",
+ "required": ["connectionId"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "connectionId": {
+ "shape": "DeleteConnectionRequestConnectionIdString",
+ "documentation": "Id of connection to delete
",
+ "location": "uri",
+ "locationName": "connectionId"
+ }
+ }
+ },
+ "DeleteConnectionRequestConnectionIdString": {
+ "type": "string",
+ "max": 1000,
+ "min": 1
+ },
+ "DeleteConnectionResponse": {
+ "type": "structure",
+ "members": {}
+ },
+ "ErrorCode": {
+ "type": "string",
+ "enum": ["QUERY_EXECUTION_NOT_FOUND", "QUERY_EXECUTION_ACCESS_DENIED"]
+ },
+ "ExecuteQueryRequest": {
+ "type": "structure",
+ "required": ["query", "queryExecutionType", "queryResponseDeliveryType", "maxItems"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "connectionId": { "shape": "ExecuteQueryRequestConnectionIdString" },
+ "databaseType": {
+ "shape": "DatabaseType",
+ "location": "querystring",
+ "locationName": "databaseType"
+ },
+ "connection": { "shape": "DatabaseConnectionConfiguration" },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "tabId": { "shape": "ExecuteQueryRequestTabIdString" },
+ "executionContext": { "shape": "ExecuteQueryRequestExecutionContextList" },
+ "query": { "shape": "ExecuteQueryRequestQueryString" },
+ "queryExecutionType": { "shape": "QueryExecutionType" },
+ "sessionId": { "shape": "ExecuteQueryRequestSessionIdString" },
+ "queryResponseDeliveryType": { "shape": "QueryResponseDeliveryType" },
+ "maxItems": { "shape": "ExecuteQueryRequestMaxItemsInteger" },
+ "limitQueryResults": { "shape": "ExecuteQueryRequestLimitQueryResultsInteger" },
+ "isExplain": { "shape": "Boolean" },
+ "ignoreHistory": { "shape": "Boolean" },
+ "timeoutMillis": { "shape": "ExecuteQueryRequestTimeoutMillisInteger" }
+ }
+ },
+ "ExecuteQueryRequestConnectionIdString": {
+ "type": "string",
+ "max": 2048,
+ "min": 32
+ },
+ "ExecuteQueryRequestExecutionContextList": {
+ "type": "list",
+ "member": { "shape": "ParentResource" },
+ "max": 100,
+ "min": 0
+ },
+ "ExecuteQueryRequestLimitQueryResultsInteger": {
+ "type": "integer",
+ "box": true,
+ "max": 1000,
+ "min": 0
+ },
+ "ExecuteQueryRequestMaxItemsInteger": {
+ "type": "integer",
+ "box": true,
+ "max": 100,
+ "min": 20
+ },
+ "ExecuteQueryRequestQueryString": {
+ "type": "string",
+ "max": 1000000,
+ "min": 0,
+ "sensitive": true
+ },
+ "ExecuteQueryRequestSessionIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 0
+ },
+ "ExecuteQueryRequestTabIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 1
+ },
+ "ExecuteQueryRequestTimeoutMillisInteger": {
+ "type": "integer",
+ "box": true,
+ "max": 120000,
+ "min": 0
+ },
+ "ExecuteQueryResponse": {
+ "type": "structure",
+ "required": ["queryExecutions"],
+ "members": {
+ "sessionId": { "shape": "ExecuteQueryResponseSessionIdString" },
+ "queryExecutions": { "shape": "QueryExecutions" },
+ "statusCode": {
+ "shape": "statusCode",
+ "location": "statusCode"
+ }
+ }
+ },
+ "ExecuteQueryResponseSessionIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 0
+ },
+ "ExportQueryResultsRequest": {
+ "type": "structure",
+ "required": ["queryExecutionId", "databaseType"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "queryExecutionId": { "shape": "ExportQueryResultsRequestQueryExecutionIdString" },
+ "databaseType": {
+ "shape": "DatabaseType",
+ "location": "querystring",
+ "locationName": "databaseType"
+ },
+ "fileType": { "shape": "FileType" }
+ }
+ },
+ "ExportQueryResultsRequestQueryExecutionIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 1
+ },
+ "ExportQueryResultsResponse": {
+ "type": "structure",
+ "required": ["queryResult", "contentType", "fileName"],
+ "members": {
+ "queryResult": { "shape": "StreamingBlob" },
+ "contentType": {
+ "shape": "String",
+ "location": "header",
+ "locationName": "Content-Type"
+ },
+ "fileName": {
+ "shape": "String",
+ "location": "header",
+ "locationName": "Content-Disposition"
+ }
+ },
+ "payload": "queryResult"
+ },
+ "FileType": {
+ "type": "string",
+ "enum": ["JSON", "CSV"]
+ },
+ "FullQueryText": {
+ "type": "string",
+ "max": 1000000,
+ "min": 0,
+ "sensitive": true
+ },
+ "GetConnectableResourcesRequest": {
+ "type": "structure",
+ "required": ["type", "maxItems", "databaseType"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "type": { "shape": "GetConnectableResourcesRequestTypeString" },
+ "maxItems": { "shape": "GetConnectableResourcesRequestMaxItemsInteger" },
+ "pageToken": { "shape": "PageToken" },
+ "databaseType": {
+ "shape": "DatabaseType",
+ "location": "querystring",
+ "locationName": "databaseType"
+ }
+ }
+ },
+ "GetConnectableResourcesRequestMaxItemsInteger": {
+ "type": "integer",
+ "box": true,
+ "max": 50,
+ "min": 20
+ },
+ "GetConnectableResourcesRequestTypeString": {
+ "type": "string",
+ "max": 150,
+ "min": 0
+ },
+ "GetConnectableResourcesResponse": {
+ "type": "structure",
+ "required": ["connectableResources"],
+ "members": {
+ "connectableResources": { "shape": "ConnectableResources" },
+ "nextToken": { "shape": "String" }
+ }
+ },
+ "GetConnectionRequest": {
+ "type": "structure",
+ "required": ["connectionId"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "connectionId": {
+ "shape": "GetConnectionRequestConnectionIdString",
+ "documentation": "Id of connection to delete
",
+ "location": "uri",
+ "locationName": "connectionId"
+ }
+ }
+ },
+ "GetConnectionRequestConnectionIdString": {
+ "type": "string",
+ "max": 1000,
+ "min": 1
+ },
+ "GetConnectionResponse": {
+ "type": "structure",
+ "members": {
+ "data": { "shape": "Connection" }
+ }
+ },
+ "GetDatabaseConfigurationsRequest": {
+ "type": "structure",
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" }
+ }
+ },
+ "GetDatabaseConfigurationsResponse": {
+ "type": "structure",
+ "members": {
+ "configurations": { "shape": "DatabaseConfigurations" }
+ }
+ },
+ "GetQueryExecutionHistoryRequest": {
+ "type": "structure",
+ "required": ["queryExecutionId"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "queryExecutionId": { "shape": "GetQueryExecutionHistoryRequestQueryExecutionIdString" },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" }
+ }
+ },
+ "GetQueryExecutionHistoryRequestQueryExecutionIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 1
+ },
+ "GetQueryExecutionHistoryResponse": {
+ "type": "structure",
+ "members": {
+ "id": { "shape": "String" },
+ "querySourceId": { "shape": "String" },
+ "queryStartTime": { "shape": "Long" },
+ "queryEndTime": { "shape": "Long" },
+ "status": { "shape": "QueryExecutionStatus" },
+ "queryText": { "shape": "FullQueryText" },
+ "serializedMetadata": { "shape": "SerializedMetadata" },
+ "serializedQueryStats": { "shape": "SerializedQueryStats" },
+ "databaseType": { "shape": "DatabaseType" }
+ }
+ },
+ "GetQueryResultRequest": {
+ "type": "structure",
+ "required": ["queryExecutionId", "databaseType"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "queryExecutionId": { "shape": "GetQueryResultRequestQueryExecutionIdString" },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "pageToken": { "shape": "PageToken" },
+ "databaseType": {
+ "shape": "DatabaseType",
+ "location": "querystring",
+ "locationName": "databaseType"
+ },
+ "pageSize": { "shape": "GetQueryResultRequestPageSizeInteger" }
+ }
+ },
+ "GetQueryResultRequestPageSizeInteger": {
+ "type": "integer",
+ "box": true,
+ "min": 0
+ },
+ "GetQueryResultRequestQueryExecutionIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 1
+ },
+ "GetQueryResultResponse": {
+ "type": "structure",
+ "members": {
+ "queryResult": { "shape": "QueryResult" },
+ "nextToken": { "shape": "String" },
+ "previousToken": { "shape": "String" }
+ }
+ },
+ "GetResourcesRequest": {
+ "type": "structure",
+ "required": ["parents", "type", "maxItems"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "connectionId": { "shape": "GetResourcesRequestConnectionIdString" },
+ "databaseType": {
+ "shape": "DatabaseType",
+ "location": "querystring",
+ "locationName": "databaseType"
+ },
+ "connection": { "shape": "DatabaseConnectionConfiguration" },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "parents": { "shape": "ParentResources" },
+ "type": { "shape": "GetResourcesRequestTypeString" },
+ "maxItems": { "shape": "GetResourcesRequestMaxItemsInteger" },
+ "pageToken": { "shape": "PageToken" },
+ "forceRefresh": { "shape": "Boolean" },
+ "forceRefreshRecursive": { "shape": "Boolean" }
+ }
+ },
+ "GetResourcesRequestConnectionIdString": {
+ "type": "string",
+ "max": 2048,
+ "min": 32
+ },
+ "GetResourcesRequestMaxItemsInteger": {
+ "type": "integer",
+ "box": true,
+ "max": 100,
+ "min": 20
+ },
+ "GetResourcesRequestTypeString": {
+ "type": "string",
+ "max": 150,
+ "min": 0
+ },
+ "GetResourcesResponse": {
+ "type": "structure",
+ "members": {
+ "resources": { "shape": "Resources" },
+ "nextToken": { "shape": "String" },
+ "statusCode": {
+ "shape": "statusCode",
+ "location": "statusCode"
+ },
+ "connectionProperties": { "shape": "ConnectionProperties" }
+ }
+ },
+ "GetTabStatesRequest": {
+ "type": "structure",
+ "required": ["tabId"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "tabId": { "shape": "String" }
+ }
+ },
+ "GetTabStatesResponse": {
+ "type": "structure",
+ "required": ["queryExecutionStates"],
+ "members": {
+ "queryExecutionStates": { "shape": "QueryExecutionStates" },
+ "sessionId": { "shape": "String" }
+ }
+ },
+ "Integer": {
+ "type": "integer",
+ "box": true
+ },
+ "InternalServerError": {
+ "type": "structure",
+ "required": ["message"],
+ "members": {
+ "message": { "shape": "String" },
+ "code": { "shape": "ErrorCode" }
+ },
+ "error": { "httpStatusCode": 500 },
+ "exception": true,
+ "fault": true
+ },
+ "KmsKeyArn": {
+ "type": "string",
+ "max": 1000,
+ "min": 0,
+ "pattern": "arn:.*"
+ },
+ "ListQueryExecutionHistoryRequest": {
+ "type": "structure",
+ "required": ["maxItems"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "maxItems": { "shape": "ListQueryExecutionHistoryRequestMaxItemsInteger" },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "pageToken": { "shape": "ListQueryExecutionHistoryRequestPageTokenString" },
+ "querySourceId": { "shape": "ListQueryExecutionHistoryRequestQuerySourceIdString" },
+ "databaseType": { "shape": "DatabaseType" },
+ "status": { "shape": "QueryExecutionStatus" },
+ "startTime": { "shape": "QueryHistoryTimestamp" },
+ "endTime": { "shape": "QueryHistoryTimestamp" },
+ "containsText": { "shape": "ListQueryExecutionHistoryRequestContainsTextString" }
+ }
+ },
+ "ListQueryExecutionHistoryRequestContainsTextString": {
+ "type": "string",
+ "max": 100,
+ "min": 0
+ },
+ "ListQueryExecutionHistoryRequestMaxItemsInteger": {
+ "type": "integer",
+ "box": true,
+ "max": 100,
+ "min": 1
+ },
+ "ListQueryExecutionHistoryRequestPageTokenString": {
+ "type": "string",
+ "max": 10000,
+ "min": 0
+ },
+ "ListQueryExecutionHistoryRequestQuerySourceIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 0
+ },
+ "ListQueryExecutionHistoryResponse": {
+ "type": "structure",
+ "required": ["items"],
+ "members": {
+ "items": { "shape": "QueryExecutionHistoryPreviews" },
+ "nextToken": { "shape": "ListQueryExecutionHistoryResponseNextTokenString" }
+ }
+ },
+ "ListQueryExecutionHistoryResponseNextTokenString": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "ListTagsForResourceRequest": {
+ "type": "structure",
+ "required": ["resourceArn"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "resourceArn": {
+ "shape": "Arn",
+ "location": "uri",
+ "locationName": "resourceArn"
+ }
+ }
+ },
+ "ListTagsForResourceResponse": {
+ "type": "structure",
+ "required": ["tags"],
+ "members": {
+ "tags": { "shape": "Tags" }
+ }
+ },
+ "Long": {
+ "type": "long",
+ "box": true
+ },
+ "PageToken": {
+ "type": "string",
+ "max": 1000,
+ "min": 0
+ },
+ "ParentResource": {
+ "type": "structure",
+ "required": ["parentId", "parentType"],
+ "members": {
+ "parentId": { "shape": "ParentResourceParentIdString" },
+ "parentType": { "shape": "ParentResourceParentTypeString" }
+ }
+ },
+ "ParentResourceParentIdString": {
+ "type": "string",
+ "max": 1000,
+ "min": 1,
+ "sensitive": true
+ },
+ "ParentResourceParentTypeString": {
+ "type": "string",
+ "max": 100,
+ "min": 1
+ },
+ "ParentResources": {
+ "type": "list",
+ "member": { "shape": "ParentResource" }
+ },
+ "PollQueryExecutionEventsRequest": {
+ "type": "structure",
+ "required": ["queryExecutionIds", "databaseType"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "queryExecutionIds": { "shape": "PollQueryExecutionEventsRequestQueryExecutionIdsList" },
+ "accountSettings": { "shape": "DatabaseConnectionAccountSettings" },
+ "databaseType": {
+ "shape": "DatabaseType",
+ "location": "querystring",
+ "locationName": "databaseType"
+ },
+ "ackIds": { "shape": "AckIds" }
+ }
+ },
+ "PollQueryExecutionEventsRequestQueryExecutionIdsList": {
+ "type": "list",
+ "member": { "shape": "PollQueryExecutionEventsRequestQueryExecutionIdsListMemberString" },
+ "max": 100,
+ "min": 1
+ },
+ "PollQueryExecutionEventsRequestQueryExecutionIdsListMemberString": {
+ "type": "string",
+ "max": 100,
+ "min": 1
+ },
+ "PollQueryExecutionEventsResponse": {
+ "type": "structure",
+ "members": {
+ "events": { "shape": "QueryExecutionEvents" }
+ }
+ },
+ "QueryCancellationStatus": {
+ "type": "string",
+ "enum": ["CANCELLED", "DOES_NOT_EXISTS", "ALREADY_FINISHED", "CANCELLATION_FAILED"]
+ },
+ "QueryExecution": {
+ "type": "structure",
+ "required": ["queryExecutionId"],
+ "members": {
+ "queryExecutionStatus": { "shape": "QueryExecutionStatus" },
+ "queryExecutionId": { "shape": "QueryExecutionQueryExecutionIdString" },
+ "queryResult": { "shape": "QueryResult" },
+ "queryText": { "shape": "QueryText" }
+ }
+ },
+ "QueryExecutionEvent": {
+ "type": "structure",
+ "required": ["queryExecutionEventType", "queryExecutionId"],
+ "members": {
+ "queryExecutionEventType": { "shape": "QueryExecutionEventType" },
+ "queryExecutionId": { "shape": "QueryExecutionEventQueryExecutionIdString" },
+ "queryExecutionStatus": { "shape": "QueryExecutionStatus" },
+ "queryResult": { "shape": "QueryResult" },
+ "nextToken": { "shape": "String" },
+ "ackId": { "shape": "String" }
+ }
+ },
+ "QueryExecutionEventQueryExecutionIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 0
+ },
+ "QueryExecutionEventType": {
+ "type": "string",
+ "enum": ["QUERY_EXECUTION_STATUS", "QUERY_EXECUTION_RESULT"]
+ },
+ "QueryExecutionEvents": {
+ "type": "list",
+ "member": { "shape": "QueryExecutionEvent" }
+ },
+ "QueryExecutionHistoryPreview": {
+ "type": "structure",
+ "members": {
+ "id": { "shape": "String" },
+ "querySourceId": { "shape": "String" },
+ "queryStartTime": { "shape": "Long" },
+ "queryEndTime": { "shape": "Long" },
+ "status": { "shape": "QueryExecutionStatus" },
+ "queryTextPreview": { "shape": "QueryTextPreview" },
+ "serializedMetadata": { "shape": "SerializedMetadata" },
+ "databaseType": { "shape": "DatabaseType" }
+ }
+ },
+ "QueryExecutionHistoryPreviews": {
+ "type": "list",
+ "member": { "shape": "QueryExecutionHistoryPreview" }
+ },
+ "QueryExecutionQueryExecutionIdString": {
+ "type": "string",
+ "max": 100,
+ "min": 0
+ },
+ "QueryExecutionState": {
+ "type": "structure",
+ "required": ["queryExecutionId", "status", "databaseType"],
+ "members": {
+ "queryExecutionId": { "shape": "String" },
+ "status": { "shape": "String" },
+ "databaseType": { "shape": "DatabaseType" }
+ }
+ },
+ "QueryExecutionStates": {
+ "type": "list",
+ "member": { "shape": "QueryExecutionState" }
+ },
+ "QueryExecutionStatus": {
+ "type": "string",
+ "enum": ["SCHEDULED", "RUNNING", "FAILED", "CANCELLED", "FINISHED"]
+ },
+ "QueryExecutionType": {
+ "type": "string",
+ "enum": ["PERSIST_SESSION", "NO_SESSION"]
+ },
+ "QueryExecutionWarning": {
+ "type": "structure",
+ "members": {
+ "message": { "shape": "QueryExecutionWarningMessage" },
+ "level": { "shape": "QueryExecutionWarningLevel" }
+ }
+ },
+ "QueryExecutionWarningLevel": {
+ "type": "string",
+ "enum": ["INFO", "WARNING"]
+ },
+ "QueryExecutionWarningMessage": {
+ "type": "string",
+ "max": 1000,
+ "min": 0,
+ "sensitive": true
+ },
+ "QueryExecutionWarnings": {
+ "type": "list",
+ "member": { "shape": "QueryExecutionWarning" }
+ },
+ "QueryExecutions": {
+ "type": "list",
+ "member": { "shape": "QueryExecution" }
+ },
+ "QueryHistoryTimestamp": {
+ "type": "long",
+ "box": true
+ },
+ "QueryResponseDeliveryType": {
+ "type": "string",
+ "enum": ["SYNC", "ASYNC"]
+ },
+ "QueryResult": {
+ "type": "structure",
+ "members": {
+ "queryExecutionStatus": { "shape": "QueryExecutionStatus" },
+ "headers": { "shape": "QueryResultHeaders" },
+ "rows": { "shape": "Rows" },
+ "affectedRows": { "shape": "Integer" },
+ "totalRowCount": { "shape": "Integer" },
+ "elapsedTime": { "shape": "Long" },
+ "errorMessage": { "shape": "QueryResultErrorMessage" },
+ "errorPosition": { "shape": "Integer" },
+ "queryResultWarningCode": { "shape": "QueryResultQueryResultWarningCodeString" },
+ "warnings": { "shape": "QueryExecutionWarnings" },
+ "queryExecutionId": { "shape": "String" },
+ "sessionId": { "shape": "String" },
+ "queryText": { "shape": "QueryText" },
+ "statementType": { "shape": "StatementType" },
+ "serializedMetadata": { "shape": "SerializedMetadata" },
+ "connectionProperties": { "shape": "ConnectionProperties" }
+ }
+ },
+ "QueryResultCellType": {
+ "type": "string",
+ "enum": ["STRING", "BOOLEAN", "INTEGER", "BIG_INTEGER", "FLOAT", "BIG_DECIMAL", "DATE", "TIME", "DATETIME"]
+ },
+ "QueryResultCellValue": {
+ "type": "string",
+ "sensitive": true
+ },
+ "QueryResultErrorMessage": {
+ "type": "string",
+ "max": 1000,
+ "min": 0,
+ "sensitive": true
+ },
+ "QueryResultHeader": {
+ "type": "structure",
+ "required": ["displayName", "type"],
+ "members": {
+ "displayName": { "shape": "QueryResultHeaderDisplayName" },
+ "type": { "shape": "QueryResultCellType" }
+ }
+ },
+ "QueryResultHeaderDisplayName": {
+ "type": "string",
+ "sensitive": true
+ },
+ "QueryResultHeaders": {
+ "type": "list",
+ "member": { "shape": "QueryResultHeader" }
+ },
+ "QueryResultQueryResultWarningCodeString": {
+ "type": "string",
+ "max": 100,
+ "min": 0
+ },
+ "QueryText": {
+ "type": "string",
+ "sensitive": true
+ },
+ "QueryTextPreview": {
+ "type": "string",
+ "max": 150,
+ "min": 0,
+ "sensitive": true
+ },
+ "Resource": {
+ "type": "structure",
+ "required": ["displayName", "identifier", "childObjectTypes"],
+ "members": {
+ "displayName": { "shape": "ResourceDisplayName" },
+ "identifier": { "shape": "ResourceIdentifier" },
+ "type": { "shape": "ResourceTypeString" },
+ "unavailable": { "shape": "Boolean" },
+ "tooltipTranslationKey": { "shape": "ResourceTooltipTranslationKeyString" },
+ "childObjectTypes": { "shape": "ChildObjectTypes" },
+ "allowedActions": { "shape": "ResourceActions" },
+ "resourceMetadata": { "shape": "ResourceMetadataItems" }
+ }
+ },
+ "ResourceAction": {
+ "type": "string",
+ "enum": ["Drop", "Truncate", "GenerateDefinition", "GenerateSelectQuery"]
+ },
+ "ResourceActions": {
+ "type": "list",
+ "member": { "shape": "ResourceAction" }
+ },
+ "ResourceDisplayName": {
+ "type": "string",
+ "max": 150,
+ "min": 0,
+ "sensitive": true
+ },
+ "ResourceIdentifier": {
+ "type": "string",
+ "max": 150,
+ "min": 0,
+ "sensitive": true
+ },
+ "ResourceMetadata": {
+ "type": "structure",
+ "members": {
+ "key": { "shape": "String" },
+ "value": { "shape": "String" }
+ }
+ },
+ "ResourceMetadataItems": {
+ "type": "list",
+ "member": { "shape": "ResourceMetadata" }
+ },
+ "ResourceNotFoundException": {
+ "type": "structure",
+ "required": ["message"],
+ "members": {
+ "message": { "shape": "String" },
+ "code": { "shape": "ErrorCode" }
+ },
+ "error": {
+ "httpStatusCode": 404,
+ "senderFault": true
+ },
+ "exception": true
+ },
+ "ResourceTooltipTranslationKeyString": {
+ "type": "string",
+ "max": 50,
+ "min": 0
+ },
+ "ResourceTypeString": {
+ "type": "string",
+ "max": 50,
+ "min": 0
+ },
+ "Resources": {
+ "type": "list",
+ "member": { "shape": "Resource" }
+ },
+ "Row": {
+ "type": "structure",
+ "members": {
+ "row": { "shape": "Columns" }
+ }
+ },
+ "Rows": {
+ "type": "list",
+ "member": { "shape": "Row" }
+ },
+ "SecretKeyArn": {
+ "type": "string",
+ "max": 1000,
+ "min": 0,
+ "pattern": "arn:.*"
+ },
+ "SerializedMetadata": {
+ "type": "string",
+ "max": 1000000,
+ "min": 0,
+ "sensitive": true
+ },
+ "SerializedQueryStats": {
+ "type": "string",
+ "max": 1000000,
+ "min": 0,
+ "sensitive": true
+ },
+ "ServiceQuotaExceededException": {
+ "type": "structure",
+ "required": ["message"],
+ "members": {
+ "message": { "shape": "String" },
+ "code": { "shape": "ErrorCode" }
+ },
+ "error": {
+ "httpStatusCode": 402,
+ "senderFault": true
+ },
+ "exception": true
+ },
+ "SqlworkbenchSource": {
+ "type": "string",
+ "enum": ["SUS", "RQEV2"]
+ },
+ "StatementType": {
+ "type": "string",
+ "enum": ["DQL", "DML", "DDL", "DCL", "Utility"]
+ },
+ "StreamingBlob": {
+ "type": "blob",
+ "streaming": true
+ },
+ "String": { "type": "string" },
+ "TagKey": {
+ "type": "string",
+ "max": 128,
+ "min": 1,
+ "pattern": "([\\p{L}\\p{Z}\\p{N}_.:/=+\\-@]*)"
+ },
+ "TagKeyList": {
+ "type": "list",
+ "member": { "shape": "TagKey" },
+ "max": 6500,
+ "min": 1
+ },
+ "TagResourceRequest": {
+ "type": "structure",
+ "required": ["resourceArn", "tags"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "resourceArn": {
+ "shape": "Arn",
+ "location": "uri",
+ "locationName": "resourceArn"
+ },
+ "tags": { "shape": "Tags" }
+ }
+ },
+ "TagResourceResponse": {
+ "type": "structure",
+ "members": {}
+ },
+ "TagValue": {
+ "type": "string",
+ "max": 256,
+ "min": 0,
+ "pattern": "([\\p{L}\\p{Z}\\p{N}_.:/=+\\-@]*)"
+ },
+ "TagrisAccessDeniedException": {
+ "type": "structure",
+ "members": {
+ "message": { "shape": "TagrisExceptionMessage" }
+ },
+ "exception": true
+ },
+ "TagrisAccountId": {
+ "type": "string",
+ "max": 12,
+ "min": 12
+ },
+ "TagrisAmazonResourceName": {
+ "type": "string",
+ "max": 1011,
+ "min": 1
+ },
+ "TagrisExceptionMessage": {
+ "type": "string",
+ "max": 2048,
+ "min": 0
+ },
+ "TagrisInternalId": {
+ "type": "string",
+ "max": 64,
+ "min": 0
+ },
+ "TagrisInternalServiceException": {
+ "type": "structure",
+ "members": {
+ "message": { "shape": "TagrisExceptionMessage" }
+ },
+ "exception": true,
+ "fault": true
+ },
+ "TagrisInvalidArnException": {
+ "type": "structure",
+ "members": {
+ "message": { "shape": "TagrisExceptionMessage" },
+ "sweepListItem": { "shape": "TagrisSweepListItem" }
+ },
+ "exception": true
+ },
+ "TagrisInvalidParameterException": {
+ "type": "structure",
+ "members": {
+ "message": { "shape": "TagrisExceptionMessage" }
+ },
+ "exception": true
+ },
+ "TagrisPartialResourcesExistResultsException": {
+ "type": "structure",
+ "members": {
+ "message": { "shape": "TagrisExceptionMessage" },
+ "resourceExistenceInformation": { "shape": "TagrisSweepListResult" }
+ },
+ "exception": true
+ },
+ "TagrisStatus": {
+ "type": "string",
+ "enum": ["ACTIVE", "NOT_ACTIVE"]
+ },
+ "TagrisSweepList": {
+ "type": "list",
+ "member": { "shape": "TagrisSweepListItem" }
+ },
+ "TagrisSweepListItem": {
+ "type": "structure",
+ "members": {
+ "TagrisAccountId": { "shape": "TagrisAccountId" },
+ "TagrisAmazonResourceName": { "shape": "TagrisAmazonResourceName" },
+ "TagrisInternalId": { "shape": "TagrisInternalId" },
+ "TagrisVersion": { "shape": "TagrisVersion" }
+ }
+ },
+ "TagrisSweepListResult": {
+ "type": "map",
+ "key": { "shape": "TagrisAmazonResourceName" },
+ "value": { "shape": "TagrisStatus" }
+ },
+ "TagrisThrottledException": {
+ "type": "structure",
+ "members": {
+ "message": { "shape": "TagrisExceptionMessage" }
+ },
+ "exception": true
+ },
+ "TagrisVerifyResourcesExistInput": {
+ "type": "structure",
+ "required": ["TagrisSweepList"],
+ "members": {
+ "TagrisSweepList": { "shape": "TagrisSweepList" }
+ }
+ },
+ "TagrisVerifyResourcesExistOutput": {
+ "type": "structure",
+ "required": ["TagrisSweepListResult"],
+ "members": {
+ "TagrisSweepListResult": { "shape": "TagrisSweepListResult" }
+ }
+ },
+ "TagrisVersion": { "type": "long" },
+ "Tags": {
+ "type": "map",
+ "key": { "shape": "TagKey" },
+ "value": { "shape": "TagValue" },
+ "max": 50,
+ "min": 1
+ },
+ "ThrottlingException": {
+ "type": "structure",
+ "required": ["message"],
+ "members": {
+ "message": { "shape": "String" },
+ "code": { "shape": "ErrorCode" }
+ },
+ "error": {
+ "httpStatusCode": 429,
+ "senderFault": true
+ },
+ "exception": true
+ },
+ "UntagResourceRequest": {
+ "type": "structure",
+ "required": ["resourceArn", "tagKeys"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "resourceArn": {
+ "shape": "Arn",
+ "location": "uri",
+ "locationName": "resourceArn"
+ },
+ "tagKeys": {
+ "shape": "TagKeyList",
+ "location": "querystring",
+ "locationName": "tagKeys"
+ }
+ }
+ },
+ "UntagResourceResponse": {
+ "type": "structure",
+ "members": {}
+ },
+ "UpdateConnectionRequest": {
+ "type": "structure",
+ "required": ["id", "authenticationType"],
+ "members": {
+ "sqlworkbenchSource": {
+ "shape": "SqlworkbenchSource",
+ "location": "header",
+ "locationName": "sqlworkbench-source"
+ },
+ "id": {
+ "shape": "UpdateConnectionRequestIdString",
+ "documentation": "Id of the connection to update
"
+ },
+ "name": {
+ "shape": "UpdateConnectionRequestNameString",
+ "documentation": "Name of the connection
"
+ },
+ "databaseName": {
+ "shape": "UpdateConnectionRequestDatabaseNameString",
+ "documentation": "Name of the database used for this connection
"
+ },
+ "authenticationType": {
+ "shape": "UpdateConnectionRequestAuthenticationTypeEnum",
+ "documentation": "Number representing the type of authentication to use (2 = IAM, 3 = Username and Password, 4 = Federated connection)
"
+ },
+ "secretArn": {
+ "shape": "UpdateConnectionRequestSecretArnString",
+ "documentation": "secretArn for redshift cluster
"
+ },
+ "clusterId": {
+ "shape": "UpdateConnectionRequestClusterIdString",
+ "documentation": "Id of the cluster used for this connection
"
+ },
+ "isServerless": {
+ "shape": "Boolean",
+ "documentation": "Is serverless connection
"
+ },
+ "dbUser": {
+ "shape": "DbUser",
+ "documentation": "User of the database used for this connection
"
+ },
+ "username": {
+ "shape": "DbUser",
+ "documentation": "Username used in the Username_Password connection type
"
+ },
+ "password": {
+ "shape": "UpdateConnectionRequestPasswordString",
+ "documentation": "Password of the user used for this connection
"
+ },
+ "host": {
+ "shape": "String",
+ "documentation": "Host address used for creating secret for Username_Password connection type
"
+ },
+ "databaseType": { "shape": "DatabaseType" },
+ "connectableResourceIdentifier": {
+ "shape": "UpdateConnectionRequestConnectableResourceIdentifierString",
+ "documentation": "Id of the connectable resource used for this connection
"
+ },
+ "connectableResourceType": {
+ "shape": "UpdateConnectionRequestConnectableResourceTypeString",
+ "documentation": "Type of the connectable resource used for this connection
"
+ }
+ }
+ },
+ "UpdateConnectionRequestAuthenticationTypeEnum": {
+ "type": "string",
+ "enum": ["2", "3", "4", "5", "6", "7", "8"],
+ "max": 1,
+ "min": 1,
+ "sensitive": true
+ },
+ "UpdateConnectionRequestClusterIdString": {
+ "type": "string",
+ "max": 63,
+ "min": 1
+ },
+ "UpdateConnectionRequestConnectableResourceIdentifierString": {
+ "type": "string",
+ "max": 63,
+ "min": 1,
+ "sensitive": true
+ },
+ "UpdateConnectionRequestConnectableResourceTypeString": {
+ "type": "string",
+ "max": 63,
+ "min": 1
+ },
+ "UpdateConnectionRequestDatabaseNameString": {
+ "type": "string",
+ "max": 64,
+ "min": 1,
+ "sensitive": true
+ },
+ "UpdateConnectionRequestIdString": {
+ "type": "string",
+ "max": 2048,
+ "min": 32
+ },
+ "UpdateConnectionRequestNameString": {
+ "type": "string",
+ "max": 512,
+ "min": 1,
+ "sensitive": true
+ },
+ "UpdateConnectionRequestPasswordString": {
+ "type": "string",
+ "max": 64,
+ "min": 8,
+ "sensitive": true
+ },
+ "UpdateConnectionRequestSecretArnString": {
+ "type": "string",
+ "max": 1000,
+ "min": 1
+ },
+ "UpdateConnectionResponse": {
+ "type": "structure",
+ "members": {
+ "data": { "shape": "Connection" }
+ }
+ },
+ "UserSettings": {
+ "type": "string",
+ "sensitive": true
+ },
+ "ValidationException": {
+ "type": "structure",
+ "required": ["message"],
+ "members": {
+ "message": { "shape": "String" },
+ "code": { "shape": "ErrorCode" }
+ },
+ "error": {
+ "httpStatusCode": 400,
+ "senderFault": true
+ },
+ "exception": true
+ },
+ "statusCode": {
+ "type": "integer",
+ "box": true,
+ "max": 500,
+ "min": 100
+ }
+ }
+}
diff --git a/packages/core/src/sagemakerunifiedstudio/shared/errors.ts b/packages/core/src/sagemakerunifiedstudio/shared/errors.ts
new file mode 100644
index 00000000000..1d582c22d74
--- /dev/null
+++ b/packages/core/src/sagemakerunifiedstudio/shared/errors.ts
@@ -0,0 +1,8 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export const errorCode = {
+ failedAuthConnecton: 'FailedAuthConnecton',
+}
diff --git a/packages/core/src/sagemakerunifiedstudio/shared/smusUtils.ts b/packages/core/src/sagemakerunifiedstudio/shared/smusUtils.ts
new file mode 100644
index 00000000000..3fe1dd9b27b
--- /dev/null
+++ b/packages/core/src/sagemakerunifiedstudio/shared/smusUtils.ts
@@ -0,0 +1,353 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { getLogger } from '../../shared/logger/logger'
+import { ToolkitError } from '../../shared/errors'
+import { isSageMaker } from '../../shared/extensionUtilities'
+import { getResourceMetadata } from './utils/resourceMetadataUtils'
+import fetch from 'node-fetch'
+
+/**
+ * Represents SSO instance information retrieved from DataZone
+ */
+export interface SsoInstanceInfo {
+ issuerUrl: string
+ ssoInstanceId: string
+ clientId: string
+ region: string
+}
+
+/**
+ * Response from DataZone /sso/login endpoint
+ */
+interface DataZoneSsoLoginResponse {
+ redirectUrl: string
+}
+
+/**
+ * Credential expiry time constants for SMUS providers (in milliseconds)
+ */
+export const SmusCredentialExpiry = {
+ /** Domain Execution Role (DER) credentials expiry time: 10 minutes */
+ derExpiryMs: 10 * 60 * 1000,
+ /** Project Role credentials expiry time: 10 minutes */
+ projectExpiryMs: 10 * 60 * 1000,
+ /** Connection credentials expiry time: 10 minutes */
+ connectionExpiryMs: 10 * 60 * 1000,
+} as const
+
+/**
+ * Error codes for SMUS-related operations
+ */
+export const SmusErrorCodes = {
+ /** Error code for when no active SMUS connection is available */
+ NoActiveConnection: 'NoActiveConnection',
+ /** Error code for when API calls timeout */
+ ApiTimeout: 'ApiTimeout',
+ /** Error code for when SMUS login fails */
+ SmusLoginFailed: 'SmusLoginFailed',
+ /** Error code for when redeeming access token fails */
+ RedeemAccessTokenFailed: 'RedeemAccessTokenFailed',
+} as const
+
+/**
+ * Timeout constants for SMUS API calls (in milliseconds)
+ */
+export const SmusTimeouts = {
+ /** Default timeout for API calls: 10 seconds */
+ apiCallTimeoutMs: 10 * 1000,
+} as const
+
+/**
+ * Interface for AWS credential objects that need validation
+ */
+interface CredentialObject {
+ accessKeyId?: unknown
+ secretAccessKey?: unknown
+ sessionToken?: unknown
+ expiration?: unknown
+}
+
+/**
+ * Validates AWS credential fields and throws appropriate errors if invalid
+ * @param credentials The credential object to validate
+ * @param errorCode The error code to use in ToolkitError
+ * @param contextMessage The context message for error messages (e.g., "API response", "project credential response")
+ * @throws ToolkitError if any credential field is invalid
+ */
+export function validateCredentialFields(
+ credentials: CredentialObject,
+ errorCode: string,
+ contextMessage: string,
+ validateExpireTime: boolean = false
+): void {
+ if (!credentials.accessKeyId || typeof credentials.accessKeyId !== 'string') {
+ throw new ToolkitError(`Invalid accessKeyId in ${contextMessage}: ${typeof credentials.accessKeyId}`, {
+ code: errorCode,
+ })
+ }
+ if (!credentials.secretAccessKey || typeof credentials.secretAccessKey !== 'string') {
+ throw new ToolkitError(`Invalid secretAccessKey in ${contextMessage}: ${typeof credentials.secretAccessKey}`, {
+ code: errorCode,
+ })
+ }
+ if (!credentials.sessionToken || typeof credentials.sessionToken !== 'string') {
+ throw new ToolkitError(`Invalid sessionToken in ${contextMessage}: ${typeof credentials.sessionToken}`, {
+ code: errorCode,
+ })
+ }
+ if (validateExpireTime) {
+ if (!credentials.expiration || !(credentials.expiration instanceof Date)) {
+ throw new ToolkitError(`Invalid expireTime in ${contextMessage}: ${typeof credentials.expiration}`, {
+ code: errorCode,
+ })
+ }
+ }
+}
+
+/**
+ * Utility class for SageMaker Unified Studio domain URL parsing and validation
+ */
+export class SmusUtils {
+ private static readonly logger = getLogger()
+
+ /**
+ * Extracts the domain ID from a SageMaker Unified Studio domain URL
+ * @param domainUrl The SageMaker Unified Studio domain URL
+ * @returns The extracted domain ID or undefined if not found
+ */
+ public static extractDomainIdFromUrl(domainUrl: string): string | undefined {
+ try {
+ // Domain URL format: https://dzd_d3hr1nfjbtwui1.sagemaker.us-east-2.on.aws
+ const url = new URL(domainUrl)
+ const hostname = url.hostname
+
+ // Extract domain ID from hostname (dzd_d3hr1nfjbtwui1 or dzd-d3hr1nfjbtwui1)
+ const domainIdMatch = hostname.match(/^(dzd[-_][a-zA-Z0-9_-]{1,36})\./)
+ return domainIdMatch?.[1]
+ } catch (error) {
+ this.logger.error('Failed to extract domain ID from URL: %s', error as Error)
+ return undefined
+ }
+ }
+
+ /**
+ * Extracts the AWS region from a SageMaker Unified Studio domain URL
+ * @param domainUrl The SageMaker Unified Studio domain URL
+ * @param fallbackRegion Fallback region if extraction fails (default: 'us-east-1')
+ * @returns The extracted AWS region or the fallback region if not found
+ */
+ public static extractRegionFromUrl(domainUrl: string, fallbackRegion: string = 'us-east-1'): string {
+ try {
+ // Domain URL formats:
+ // - https://dzd_d3hr1nfjbtwui1.sagemaker.us-east-2.on.aws
+ // - https://dzd_4gickdfsxtoxg0.sagemaker-gamma.us-west-2.on.aws
+ const url = new URL(domainUrl)
+ const hostname = url.hostname
+
+ // Extract region from hostname, handling both prod and non-prod stages
+ // Pattern matches: .sagemaker[-stage].{region}.on.aws
+ const regionMatch = hostname.match(/\.sagemaker(?:-[a-z]+)?\.([a-z0-9-]+)\.on\.aws$/)
+ return regionMatch?.[1] || fallbackRegion
+ } catch (error) {
+ this.logger.error('Failed to extract region from URL: %s', error as Error)
+ return fallbackRegion
+ }
+ }
+
+ /**
+ * Extracts both domain ID and region from a SageMaker Unified Studio domain URL
+ * @param domainUrl The SageMaker Unified Studio domain URL
+ * @param fallbackRegion Fallback region if extraction fails (default: 'us-east-1')
+ * @returns Object containing domainId and region
+ */
+ public static extractDomainInfoFromUrl(
+ domainUrl: string,
+ fallbackRegion: string = 'us-east-1'
+ ): { domainId: string | undefined; region: string } {
+ return {
+ domainId: this.extractDomainIdFromUrl(domainUrl),
+ region: this.extractRegionFromUrl(domainUrl, fallbackRegion),
+ }
+ }
+
+ /**
+ * Validates the domain URL format for SageMaker Unified Studio
+ * @param value The URL to validate
+ * @returns Error message if invalid, undefined if valid
+ */
+ public static validateDomainUrl(value: string): string | undefined {
+ if (!value || value.trim() === '') {
+ return 'Domain URL is required'
+ }
+
+ const trimmedValue = value.trim()
+
+ // Check HTTPS requirement
+ if (!trimmedValue.startsWith('https://')) {
+ return 'Domain URL must use HTTPS (https://)'
+ }
+
+ // Check basic URL format
+ try {
+ const url = new URL(trimmedValue)
+
+ // Check if it looks like a SageMaker Unified Studio domain
+ if (!url.hostname.includes('sagemaker') || !url.hostname.includes('on.aws')) {
+ return 'URL must be a valid SageMaker Unified Studio domain (e.g., https://dzd_xxxxxxxxx.sagemaker.us-east-1.on.aws)'
+ }
+
+ // Extract domain ID to validate
+ const domainId = this.extractDomainIdFromUrl(trimmedValue)
+
+ if (!domainId) {
+ return 'URL must contain a valid domain ID (starting with dzd- or dzd_)'
+ }
+
+ return undefined // Valid
+ } catch (err) {
+ return 'Invalid URL format'
+ }
+ }
+
+ /**
+ * Makes HTTP call to DataZone /sso/login endpoint
+ * @param domainUrl The SageMaker Unified Studio domain URL
+ * @param domainId The extracted domain ID
+ * @returns Promise resolving to the login response
+ * @throws ToolkitError if the API call fails
+ */
+ private static async callDataZoneLogin(domainUrl: string, domainId: string): Promise {
+ const loginUrl = new URL('/sso/login', domainUrl)
+ const requestBody = {
+ domainId: domainId,
+ }
+
+ try {
+ const response = await fetch(loginUrl.toString(), {
+ method: 'POST',
+ headers: {
+ 'Content-Type': 'application/json',
+ Accept: 'application/json',
+ 'User-Agent': 'aws-toolkit-vscode',
+ },
+ body: JSON.stringify(requestBody),
+ timeout: SmusTimeouts.apiCallTimeoutMs,
+ })
+
+ if (!response.ok) {
+ throw new ToolkitError(`SMUS login failed: ${response.status} ${response.statusText}`, {
+ code: SmusErrorCodes.SmusLoginFailed,
+ })
+ }
+
+ return (await response.json()) as DataZoneSsoLoginResponse
+ } catch (error) {
+ // Handle timeout errors specifically
+ if (error instanceof Error && (error.name === 'AbortError' || error.message.includes('timeout'))) {
+ throw new ToolkitError(
+ `DataZone login request timed out after ${SmusTimeouts.apiCallTimeoutMs / 1000} seconds`,
+ {
+ code: SmusErrorCodes.ApiTimeout,
+ cause: error,
+ }
+ )
+ }
+ // Re-throw other errors as-is
+ throw error
+ }
+ }
+
+ /**
+ * Gets SSO instance information by calling DataZone /sso/login endpoint
+ * This extracts the proper SSO instance ID and issuer URL needed for OAuth client registration
+ *
+ * @param domainUrl The SageMaker Unified Studio domain URL
+ * @returns Promise resolving to SSO instance information
+ * @throws ToolkitError if the API call fails or response is invalid
+ */
+ public static async getSsoInstanceInfo(domainUrl: string): Promise {
+ try {
+ this.logger.info(`SMUS Auth: Getting SSO instance info from DataZone for domainurl: ${domainUrl}`)
+
+ // Extract domain ID from the domain URL
+ const domainId = this.extractDomainIdFromUrl(domainUrl)
+ if (!domainId) {
+ throw new ToolkitError('Invalid domain URL format', { code: 'InvalidDomainUrl' })
+ }
+
+ // Call DataZone /sso/login endpoint to get redirect URL with SSO instance info
+ const loginData = await this.callDataZoneLogin(domainUrl, domainId)
+ if (!loginData.redirectUrl) {
+ throw new ToolkitError('No redirect URL received from DataZone login', { code: 'InvalidLoginResponse' })
+ }
+
+ // Parse the redirect URL to extract SSO instance information
+ const redirectUrl = new URL(loginData.redirectUrl)
+ const clientIdParam = redirectUrl.searchParams.get('client_id')
+ if (!clientIdParam) {
+ throw new ToolkitError('No client_id found in DataZone redirect URL', { code: 'InvalidRedirectUrl' })
+ }
+
+ // Decode the client_id ARN: arn:aws:sso::785498918019:application/ssoins-6684636af7e1a207/apl-5f60548b7f5677a2
+ const decodedClientId = decodeURIComponent(clientIdParam)
+ const arnParts = decodedClientId.split('/')
+ if (arnParts.length < 2) {
+ throw new ToolkitError('Invalid client_id ARN format', { code: 'InvalidArnFormat' })
+ }
+
+ const ssoInstanceId = arnParts[1] // Extract ssoins-6684636af7e1a207
+ const issuerUrl = `https://identitycenter.amazonaws.com/${ssoInstanceId}`
+
+ // Extract region from domain URL
+ const region = this.extractRegionFromUrl(domainUrl)
+
+ this.logger.info('SMUS Auth: Extracted SSO instance info: %s', ssoInstanceId)
+
+ return {
+ issuerUrl,
+ ssoInstanceId,
+ clientId: decodedClientId,
+ region,
+ }
+ } catch (error) {
+ const errorMsg = error instanceof Error ? error.message : 'Unknown error'
+ this.logger.error('SMUS Auth: Failed to get SSO instance info: %s', errorMsg)
+
+ if (error instanceof ToolkitError) {
+ throw error
+ }
+
+ throw new ToolkitError(`Failed to get SSO instance info: ${errorMsg}`, {
+ code: 'SsoInstanceInfoFailed',
+ cause: error instanceof Error ? error : undefined,
+ })
+ }
+ }
+ /**
+ * Extracts SSO ID from a user ID in the format "user-"
+ * @param userId The user ID to extract SSO ID from
+ * @returns The extracted SSO ID
+ * @throws Error if the userId format is invalid
+ */
+ public static extractSSOIdFromUserId(userId: string): string {
+ const match = userId.match(/user-(.+)$/)
+ if (!match) {
+ this.logger.error(`Invalid UserId format: ${userId}`)
+ throw new Error(`Invalid UserId format: ${userId}`)
+ }
+ return match[1]
+ }
+
+ /**
+ * Checks if we're in SMUS space environment (should hide certain UI elements)
+ * @returns True if in SMUS space environment with DataZone domain ID
+ */
+ public static isInSmusSpaceEnvironment(): boolean {
+ const isSMUSspace = isSageMaker('SMUS') || isSageMaker('SMUS-SPACE-REMOTE-ACCESS')
+ const resourceMetadata = getResourceMetadata()
+ return isSMUSspace && !!resourceMetadata?.AdditionalMetadata?.DataZoneDomainId
+ }
+}
diff --git a/packages/core/src/sagemakerunifiedstudio/shared/utils/resourceMetadataUtils.ts b/packages/core/src/sagemakerunifiedstudio/shared/utils/resourceMetadataUtils.ts
new file mode 100644
index 00000000000..61ce0430ecd
--- /dev/null
+++ b/packages/core/src/sagemakerunifiedstudio/shared/utils/resourceMetadataUtils.ts
@@ -0,0 +1,93 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { fs } from '../../../shared/fs/fs'
+import { getLogger } from '../../../shared/logger/logger'
+import { isSageMaker } from '../../../shared/extensionUtilities'
+
+/**
+ * Resource metadata schema used by `resource-metadata.json` in SageMaker Unified Studio spaces
+ */
+export type ResourceMetadata = {
+ AppType?: string
+ DomainId?: string
+ SpaceName?: string
+ UserProfileName?: string
+ ExecutionRoleArn?: string
+ ResourceArn?: string
+ ResourceName?: string
+ AppImageVersion?: string
+ AdditionalMetadata?: {
+ DataZoneDomainId?: string
+ DataZoneDomainRegion?: string
+ DataZoneEndpoint?: string
+ DataZoneEnvironmentId?: string
+ DataZoneProjectId?: string
+ DataZoneScopeName?: string
+ DataZoneStage?: string
+ DataZoneUserId?: string
+ PrivateSubnets?: string
+ ProjectS3Path?: string
+ SecurityGroup?: string
+ }
+ ResourceArnCaseSensitive?: string
+ IpAddressType?: string
+} & Record
+
+const resourceMetadataPath = '/opt/ml/metadata/resource-metadata.json'
+let resourceMetadata: ResourceMetadata | undefined = undefined
+
+/**
+ * Gets the cached resource metadata (must be initialized with `initializeResourceMetadata()` first)
+ * @returns ResourceMetadata object or undefined if not yet initialized
+ */
+export function getResourceMetadata(): ResourceMetadata | undefined {
+ return resourceMetadata
+}
+
+/**
+ * Initializes resource metadata by reading and parsing the resource-metadata.json file
+ */
+export async function initializeResourceMetadata(): Promise {
+ const logger = getLogger()
+
+ if (!isSageMaker('SMUS') && !isSageMaker('SMUS-SPACE-REMOTE-ACCESS')) {
+ logger.debug(`Not in SageMaker Unified Studio space, skipping initialization of resource metadata`)
+ return
+ }
+
+ try {
+ if (!(await resourceMetadataFileExists())) {
+ logger.debug(`Resource metadata file not found at: ${resourceMetadataPath}`)
+ }
+
+ const fileContent = await fs.readFileText(resourceMetadataPath)
+ resourceMetadata = JSON.parse(fileContent) as ResourceMetadata
+ logger.debug(`Successfully read resource metadata from: ${resourceMetadataPath}`)
+ } catch (error) {
+ logger.error(`Failed to read or parse resource metadata file: ${error as Error}`)
+ }
+}
+
+/**
+ * Checks if the resource-metadata.json file exists
+ * @returns True if the file exists, false otherwise
+ */
+export async function resourceMetadataFileExists(): Promise {
+ try {
+ return await fs.existsFile(resourceMetadataPath)
+ } catch (error) {
+ const logger = getLogger()
+ logger.error(`Failed to check if resource metadata file exists: ${error as Error}`)
+ return false
+ }
+}
+
+/**
+ * Resets the cached resource metadata
+ */
+export function resetResourceMetadata(): void {
+ resourceMetadata = undefined
+}
diff --git a/packages/core/src/shared/awsClientBuilder.ts b/packages/core/src/shared/awsClientBuilder.ts
index bdec40957cb..b6464317a3e 100644
--- a/packages/core/src/shared/awsClientBuilder.ts
+++ b/packages/core/src/shared/awsClientBuilder.ts
@@ -82,8 +82,11 @@ export class DefaultAWSClientBuilder implements AWSClientBuilder {
const listeners = Array.isArray(onRequest) ? onRequest : [onRequest]
const opt = { ...options }
delete opt.onRequestSetup
+ if (opt.credentialProvider) {
+ opt.credentials = await opt.credentialProvider.resolvePromise()
+ }
- if (!opt.credentials && !opt.token) {
+ if (!opt.credentials && !opt.token && !opt.credentialProvider) {
const shim = this.awsContext.credentialsShim
if (!shim) {
diff --git a/packages/core/src/shared/clients/clientWrapper.ts b/packages/core/src/shared/clients/clientWrapper.ts
index a90d009eb18..beb117a9bf6 100644
--- a/packages/core/src/shared/clients/clientWrapper.ts
+++ b/packages/core/src/shared/clients/clientWrapper.ts
@@ -19,22 +19,13 @@ export abstract class ClientWrapper implements vscode.Dispo
public constructor(
public readonly regionCode: string,
- private readonly clientType: AwsClientConstructor,
- private readonly isSageMaker: boolean = false
+ private readonly clientType: AwsClientConstructor
) {}
protected getClient(ignoreCache: boolean = false) {
const args = {
serviceClient: this.clientType,
region: this.regionCode,
- ...(this.isSageMaker
- ? {
- clientOptions: {
- endpoint: `https://sagemaker.${this.regionCode}.amazonaws.com`,
- region: this.regionCode,
- },
- }
- : {}),
}
return ignoreCache
? globals.sdkClientBuilderV3.createAwsService(args)
diff --git a/packages/core/src/shared/clients/sagemaker.ts b/packages/core/src/shared/clients/sagemaker.ts
index 8a8e138dd85..ff086ed1d9e 100644
--- a/packages/core/src/shared/clients/sagemaker.ts
+++ b/packages/core/src/shared/clients/sagemaker.ts
@@ -6,6 +6,7 @@
import * as vscode from 'vscode'
import {
AppDetails,
+ AppType,
CreateAppCommand,
CreateAppCommandInput,
CreateAppCommandOutput,
@@ -48,14 +49,42 @@ import { getDomainSpaceKey } from '../../awsService/sagemaker/utils'
import { getLogger } from '../logger/logger'
import { ToolkitError } from '../errors'
import { yes, no, continueText, cancel } from '../localizedText'
+import { AwsCredentialIdentity } from '@aws-sdk/types'
+import globals from '../extensionGlobals'
export interface SagemakerSpaceApp extends SpaceDetails {
App?: AppDetails
DomainSpaceKey: string
}
+
export class SagemakerClient extends ClientWrapper {
- public constructor(public override readonly regionCode: string) {
- super(regionCode, SageMakerClient, true)
+ public constructor(
+ public override readonly regionCode: string,
+ private readonly credentialsProvider?: () => Promise
+ ) {
+ super(regionCode, SageMakerClient)
+ }
+
+ protected override getClient(ignoreCache: boolean = false) {
+ if (!this.client || ignoreCache) {
+ const args = {
+ serviceClient: SageMakerClient,
+ region: this.regionCode,
+ clientOptions: {
+ endpoint: `https://sagemaker.${this.regionCode}.amazonaws.com`,
+ region: this.regionCode,
+ ...(this.credentialsProvider && { credentials: this.credentialsProvider }),
+ },
+ }
+ this.client = globals.sdkClientBuilderV3.createAwsService(args)
+ }
+ return this.client
+ }
+
+ public override dispose() {
+ getLogger().debug('SagemakerClient: Disposing client %O', this.client)
+ this.client?.destroy()
+ this.client = undefined
}
public listSpaces(request: ListSpacesCommandInput = {}): AsyncCollection {
@@ -200,27 +229,37 @@ export class SagemakerClient extends ClientWrapper {
}
}
- public async fetchSpaceAppsAndDomains(): Promise<
- [Map, Map]
- > {
- try {
- const appMap: Map = await this.listApps()
- .flatten()
- .filter((app) => !!app.DomainId && !!app.SpaceName)
- .filter((app) => app.AppType === 'JupyterLab' || app.AppType === 'CodeEditor')
- .toMap((app) => getDomainSpaceKey(app.DomainId || '', app.SpaceName || ''))
-
- const spaceApps: Map = await this.listSpaces()
- .flatten()
- .filter((space) => !!space.DomainId && !!space.SpaceName)
- .map((space) => {
- const key = getDomainSpaceKey(space.DomainId || '', space.SpaceName || '')
- return { ...space, App: appMap.get(key), DomainSpaceKey: key }
- })
- .toMap((space) => getDomainSpaceKey(space.DomainId || '', space.SpaceName || ''))
+ public async listSpaceApps(domainId?: string): Promise> {
+ // Create options object conditionally if domainId is provided
+ const options = domainId ? { DomainIdEquals: domainId } : undefined
+
+ const appMap: Map = await this.listApps(options)
+ .flatten()
+ .filter((app) => !!app.DomainId && !!app.SpaceName)
+ .filter((app) => app.AppType === AppType.JupyterLab || app.AppType === AppType.CodeEditor)
+ .toMap((app) => getDomainSpaceKey(app.DomainId || '', app.SpaceName || ''))
+
+ const spaceApps: Map = await this.listSpaces(options)
+ .flatten()
+ .filter((space) => !!space.DomainId && !!space.SpaceName)
+ .map((space) => {
+ const key = getDomainSpaceKey(space.DomainId || '', space.SpaceName || '')
+ return { ...space, App: appMap.get(key), DomainSpaceKey: key }
+ })
+ .toMap((space) => getDomainSpaceKey(space.DomainId || '', space.SpaceName || ''))
+ return spaceApps
+ }
+ public async fetchSpaceAppsAndDomains(
+ domainId?: string,
+ filterSmusDomains: boolean = true
+ ): Promise<[Map, Map]> {
+ try {
+ const spaceApps = await this.listSpaceApps(domainId)
// Get de-duped list of domain IDs for all of the spaces
- const domainIds: string[] = [...new Set([...spaceApps].map(([_, spaceApp]) => spaceApp.DomainId || ''))]
+ const domainIds: string[] = domainId
+ ? [domainId]
+ : [...new Set([...spaceApps].map(([_, spaceApp]) => spaceApp.DomainId || ''))]
// Get details for each domain
const domains: [string, DescribeDomainResponse][] = await Promise.all(
@@ -235,9 +274,11 @@ export class SagemakerClient extends ClientWrapper {
const filteredSpaceApps = new Map(
[...spaceApps]
- // Filter out SageMaker Unified Studio domains
- .filter(([_, spaceApp]) =>
- isEmpty(domainsMap.get(spaceApp.DomainId || '')?.DomainSettings?.UnifiedStudioSettings)
+ // Filter out SageMaker Unified Studio domains only if filterSmusDomains is true
+ .filter(
+ ([_, spaceApp]) =>
+ !filterSmusDomains ||
+ isEmpty(domainsMap.get(spaceApp.DomainId || '')?.DomainSettings?.UnifiedStudioSettings)
)
)
diff --git a/packages/core/src/shared/extensionUtilities.ts b/packages/core/src/shared/extensionUtilities.ts
index 80bedf1e0f6..b8b5780c612 100644
--- a/packages/core/src/shared/extensionUtilities.ts
+++ b/packages/core/src/shared/extensionUtilities.ts
@@ -188,7 +188,7 @@ export function isCloud9(flavor: 'classic' | 'codecatalyst' | 'any' = 'any'): bo
* @param appName to identify the proper SM instance
* @returns true if the current system is SageMaker(SMAI or SMUS)
*/
-export function isSageMaker(appName: 'SMAI' | 'SMUS' = 'SMAI'): boolean {
+export function isSageMaker(appName: 'SMAI' | 'SMUS' | 'SMUS-SPACE-REMOTE-ACCESS' = 'SMAI'): boolean {
// Check for SageMaker-specific environment variables first
let hasSMEnvVars: boolean = false
if (hasSageMakerEnvVars()) {
@@ -201,6 +201,9 @@ export function isSageMaker(appName: 'SMAI' | 'SMUS' = 'SMAI'): boolean {
return vscode.env.appName === sageMakerAppname && hasSMEnvVars
case 'SMUS':
return vscode.env.appName === sageMakerAppname && isSageMakerUnifiedStudio() && hasSMEnvVars
+ case 'SMUS-SPACE-REMOTE-ACCESS':
+ // When is true, the AWS toolkit is running in remote SSH conenction to SageMaker Unified Studio space
+ return vscode.env.appName !== sageMakerAppname && isSageMakerUnifiedStudio() && hasSMEnvVars
default:
return false
}
diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts
index e8e6a3bff44..0e9ab819046 100644
--- a/packages/core/src/shared/globalState.ts
+++ b/packages/core/src/shared/globalState.ts
@@ -8,7 +8,7 @@ import { getLogger } from './logger/logger'
import * as redshift from '../awsService/redshift/models/models'
import { TypeConstructor, cast } from './utilities/typeConstructors'
-type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
+type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId' | 'smus'
export type ToolIdStateKey = `${ToolId}.savedConnectionId`
export type JsonSchemasKey = 'devfileSchemaVersion' | 'samAndCfnSchemaVersion'
diff --git a/packages/core/src/shared/telemetry/vscodeTelemetry.json b/packages/core/src/shared/telemetry/vscodeTelemetry.json
index 1128eef8ab6..1ca57d4669f 100644
--- a/packages/core/src/shared/telemetry/vscodeTelemetry.json
+++ b/packages/core/src/shared/telemetry/vscodeTelemetry.json
@@ -238,6 +238,46 @@
"name": "executedCount",
"type": "int",
"description": "The number of executed operations"
+ },
+ {
+ "name": "smusDomainId",
+ "type": "string",
+ "description": "SMUS domain identifier"
+ },
+ {
+ "name": "smusProjectId",
+ "type": "string",
+ "description": "SMUS project identifier"
+ },
+ {
+ "name": "smusSpaceKey",
+ "type": "string",
+ "description": "SMUS space composite key consisting of domainId and spaceName"
+ },
+ {
+ "name": "smusToolkitEnv",
+ "type": "string",
+ "description": "The environment user is running SMUS extension against"
+ },
+ {
+ "name": "smusDomainRegion",
+ "type": "string",
+ "description": "The SMUS domain region"
+ },
+ {
+ "name": "smusProjectRegion",
+ "type": "string",
+ "description": "The SMUS project region"
+ },
+ {
+ "name": "smusConnectionId",
+ "type": "string",
+ "description": "SMUS connection identifier"
+ },
+ {
+ "name": "smusConnectionType",
+ "type": "string",
+ "description": "SMUS connection type"
}
],
"metrics": [
@@ -1257,6 +1297,201 @@
"required": false
}
]
+ },
+ {
+ "name": "smus_login",
+ "description": "Emitted whenever a user signin to SMUS",
+ "metadata": [
+ {
+ "type": "smusDomainId",
+ "required": false
+ }
+ ]
+ },
+ {
+ "name": "smus_signOut",
+ "description": "Emitted whenever a user signouts SMUS",
+ "metadata": [
+ {
+ "type": "smusDomainId",
+ "required": false
+ }
+ ]
+ },
+ {
+ "name": "smus_accessProject",
+ "description": "Emitted whenever a user accesses a SMUS project",
+ "metadata": [
+ {
+ "type": "smusDomainId",
+ "required": false
+ },
+ {
+ "type": "smusProjectId",
+ "required": false
+ },
+ {
+ "type": "smusDomainRegion",
+ "required": false
+ }
+ ]
+ },
+ {
+ "name": "smus_renderProjectChildrenNode",
+ "description": "Emitted whenever children node of project is rendered",
+ "metadata": [
+ {
+ "type": "smusToolkitEnv",
+ "required": false
+ },
+ {
+ "type": "smusDomainId",
+ "required": false
+ },
+ {
+ "type": "smusProjectId",
+ "required": false
+ },
+ {
+ "type": "smusDomainRegion",
+ "required": false
+ }
+ ],
+ "passive": true
+ },
+ {
+ "name": "smus_startSpace",
+ "description": "Emitted whenever a user starts a SMUS space",
+ "metadata": [
+ {
+ "type": "smusDomainId",
+ "required": false
+ },
+ {
+ "type": "smusProjectId",
+ "required": false
+ },
+ {
+ "type": "smusSpaceKey",
+ "required": false
+ },
+ {
+ "type": "smusDomainRegion",
+ "required": false
+ }
+ ]
+ },
+ {
+ "name": "smus_stopSpace",
+ "description": "Emitted whenever a user stop a SMUS space",
+ "metadata": [
+ {
+ "type": "smusDomainId",
+ "required": false
+ },
+ {
+ "type": "smusProjectId",
+ "required": false
+ },
+ {
+ "type": "smusSpaceKey",
+ "required": false
+ },
+ {
+ "type": "smusDomainRegion",
+ "required": false
+ }
+ ]
+ },
+ {
+ "name": "smus_renderS3Node",
+ "description": "Emitted whenever rendering a s3 node",
+ "metadata": [
+ {
+ "type": "smusToolkitEnv",
+ "required": false
+ },
+ {
+ "type": "smusDomainId",
+ "required": false
+ },
+ {
+ "type": "smusProjectId",
+ "required": false
+ },
+ {
+ "type": "smusProjectRegion",
+ "required": false
+ },
+ {
+ "type": "smusConnectionId",
+ "required": false
+ },
+ {
+ "type": "smusConnectionType",
+ "required": false
+ }
+ ]
+ },
+ {
+ "name": "smus_renderRedshiftNode",
+ "description": "Emitted whenever rendering a Redshift node",
+ "metadata": [
+ {
+ "type": "smusToolkitEnv",
+ "required": false
+ },
+ {
+ "type": "smusDomainId",
+ "required": false
+ },
+ {
+ "type": "smusProjectId",
+ "required": false
+ },
+ {
+ "type": "smusProjectRegion",
+ "required": false
+ },
+ {
+ "type": "smusConnectionId",
+ "required": false
+ },
+ {
+ "type": "smusConnectionType",
+ "required": false
+ }
+ ]
+ },
+ {
+ "name": "smus_renderLakehouseNode",
+ "description": "Emitted whenever rendering a Lakehouse node",
+ "metadata": [
+ {
+ "type": "smusToolkitEnv",
+ "required": false
+ },
+ {
+ "type": "smusDomainId",
+ "required": false
+ },
+ {
+ "type": "smusProjectId",
+ "required": false
+ },
+ {
+ "type": "smusProjectRegion",
+ "required": false
+ },
+ {
+ "type": "smusConnectionId",
+ "required": false
+ },
+ {
+ "type": "smusConnectionType",
+ "required": false
+ }
+ ]
}
]
}
diff --git a/packages/core/src/shared/vscode/setContext.ts b/packages/core/src/shared/vscode/setContext.ts
index 7cfaf4092f8..3d45d93e14a 100644
--- a/packages/core/src/shared/vscode/setContext.ts
+++ b/packages/core/src/shared/vscode/setContext.ts
@@ -30,6 +30,8 @@ export type contextKey =
| 'aws.stepFunctions.isWorkflowStudioFocused'
| 'aws.toolkit.notifications.show'
| 'aws.amazonq.editSuggestionActive'
+ | 'aws.smus.connected'
+ | 'aws.smus.inSmusSpaceEnvironment'
// Deprecated/legacy names. New keys should start with "aws.".
| 'codewhisperer.activeLine'
| 'gumby.isPlanAvailable'
diff --git a/packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts b/packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts
index 06f19a5e890..3134f11e5e0 100644
--- a/packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts
+++ b/packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts
@@ -5,10 +5,22 @@
import * as sinon from 'sinon'
import * as assert from 'assert'
-import { persistLocalCredentials, persistSSMConnection } from '../../../awsService/sagemaker/credentialMapping'
+import {
+ persistLocalCredentials,
+ persistSSMConnection,
+ persistSmusProjectCreds,
+ loadMappings,
+ saveMappings,
+ setSpaceIamProfile,
+ setSpaceSsoProfile,
+ setSmusSpaceSsoProfile,
+ setSpaceCredentials,
+} from '../../../awsService/sagemaker/credentialMapping'
import { Auth } from '../../../auth'
import { DevSettings, fs } from '../../../shared'
import globals from '../../../shared/extensionGlobals'
+import { SagemakerUnifiedStudioSpaceNode } from '../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
+import { SageMakerUnifiedStudioSpacesParentNode } from '../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode'
describe('credentialMapping', () => {
describe('persistLocalCredentials', () => {
@@ -207,4 +219,230 @@ describe('credentialMapping', () => {
})
})
})
+
+ describe('persistSmusProjectCreds', () => {
+ const appArn = 'arn:aws:sagemaker:us-west-2:123456789012:space/d-f0lwireyzpjp/test-space'
+ const projectId = 'test-project-id'
+ let sandbox: sinon.SinonSandbox
+ let mockNode: sinon.SinonStubbedInstance
+ let mockParent: sinon.SinonStubbedInstance
+
+ beforeEach(() => {
+ sandbox = sinon.createSandbox()
+ mockNode = sandbox.createStubInstance(SagemakerUnifiedStudioSpaceNode)
+ mockParent = sandbox.createStubInstance(SageMakerUnifiedStudioSpacesParentNode)
+ })
+
+ afterEach(() => {
+ sandbox.restore()
+ })
+
+ it('persists SMUS project credentials', async () => {
+ const mockCredentialProvider = {
+ getCredentials: sandbox.stub().resolves(),
+ startProactiveCredentialRefresh: sandbox.stub(),
+ }
+
+ const mockAuthProvider = {
+ getProjectCredentialProvider: sandbox.stub().resolves(mockCredentialProvider),
+ }
+
+ mockNode.getParent.returns(mockParent as any)
+ mockParent.getAuthProvider.returns(mockAuthProvider as any)
+ mockParent.getProjectId.returns(projectId)
+
+ sandbox.stub(fs, 'existsFile').resolves(false)
+ const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+
+ await persistSmusProjectCreds(appArn, mockNode as any)
+
+ assert.ok(writeStub.calledOnce)
+ const raw = writeStub.firstCall.args[1]
+ const data = JSON.parse(typeof raw === 'string' ? raw : raw.toString())
+ assert.deepStrictEqual(data.localCredential?.[appArn], {
+ type: 'sso',
+ smusProjectId: projectId,
+ })
+
+ // Verify the correct methods were called
+ assert.ok(mockAuthProvider.getProjectCredentialProvider.calledWith(projectId))
+ assert.ok(mockCredentialProvider.getCredentials.calledOnce)
+ assert.ok(mockCredentialProvider.startProactiveCredentialRefresh.calledOnce)
+ })
+ })
+
+ describe('loadMappings', () => {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(() => {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(() => {
+ sandbox.restore()
+ })
+
+ it('returns empty object when file does not exist', async () => {
+ sandbox.stub(fs, 'existsFile').resolves(false)
+
+ const result = await loadMappings()
+
+ assert.deepStrictEqual(result, {})
+ })
+
+ it('loads and parses existing mappings', async () => {
+ const mockData = { localCredential: { 'test-arn': { type: 'iam' as const, profileName: 'test' } } }
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(mockData))
+
+ const result = await loadMappings()
+
+ assert.deepStrictEqual(result, mockData)
+ })
+
+ it('returns empty object on parse error', async () => {
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves('invalid json')
+
+ const result = await loadMappings()
+
+ assert.deepStrictEqual(result, {})
+ })
+ })
+
+ describe('saveMappings', () => {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(() => {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(() => {
+ sandbox.restore()
+ })
+
+ it('saves mappings to file', async () => {
+ const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+ const testData = { localCredential: { 'test-arn': { type: 'iam' as const, profileName: 'test' } } }
+
+ await saveMappings(testData)
+
+ assert.ok(writeStub.calledOnce)
+ const [, content, options] = writeStub.firstCall.args
+ assert.strictEqual(content, JSON.stringify(testData, undefined, 2))
+ assert.deepStrictEqual(options, { mode: 0o600, atomic: true })
+ })
+ })
+
+ describe('setSpaceIamProfile', () => {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(() => {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(() => {
+ sandbox.restore()
+ })
+
+ it('sets IAM profile for space', async () => {
+ sandbox.stub(fs, 'existsFile').resolves(false)
+ const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+
+ await setSpaceIamProfile('test-space', 'test-profile')
+
+ const raw = writeStub.firstCall.args[1]
+ const data = JSON.parse(typeof raw === 'string' ? raw : raw.toString())
+ assert.deepStrictEqual(data.localCredential?.['test-space'], {
+ type: 'iam',
+ profileName: 'test-profile',
+ })
+ })
+ })
+
+ describe('setSpaceSsoProfile', () => {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(() => {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(() => {
+ sandbox.restore()
+ })
+
+ it('sets SSO profile for space', async () => {
+ sandbox.stub(fs, 'existsFile').resolves(false)
+ const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+
+ await setSpaceSsoProfile('test-space', 'access-key', 'secret', 'token')
+
+ const raw = writeStub.firstCall.args[1]
+ const data = JSON.parse(typeof raw === 'string' ? raw : raw.toString())
+ assert.deepStrictEqual(data.localCredential?.['test-space'], {
+ type: 'sso',
+ accessKey: 'access-key',
+ secret: 'secret',
+ token: 'token',
+ })
+ })
+ })
+
+ describe('setSmusSpaceSsoProfile', () => {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(() => {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(() => {
+ sandbox.restore()
+ })
+
+ it('sets SMUS SSO profile for space', async () => {
+ sandbox.stub(fs, 'existsFile').resolves(false)
+ const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+
+ await setSmusSpaceSsoProfile('test-space', 'project-id')
+
+ const raw = writeStub.firstCall.args[1]
+ const data = JSON.parse(typeof raw === 'string' ? raw : raw.toString())
+ assert.deepStrictEqual(data.localCredential?.['test-space'], {
+ type: 'sso',
+ smusProjectId: 'project-id',
+ })
+ })
+ })
+
+ describe('setSpaceCredentials', () => {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(() => {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(() => {
+ sandbox.restore()
+ })
+
+ it('sets space credentials with refresh URL', async () => {
+ sandbox.stub(fs, 'existsFile').resolves(false)
+ const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+ const credentials = { sessionId: 'sess', url: 'ws://test', token: 'token' }
+
+ await setSpaceCredentials('test-space', 'https://refresh.url', credentials)
+
+ const raw = writeStub.firstCall.args[1]
+ const data = JSON.parse(typeof raw === 'string' ? raw : raw.toString())
+ assert.deepStrictEqual(data.deepLink?.['test-space'], {
+ refreshUrl: 'https://refresh.url',
+ requests: {
+ 'initial-connection': {
+ ...credentials,
+ status: 'fresh',
+ },
+ },
+ })
+ })
+ })
})
diff --git a/packages/core/src/test/awsService/sagemaker/detached-server/credentials.test.ts b/packages/core/src/test/awsService/sagemaker/detached-server/credentials.test.ts
index a979c2186d3..3db189f8390 100644
--- a/packages/core/src/test/awsService/sagemaker/detached-server/credentials.test.ts
+++ b/packages/core/src/test/awsService/sagemaker/detached-server/credentials.test.ts
@@ -73,6 +73,69 @@ describe('resolveCredentialsFor', () => {
})
})
+ it('resolves SSO credentials with SMUS project ID', async () => {
+ sinon.stub(utils, 'readMapping').resolves({
+ localCredential: {
+ [connectionId]: {
+ type: 'sso',
+ smusProjectId: 'project123',
+ },
+ },
+ smusProjects: {
+ project123: {
+ accessKey: 'smus-key',
+ secret: 'smus-secret',
+ token: 'smus-token',
+ },
+ },
+ })
+
+ const creds = await resolveCredentialsFor(connectionId)
+ assert.deepStrictEqual(creds, {
+ accessKeyId: 'smus-key',
+ secretAccessKey: 'smus-secret',
+ sessionToken: 'smus-token',
+ })
+ })
+
+ it('throws if SMUS project credentials are missing', async () => {
+ sinon.stub(utils, 'readMapping').resolves({
+ localCredential: {
+ [connectionId]: {
+ type: 'sso',
+ smusProjectId: 'project123',
+ },
+ },
+ smusProjects: {
+ project123: {
+ accessKey: '',
+ secret: 'smus-secret',
+ token: 'smus-token',
+ },
+ },
+ })
+
+ await assert.rejects(() => resolveCredentialsFor(connectionId), {
+ message: `Missing ProjectRole credentials for SMUS Space "${connectionId}"`,
+ })
+ })
+
+ it('throws if SMUS project is not found', async () => {
+ sinon.stub(utils, 'readMapping').resolves({
+ localCredential: {
+ [connectionId]: {
+ type: 'sso',
+ smusProjectId: 'nonexistent',
+ },
+ },
+ smusProjects: {},
+ })
+
+ await assert.rejects(() => resolveCredentialsFor(connectionId), {
+ message: `Missing ProjectRole credentials for SMUS Space "${connectionId}"`,
+ })
+ })
+
it('throws for unsupported profile types', async () => {
sinon.stub(utils, 'readMapping').resolves({
localCredential: {
diff --git a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerParentNode.test.ts b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerParentNode.test.ts
index 8fccfe4bfd9..b7cf98496fc 100644
--- a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerParentNode.test.ts
+++ b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerParentNode.test.ts
@@ -23,7 +23,7 @@ describe('sagemakerParentNode', function () {
let testNode: SagemakerParentNode
let client: SagemakerClient
let fetchSpaceAppsAndDomainsStub: sinon.SinonStub<
- [],
+ [domainId?: string | undefined, filterSmusDomains?: boolean | undefined],
Promise<[Map, Map]>
>
let getCallerIdentityStub: sinon.SinonStub<[], Promise>
diff --git a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerSpaceNode.test.ts b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerSpaceNode.test.ts
index 57b4d7a80c6..b0fc6d78c0f 100644
--- a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerSpaceNode.test.ts
+++ b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerSpaceNode.test.ts
@@ -69,10 +69,9 @@ describe('SagemakerSpaceNode', function () {
})
it('returns ARN from describeApp', async function () {
- describeAppStub.resolves({ AppArn: 'arn:aws:sagemaker:1234:app/TestApp' })
+ describeAppStub.resolves({ AppArn: 'arn:aws:sagemaker:1234:app/TestApp', $metadata: {} })
- const node = new SagemakerSpaceNode(testParent, client, testRegion, testSpaceApp)
- const arn = await node.getAppArn()
+ const arn = await testSpaceAppNode.getAppArn()
assert.strictEqual(arn, 'arn:aws:sagemaker:1234:app/TestApp')
sinon.assert.calledOnce(describeAppStub)
@@ -84,10 +83,40 @@ describe('SagemakerSpaceNode', function () {
})
})
- it('updates status with new spaceApp', async function () {
- const newStatus = 'Starting'
+ it('returns space ARN from describeSpace', async function () {
+ const describeSpaceStub = sinon.stub(SagemakerClient.prototype, 'describeSpace')
+ describeSpaceStub.resolves({ SpaceArn: 'arn:aws:sagemaker:1234:space/TestSpace', $metadata: {} })
+
+ const arn = await testSpaceAppNode.getSpaceArn()
+
+ assert.strictEqual(arn, 'arn:aws:sagemaker:1234:space/TestSpace')
+ sinon.assert.calledOnce(describeSpaceStub)
+ })
+
+ it('updates status with new spaceApp', function () {
const newSpaceApp = { ...testSpaceApp, App: { AppName: 'TestApp', Status: 'Pending' } } as SagemakerSpaceApp
testSpaceAppNode.updateSpace(newSpaceApp)
- assert.strictEqual(testSpaceAppNode.getStatus(), newStatus)
+ assert.strictEqual(testSpaceAppNode.getStatus(), 'Starting')
+ })
+
+ it('delegates to SagemakerSpace for properties', function () {
+ const node = new SagemakerSpaceNode(testParent, client, testRegion, testSpaceApp)
+
+ // Verify that properties are managed by SagemakerSpace
+ assert.strictEqual(node.name, 'TestSpace')
+ assert.strictEqual(node.label, 'TestSpace (Running)')
+ assert.strictEqual(node.description, 'Private space')
+ assert.ok(node.tooltip instanceof vscode.MarkdownString)
+ })
+
+ it('updates space app status', async function () {
+ const describeSpaceStub = sinon.stub(SagemakerClient.prototype, 'describeSpace')
+ describeSpaceStub.resolves({ SpaceName: 'TestSpace', Status: 'InService', $metadata: {} })
+ describeAppStub.resolves({ AppName: 'TestApp', Status: 'InService', $metadata: {} })
+
+ await testSpaceAppNode.updateSpaceAppStatus()
+
+ sinon.assert.calledOnce(describeSpaceStub)
+ sinon.assert.calledOnce(describeAppStub)
})
})
diff --git a/packages/core/src/test/awsService/sagemaker/model.test.ts b/packages/core/src/test/awsService/sagemaker/model.test.ts
index 892baf2f77b..e6a9637ed15 100644
--- a/packages/core/src/test/awsService/sagemaker/model.test.ts
+++ b/packages/core/src/test/awsService/sagemaker/model.test.ts
@@ -59,6 +59,30 @@ describe('SageMaker Model', () => {
assert.ok(existsStub.callCount >= 3, 'should have retried for file existence')
})
+
+ it('throws ToolkitError when info file never appears', async function () {
+ sandbox.stub(fs, 'existsFile').resolves(false)
+ sandbox.stub(require('fs'), 'openSync').returns(42)
+ sandbox.replace(
+ require('../../../awsService/sagemaker/model'),
+ 'stopLocalServer',
+ sandbox.stub().resolves()
+ )
+ sandbox.replace(
+ require('../../../awsService/sagemaker/utils'),
+ 'spawnDetachedServer',
+ sandbox.stub().returns({ unref: sandbox.stub() })
+ )
+ sandbox.stub(DevSettings.instance, 'get').returns({})
+
+ try {
+ await startLocalServer(ctx)
+ assert.ok(false, 'Expected error not thrown')
+ } catch (err) {
+ assert.ok(err instanceof ToolkitError)
+ assert.ok(err.message.includes('Timed out waiting for local server info file'))
+ }
+ })
})
describe('stopLocalServer', function () {
@@ -106,6 +130,17 @@ describe('SageMaker Model', () => {
}
})
+ it('logs warning when process not found (ESRCH)', async function () {
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(validJson)
+ sandbox.stub(fs, 'delete').resolves()
+ sandbox.stub(process, 'kill').throws({ code: 'ESRCH', message: 'no such process' })
+
+ await stopLocalServer(ctx)
+
+ assertLogsContain(`no process found with PID ${validPid}. It may have already exited.`, false, 'warn')
+ })
+
it('throws ToolkitError when killing process fails for another reason', async function () {
sandbox.stub(fs, 'existsFile').resolves(true)
sandbox.stub(fs, 'readFileText').resolves(validJson)
@@ -120,6 +155,27 @@ describe('SageMaker Model', () => {
assert.strictEqual(err.message, 'failed to stop local server')
}
})
+
+ it('logs warning when PID is invalid', async function () {
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify({ pid: 'invalid' }))
+ sandbox.stub(fs, 'delete').resolves()
+
+ await stopLocalServer(ctx)
+
+ assertLogsContain('no valid PID found in info file.', false, 'warn')
+ })
+
+ it('logs warning when file deletion fails', async function () {
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(validJson)
+ sandbox.stub(process, 'kill').returns(true)
+ sandbox.stub(fs, 'delete').rejects(new Error('delete failed'))
+
+ await stopLocalServer(ctx)
+
+ assertLogsContain('could not delete info file: delete failed', false, 'warn')
+ })
})
describe('removeKnownHost', function () {
@@ -152,6 +208,38 @@ describe('SageMaker Model', () => {
sinon.match((value: string) => value.trim() === expectedOutput),
{ atomic: true }
)
+ assertLogsContain(`Removed '${hostname}' from known_hosts`, false, 'debug')
+ })
+
+ it('handles hostname in comma-separated list', async function () {
+ sandbox.stub(fs, 'existsFile').resolves(true)
+
+ const inputContent = `host1,${hostname},host2 ssh-rsa AAAA\nother.host ssh-rsa BBBB`
+ const expectedOutput = `other.host ssh-rsa BBBB`
+
+ sandbox.stub(fs, 'readFileText').resolves(inputContent)
+ const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+
+ await removeKnownHost(hostname)
+
+ sinon.assert.calledWith(
+ writeStub,
+ knownHostsPath,
+ sinon.match((value: string) => value.trim() === expectedOutput),
+ { atomic: true }
+ )
+ })
+
+ it('does not write file when hostname not found', async function () {
+ sandbox.stub(fs, 'existsFile').resolves(true)
+
+ const inputContent = `other.host ssh-rsa AAAA\nsome.other.com ssh-rsa BBBB`
+ sandbox.stub(fs, 'readFileText').resolves(inputContent)
+ const writeStub = sandbox.stub(fs, 'writeFile').resolves()
+
+ await removeKnownHost(hostname)
+
+ sinon.assert.notCalled(writeStub)
})
it('logs warning when known_hosts does not exist', async function () {
diff --git a/packages/core/src/test/awsService/sagemaker/sagemakerSpace.test.ts b/packages/core/src/test/awsService/sagemaker/sagemakerSpace.test.ts
new file mode 100644
index 00000000000..2a52b08a3a6
--- /dev/null
+++ b/packages/core/src/test/awsService/sagemaker/sagemakerSpace.test.ts
@@ -0,0 +1,129 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import { SagemakerSpace } from '../../../awsService/sagemaker/sagemakerSpace'
+import { SagemakerClient, SagemakerSpaceApp } from '../../../shared/clients/sagemaker'
+import sinon from 'sinon'
+
+describe('SagemakerSpace', function () {
+ let mockClient: sinon.SinonStubbedInstance
+ let mockSpaceApp: SagemakerSpaceApp
+
+ beforeEach(function () {
+ mockClient = sinon.createStubInstance(SagemakerClient)
+ mockSpaceApp = {
+ SpaceName: 'test-space',
+ Status: 'InService',
+ DomainId: 'test-domain',
+ DomainSpaceKey: 'test-key',
+ SpaceSettingsSummary: {
+ AppType: 'JupyterLab',
+ RemoteAccess: 'ENABLED',
+ },
+ }
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('updateSpaceAppStatus', function () {
+ it('should correctly map DescribeSpace API response to SagemakerSpaceApp type', async function () {
+ // Mock DescribeSpace response (uses full property names)
+ const mockDescribeSpaceResponse = {
+ SpaceName: 'updated-space',
+ Status: 'InService',
+ DomainId: 'test-domain',
+ SpaceSettings: {
+ // Note: 'SpaceSettings' not 'SpaceSettingsSummary'
+ AppType: 'CodeEditor',
+ RemoteAccess: 'DISABLED',
+ },
+ OwnershipSettings: {
+ OwnerUserProfileName: 'test-user',
+ },
+ SpaceSharingSettings: {
+ SharingType: 'Private',
+ },
+ $metadata: { requestId: 'test-request-id' },
+ }
+
+ // Mock DescribeApp response
+ const mockDescribeAppResponse = {
+ AppName: 'test-app',
+ Status: 'InService',
+ ResourceSpec: {
+ InstanceType: 'ml.t3.medium',
+ },
+ $metadata: { requestId: 'test-request-id' },
+ }
+
+ mockClient.describeSpace.resolves(mockDescribeSpaceResponse)
+ mockClient.describeApp.resolves(mockDescribeAppResponse)
+
+ const space = new SagemakerSpace(mockClient as any, 'us-east-1', mockSpaceApp)
+ const updateSpaceSpy = sinon.spy(space, 'updateSpace')
+
+ await space.updateSpaceAppStatus()
+
+ // Verify updateSpace was called with correctly mapped properties
+ assert.ok(updateSpaceSpy.calledOnce)
+ const updateSpaceArgs = updateSpaceSpy.getCall(0).args[0]
+
+ // Verify property name mapping from DescribeSpace to SagemakerSpaceApp
+ assert.strictEqual(updateSpaceArgs.SpaceSettingsSummary?.AppType, 'CodeEditor')
+ assert.strictEqual(updateSpaceArgs.SpaceSettingsSummary?.RemoteAccess, 'DISABLED')
+ assert.strictEqual(updateSpaceArgs.OwnershipSettingsSummary?.OwnerUserProfileName, 'test-user')
+ assert.strictEqual(updateSpaceArgs.SpaceSharingSettingsSummary?.SharingType, 'Private')
+
+ // Verify other properties are preserved
+ assert.strictEqual(updateSpaceArgs.SpaceName, 'updated-space')
+ assert.strictEqual(updateSpaceArgs.Status, 'InService')
+ assert.strictEqual(updateSpaceArgs.DomainId, 'test-domain')
+ assert.strictEqual(updateSpaceArgs.App, mockDescribeAppResponse)
+ assert.strictEqual(updateSpaceArgs.DomainSpaceKey, 'test-key')
+
+ // Verify original API property names are not present
+ assert.ok(!('SpaceSettings' in updateSpaceArgs))
+ assert.ok(!('OwnershipSettings' in updateSpaceArgs))
+ assert.ok(!('SpaceSharingSettings' in updateSpaceArgs))
+ })
+
+ it('should handle missing optional properties gracefully', async function () {
+ // Mock minimal DescribeSpace response
+ const mockDescribeSpaceResponse = {
+ SpaceName: 'minimal-space',
+ Status: 'InService',
+ DomainId: 'test-domain',
+ $metadata: { requestId: 'test-request-id' },
+ // No SpaceSettings, OwnershipSettings, or SpaceSharingSettings
+ }
+
+ const mockDescribeAppResponse = {
+ AppName: 'test-app',
+ Status: 'InService',
+ $metadata: { requestId: 'test-request-id' },
+ }
+
+ mockClient.describeSpace.resolves(mockDescribeSpaceResponse)
+ mockClient.describeApp.resolves(mockDescribeAppResponse)
+
+ const space = new SagemakerSpace(mockClient as any, 'us-east-1', mockSpaceApp)
+ const updateSpaceSpy = sinon.spy(space, 'updateSpace')
+
+ await space.updateSpaceAppStatus()
+
+ // Should not throw and should handle undefined properties
+ assert.ok(updateSpaceSpy.calledOnce)
+ const updateSpaceArgs = updateSpaceSpy.getCall(0).args[0]
+
+ assert.strictEqual(updateSpaceArgs.SpaceName, 'minimal-space')
+ assert.strictEqual(updateSpaceArgs.SpaceSettingsSummary, undefined)
+ assert.strictEqual(updateSpaceArgs.OwnershipSettingsSummary, undefined)
+ assert.strictEqual(updateSpaceArgs.SpaceSharingSettingsSummary, undefined)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/auth/connectionCredentialsProvider.test.ts b/packages/core/src/test/sagemakerunifiedstudio/auth/connectionCredentialsProvider.test.ts
new file mode 100644
index 00000000000..951e391d181
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/auth/connectionCredentialsProvider.test.ts
@@ -0,0 +1,215 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import { ConnectionCredentialsProvider } from '../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+import { SmusAuthenticationProvider } from '../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+import { DataZoneClient } from '../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { ToolkitError } from '../../../shared/errors'
+
+describe('ConnectionCredentialsProvider', function () {
+ let mockAuthProvider: sinon.SinonStubbedInstance
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+ let connectionProvider: ConnectionCredentialsProvider
+ let dataZoneClientStub: sinon.SinonStub
+
+ const testConnectionId = 'conn-123456'
+ const testDomainId = 'dzd_testdomain'
+ const testRegion = 'us-east-2'
+
+ const mockConnectionCredentials = {
+ accessKeyId: 'AKIA-CONNECTION-KEY',
+ secretAccessKey: 'connection-secret-key',
+ sessionToken: 'connection-session-token',
+ expiration: new Date(Date.now() + 3600000), // 1 hour from now
+ }
+
+ const mockGetConnectionResponse = {
+ connectionId: testConnectionId,
+ name: 'Test Connection',
+ type: 'S3',
+ domainId: testDomainId,
+ projectId: 'project-123',
+ connectionCredentials: mockConnectionCredentials,
+ }
+
+ beforeEach(function () {
+ // Mock auth provider
+ mockAuthProvider = {
+ isConnected: sinon.stub().returns(true),
+ getDomainId: sinon.stub().returns(testDomainId),
+ getDomainRegion: sinon.stub().returns(testRegion),
+ activeConnection: {
+ ssoRegion: testRegion,
+ },
+ } as any
+
+ // Mock DataZone client
+ mockDataZoneClient = {
+ getConnection: sinon.stub().resolves(mockGetConnectionResponse),
+ } as any
+
+ // Stub DataZoneClient.getInstance
+ dataZoneClientStub = sinon.stub(DataZoneClient, 'getInstance').resolves(mockDataZoneClient as any)
+
+ connectionProvider = new ConnectionCredentialsProvider(mockAuthProvider as any, testConnectionId)
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('should create provider with correct properties', function () {
+ assert.strictEqual(connectionProvider.getConnectionId(), testConnectionId)
+ assert.strictEqual(connectionProvider.getDefaultRegion(), testRegion)
+ })
+ })
+
+ describe('getCredentialsId', function () {
+ it('should return correct credentials ID', function () {
+ const credentialsId = connectionProvider.getCredentialsId()
+ assert.strictEqual(credentialsId.credentialSource, 'temp')
+ assert.strictEqual(credentialsId.credentialTypeId, `${testDomainId}:${testConnectionId}`)
+ })
+ })
+
+ describe('getHashCode', function () {
+ it('should return correct hash code', function () {
+ const hashCode = connectionProvider.getHashCode()
+ assert.strictEqual(hashCode, `smus-connection:${testDomainId}:${testConnectionId}`)
+ })
+ })
+
+ describe('isAvailable', function () {
+ it('should return true when auth provider is connected', async function () {
+ mockAuthProvider.isConnected.returns(true)
+ const isAvailable = await connectionProvider.isAvailable()
+ assert.strictEqual(isAvailable, true)
+ })
+
+ it('should return false when auth provider is not connected', async function () {
+ mockAuthProvider.isConnected.returns(false)
+ const isAvailable = await connectionProvider.isAvailable()
+ assert.strictEqual(isAvailable, false)
+ })
+
+ it('should return false when auth provider throws error', async function () {
+ mockAuthProvider.isConnected.throws(new Error('Connection error'))
+ const isAvailable = await connectionProvider.isAvailable()
+ assert.strictEqual(isAvailable, false)
+ })
+ })
+
+ describe('canAutoConnect', function () {
+ it('should return false', async function () {
+ const canAutoConnect = await connectionProvider.canAutoConnect()
+ assert.strictEqual(canAutoConnect, false)
+ })
+ })
+
+ describe('getCredentials', function () {
+ it('should fetch and return connection credentials', async function () {
+ const credentials = await connectionProvider.getCredentials()
+
+ assert.strictEqual(credentials.accessKeyId, mockConnectionCredentials.accessKeyId)
+ assert.strictEqual(credentials.secretAccessKey, mockConnectionCredentials.secretAccessKey)
+ assert.strictEqual(credentials.sessionToken, mockConnectionCredentials.sessionToken)
+ assert(credentials.expiration instanceof Date)
+
+ // Verify DataZone client was called correctly
+ sinon.assert.calledOnce(dataZoneClientStub)
+ sinon.assert.calledWith(mockDataZoneClient.getConnection, {
+ domainIdentifier: testDomainId,
+ identifier: testConnectionId,
+ withSecret: true,
+ })
+ })
+
+ it('should use cached credentials on subsequent calls', async function () {
+ // First call
+ const credentials1 = await connectionProvider.getCredentials()
+ // Second call
+ const credentials2 = await connectionProvider.getCredentials()
+
+ assert.strictEqual(credentials1, credentials2)
+ // DataZone client should only be called once due to caching
+ sinon.assert.calledOnce(mockDataZoneClient.getConnection)
+ })
+
+ it('should throw error when no connection credentials available', async function () {
+ mockDataZoneClient.getConnection.resolves({
+ ...mockGetConnectionResponse,
+ connectionCredentials: undefined,
+ })
+
+ await assert.rejects(
+ () => connectionProvider.getCredentials(),
+ (err: ToolkitError) => {
+ assert.strictEqual(err.code, 'NoConnectionCredentials')
+ return true
+ }
+ )
+ })
+
+ it('should throw error when connection credentials are invalid', async function () {
+ mockDataZoneClient.getConnection.resolves({
+ ...mockGetConnectionResponse,
+ connectionCredentials: {
+ accessKeyId: '', // Invalid empty string
+ secretAccessKey: 'valid-secret',
+ sessionToken: 'valid-token',
+ },
+ })
+
+ await assert.rejects(
+ () => connectionProvider.getCredentials(),
+ (err: ToolkitError) => {
+ assert.strictEqual(err.code, 'InvalidConnectionCredentials')
+ return true
+ }
+ )
+ })
+
+ it('should throw error when DataZone client fails', async function () {
+ const dataZoneError = new Error('DataZone API error')
+ mockDataZoneClient.getConnection.rejects(dataZoneError)
+
+ await assert.rejects(
+ () => connectionProvider.getCredentials(),
+ (err: ToolkitError) => {
+ assert.strictEqual(err.code, 'ConnectionCredentialsFetchFailed')
+ return true
+ }
+ )
+ })
+ })
+
+ describe('invalidate', function () {
+ it('should clear cached credentials', async function () {
+ // Get credentials to populate cache
+ await connectionProvider.getCredentials()
+ sinon.assert.calledOnce(mockDataZoneClient.getConnection)
+
+ // Invalidate cache
+ connectionProvider.invalidate()
+
+ // Get credentials again - should make new API call
+ await connectionProvider.getCredentials()
+ sinon.assert.calledTwice(mockDataZoneClient.getConnection)
+ })
+ })
+
+ describe('provider metadata', function () {
+ it('should return correct provider type', function () {
+ assert.strictEqual(connectionProvider.getProviderType(), 'temp')
+ })
+
+ it('should return correct telemetry type', function () {
+ assert.strictEqual(connectionProvider.getTelemetryType(), 'other')
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/auth/domainExecRoleCredentialsProvider.test.ts b/packages/core/src/test/sagemakerunifiedstudio/auth/domainExecRoleCredentialsProvider.test.ts
new file mode 100644
index 00000000000..7e8cdd8632d
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/auth/domainExecRoleCredentialsProvider.test.ts
@@ -0,0 +1,583 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import { DomainExecRoleCredentialsProvider } from '../../../sagemakerunifiedstudio/auth/providers/domainExecRoleCredentialsProvider'
+import { ToolkitError } from '../../../shared/errors'
+import fetch from 'node-fetch'
+import { SmusTimeouts } from '../../../sagemakerunifiedstudio/shared/smusUtils'
+
+describe('DomainExecRoleCredentialsProvider', function () {
+ let derProvider: DomainExecRoleCredentialsProvider
+ let mockGetAccessToken: sinon.SinonStub
+ let fetchStub: sinon.SinonStub
+
+ const testDomainId = 'dzd_testdomain'
+ const testDomainUrl = 'https://test-domain.sagemaker.us-east-2.on.aws'
+ const testSsoRegion = 'us-east-2'
+ const testAccessToken = 'test-access-token-12345'
+
+ const mockCredentialsResponse = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ },
+ }
+
+ beforeEach(function () {
+ // Mock access token function
+ mockGetAccessToken = sinon.stub().resolves(testAccessToken)
+
+ // Mock fetch
+ fetchStub = sinon.stub(fetch, 'default' as any).resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(mockCredentialsResponse)),
+ json: sinon.stub().resolves(mockCredentialsResponse),
+ } as any)
+
+ derProvider = new DomainExecRoleCredentialsProvider(
+ testDomainUrl,
+ testDomainId,
+ testSsoRegion,
+ mockGetAccessToken
+ )
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('should initialize with correct properties', function () {
+ assert.strictEqual(derProvider.getDomainId(), testDomainId)
+ assert.strictEqual(derProvider.getDomainUrl(), testDomainUrl)
+ assert.strictEqual(derProvider.getDefaultRegion(), testSsoRegion)
+ })
+ })
+
+ describe('getCredentialsId', function () {
+ it('should return correct credentials ID', function () {
+ const credentialsId = derProvider.getCredentialsId()
+ assert.strictEqual(credentialsId.credentialSource, 'sso')
+ assert.strictEqual(credentialsId.credentialTypeId, testDomainId)
+ })
+ })
+
+ describe('getProviderType', function () {
+ it('should return sso provider type', function () {
+ assert.strictEqual(derProvider.getProviderType(), 'sso')
+ })
+ })
+
+ describe('getTelemetryType', function () {
+ it('should return ssoProfile telemetry type', function () {
+ assert.strictEqual(derProvider.getTelemetryType(), 'ssoProfile')
+ })
+ })
+
+ describe('getHashCode', function () {
+ it('should return correct hash code', function () {
+ const hashCode = derProvider.getHashCode()
+ assert.strictEqual(hashCode, `smus-der:${testDomainId}:${testSsoRegion}`)
+ })
+ })
+
+ describe('canAutoConnect', function () {
+ it('should return false', async function () {
+ const result = await derProvider.canAutoConnect()
+ assert.strictEqual(result, false)
+ })
+ })
+
+ describe('isAvailable', function () {
+ it('should return true when access token is available', async function () {
+ const result = await derProvider.isAvailable()
+ assert.strictEqual(result, true)
+ assert.ok(mockGetAccessToken.called)
+ })
+
+ it('should return false when access token throws error', async function () {
+ mockGetAccessToken.rejects(new Error('Token error'))
+ const result = await derProvider.isAvailable()
+ assert.strictEqual(result, false)
+ })
+ })
+
+ describe('getCredentials', function () {
+ it('should fetch and cache DER credentials', async function () {
+ const credentials = await derProvider.getCredentials()
+
+ // Verify access token was fetched
+ assert.ok(mockGetAccessToken.called)
+
+ // Verify fetch was called with correct parameters
+ assert.ok(fetchStub.called)
+ const fetchCall = fetchStub.firstCall
+ assert.strictEqual(fetchCall.args[0], `${testDomainUrl}/sso/redeem-token`)
+
+ const fetchOptions = fetchCall.args[1]
+ assert.strictEqual(fetchOptions.method, 'POST')
+ assert.strictEqual(fetchOptions.headers['Content-Type'], 'application/json')
+ assert.strictEqual(fetchOptions.headers['Accept'], 'application/json')
+ assert.strictEqual(fetchOptions.headers['User-Agent'], 'aws-toolkit-vscode')
+
+ const requestBody = JSON.parse(fetchOptions.body)
+ assert.strictEqual(requestBody.domainId, testDomainId)
+ assert.strictEqual(requestBody.accessToken, testAccessToken)
+
+ // Verify timeout is set
+ assert.strictEqual(fetchOptions.timeout, SmusTimeouts.apiCallTimeoutMs)
+ assert.strictEqual(fetchOptions.timeout, 10000) // 10 seconds
+
+ // Verify returned credentials
+ assert.strictEqual(credentials.accessKeyId, mockCredentialsResponse.credentials.accessKeyId)
+ assert.strictEqual(credentials.secretAccessKey, mockCredentialsResponse.credentials.secretAccessKey)
+ assert.strictEqual(credentials.sessionToken, mockCredentialsResponse.credentials.sessionToken)
+ assert.ok(credentials.expiration)
+ })
+
+ it('should use cached credentials when available', async function () {
+ // First call should fetch credentials
+ const credentials1 = await derProvider.getCredentials()
+
+ // Second call should use cache
+ const credentials2 = await derProvider.getCredentials()
+
+ // Fetch should only be called once
+ assert.strictEqual(fetchStub.callCount, 1)
+ assert.strictEqual(mockGetAccessToken.callCount, 1)
+
+ // Credentials should be the same
+ assert.strictEqual(credentials1, credentials2)
+ })
+
+ it('should handle missing access token', async function () {
+ mockGetAccessToken.resolves('')
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'DerCredentialsFetchFailed' && err.message.includes('No access token available')
+ }
+ )
+ })
+
+ it('should handle HTTP errors from redeem token API', async function () {
+ fetchStub.resolves({
+ ok: false,
+ status: 401,
+ statusText: 'Unauthorized',
+ text: sinon.stub().resolves('Invalid token'),
+ } as any)
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'DerCredentialsFetchFailed' && err.message.includes('401')
+ }
+ )
+ })
+
+ it('should handle timeout errors', async function () {
+ const timeoutError = new Error('Request timeout')
+ timeoutError.name = 'AbortError'
+ fetchStub.rejects(timeoutError)
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return (
+ err.code === 'DerCredentialsFetchFailed' && err.message.includes('timed out after 10 seconds')
+ )
+ }
+ )
+ })
+
+ it('should handle network errors', async function () {
+ const networkError = new Error('Network error')
+ fetchStub.rejects(networkError)
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'DerCredentialsFetchFailed'
+ }
+ )
+ })
+
+ it('should handle missing credentials object in response', async function () {
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify({})),
+ json: sinon.stub().resolves({}), // Missing credentials object
+ } as any)
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return (
+ err.code === 'DerCredentialsFetchFailed' && err.message.includes('Missing credentials object')
+ )
+ }
+ )
+ })
+
+ it('should handle invalid accessKeyId in response', async function () {
+ const invalidResponse = {
+ credentials: {
+ accessKeyId: '', // Invalid empty string
+ secretAccessKey: 'valid-secret',
+ sessionToken: 'valid-token',
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(invalidResponse)),
+ json: sinon.stub().resolves(invalidResponse),
+ } as any)
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'DerCredentialsFetchFailed' && err.message.includes('Invalid accessKeyId')
+ }
+ )
+ })
+
+ it('should handle invalid secretAccessKey in response', async function () {
+ const invalidResponse = {
+ credentials: {
+ accessKeyId: 'valid-key',
+ secretAccessKey: undefined, // Invalid null value
+ sessionToken: 'valid-token',
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(invalidResponse)),
+ json: sinon.stub().resolves(invalidResponse),
+ } as any)
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'DerCredentialsFetchFailed' && err.message.includes('Invalid secretAccessKey')
+ }
+ )
+ })
+
+ it('should handle invalid sessionToken in response', async function () {
+ const invalidResponse = {
+ credentials: {
+ accessKeyId: 'valid-key',
+ secretAccessKey: 'valid-secret',
+ sessionToken: undefined, // Invalid undefined value
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(invalidResponse)),
+ json: sinon.stub().resolves(invalidResponse),
+ } as any)
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'DerCredentialsFetchFailed' && err.message.includes('Invalid sessionToken')
+ }
+ )
+ })
+
+ it('should set default expiration when not provided in response', async function () {
+ const credentials = await derProvider.getCredentials()
+
+ // Should have expiration set to 10 mins from now
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = Date.now() + 10 * 60 * 1000 // 10 minutes
+ const timeDiff = Math.abs(expirationTime - expectedTime)
+ assert.ok(timeDiff < 5000, 'Expiration should be 10 mins from now')
+ })
+
+ it('should use expiration from API response when provided as ISO string', async function () {
+ const futureExpiration = new Date(Date.now() + 2 * 60 * 60 * 1000) // 2 hours from now
+ const responseWithExpiration = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ expiration: futureExpiration.toISOString(), // API returns expiration as ISO string
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(responseWithExpiration)),
+ json: sinon.stub().resolves(responseWithExpiration),
+ } as any)
+
+ const credentials = await derProvider.getCredentials()
+
+ // Should use the expiration from the API response
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = futureExpiration.getTime()
+ const timeDiff = Math.abs(expirationTime - expectedTime)
+ assert.ok(timeDiff < 1000, 'Should use expiration from API response')
+ })
+
+ it('should handle epoch timestamp in seconds from API response', async function () {
+ const futureTime = Math.floor(Date.now() / 1000) + 3600 // 1 hour from now in seconds
+ const responseWithEpochExpiration = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ expiration: futureTime.toString(), // Epoch timestamp in seconds as string
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(responseWithEpochExpiration)),
+ json: sinon.stub().resolves(responseWithEpochExpiration),
+ } as any)
+
+ const credentials = await derProvider.getCredentials()
+
+ // Should correctly parse epoch timestamp and convert to Date
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = futureTime * 1000 // Convert to milliseconds
+ const timeDiff = Math.abs(expirationTime - expectedTime)
+ assert.ok(timeDiff < 1000, 'Should correctly parse epoch timestamp in seconds')
+ })
+
+ it('should handle epoch timestamp as number from API response', async function () {
+ const futureTime = Math.floor(Date.now() / 1000) + 7200 // 2 hours from now in seconds
+ const responseWithEpochExpiration = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ expiration: futureTime, // Epoch timestamp in seconds as number
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(responseWithEpochExpiration)),
+ json: sinon.stub().resolves(responseWithEpochExpiration),
+ } as any)
+
+ const credentials = await derProvider.getCredentials()
+
+ // Should correctly parse epoch timestamp and convert to Date
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = futureTime * 1000 // Convert to milliseconds
+ const timeDiff = Math.abs(expirationTime - expectedTime)
+ assert.ok(timeDiff < 1000, 'Should correctly parse epoch timestamp as number')
+ })
+
+ it('should handle zero epoch timestamp gracefully', async function () {
+ const responseWithZeroExpiration = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ expiration: '0', // Zero is not > 0, so treated as ISO string "0" which represents year 0
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(responseWithZeroExpiration)),
+ json: sinon.stub().resolves(responseWithZeroExpiration),
+ } as any)
+
+ const credentials = await derProvider.getCredentials()
+
+ // "0" is parsed as a valid date (year 0), not as an invalid date
+ // So it should use the parsed date, not the default expiration
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = new Date('0').getTime() // Year 0
+ assert.strictEqual(expirationTime, expectedTime, 'Should use parsed date for year 0')
+ })
+
+ it('should handle negative epoch timestamp gracefully', async function () {
+ const responseWithNegativeExpiration = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ expiration: '-1', // Negative is not > 0, so treated as ISO string "-1" which represents year -1
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(responseWithNegativeExpiration)),
+ json: sinon.stub().resolves(responseWithNegativeExpiration),
+ } as any)
+
+ const credentials = await derProvider.getCredentials()
+
+ // "-1" is parsed as a valid date (year -1), not as an invalid date
+ // So it should use the parsed date, not the default expiration
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = new Date('-1').getTime() // Year -1
+ assert.strictEqual(expirationTime, expectedTime, 'Should use parsed date for year -1')
+ })
+
+ it('should handle JSON parsing errors', async function () {
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves('invalid json'),
+ json: sinon.stub().rejects(new Error('Invalid JSON')),
+ } as any)
+
+ await assert.rejects(
+ () => derProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'DerCredentialsFetchFailed'
+ }
+ )
+ })
+
+ it('should handle invalid expiration string in response', async function () {
+ const responseWithInvalidExpiration = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ expiration: 'invalid-date-string', // Invalid date string
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(responseWithInvalidExpiration)),
+ json: sinon.stub().resolves(responseWithInvalidExpiration),
+ } as any)
+
+ const credentials = await derProvider.getCredentials()
+
+ // Should fall back to default expiration when date parsing fails
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+
+ // Should be a valid timestamp (not NaN) using the default expiration
+ assert.ok(!isNaN(expirationTime), 'Should have valid expiration timestamp')
+
+ // Should be close to now + 10 minutes (default expiration)
+ const expectedTime = Date.now() + 10 * 60 * 1000
+ const timeDiff = Math.abs(expirationTime - expectedTime)
+ assert.ok(timeDiff < 5000, 'Should fall back to default expiration for invalid date string')
+ })
+
+ it('should handle empty expiration string in response', async function () {
+ const responseWithEmptyExpiration = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ expiration: '', // Empty string
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(responseWithEmptyExpiration)),
+ json: sinon.stub().resolves(responseWithEmptyExpiration),
+ } as any)
+
+ const credentials = await derProvider.getCredentials()
+
+ // Should fall back to default expiration for empty string
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = Date.now() + 10 * 60 * 1000 // Default 10 minutes
+ const timeDiff = Math.abs(expirationTime - expectedTime)
+ assert.ok(timeDiff < 5000, 'Should use default expiration for empty string')
+ })
+
+ it('should handle non-numeric string that looks like a number', async function () {
+ const responseWithInvalidNumber = {
+ credentials: {
+ accessKeyId: 'AKIA-DER-KEY',
+ secretAccessKey: 'der-secret-key',
+ sessionToken: 'der-session-token',
+ expiration: '123abc', // Non-numeric string
+ },
+ }
+
+ fetchStub.resolves({
+ ok: true,
+ status: 200,
+ statusText: 'OK',
+ text: sinon.stub().resolves(JSON.stringify(responseWithInvalidNumber)),
+ json: sinon.stub().resolves(responseWithInvalidNumber),
+ } as any)
+
+ const credentials = await derProvider.getCredentials()
+
+ // Should fall back to default expiration for invalid numeric string
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = Date.now() + 10 * 60 * 1000 // Default 10 minutes
+ const timeDiff = Math.abs(expirationTime - expectedTime)
+ assert.ok(timeDiff < 5000, 'Should use default expiration for invalid numeric string')
+ })
+ })
+
+ describe('invalidate', function () {
+ it('should clear cache and force fresh fetch on next call', async function () {
+ // First call to populate cache
+ await derProvider.getCredentials()
+ assert.strictEqual(fetchStub.callCount, 1)
+
+ // Invalidate should clear cache
+ derProvider.invalidate()
+
+ // Next call should fetch fresh credentials
+ await derProvider.getCredentials()
+ assert.strictEqual(fetchStub.callCount, 2)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/auth/model.test.ts b/packages/core/src/test/sagemakerunifiedstudio/auth/model.test.ts
new file mode 100644
index 00000000000..a6ca72736e9
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/auth/model.test.ts
@@ -0,0 +1,232 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import {
+ createSmusProfile,
+ isValidSmusConnection,
+ scopeSmus,
+ SmusConnection,
+} from '../../../sagemakerunifiedstudio/auth/model'
+import { SsoConnection } from '../../../auth/connection'
+
+describe('SMUS Auth Model', function () {
+ const testDomainUrl = 'https://dzd_domainId.sagemaker.us-east-2.on.aws'
+ const testDomainId = 'dzd_domainId'
+ const testStartUrl = 'https://identitycenter.amazonaws.com/ssoins-testInstanceId'
+ const testRegion = 'us-east-2'
+
+ describe('scopeSmus', function () {
+ it('should have correct scope value', function () {
+ assert.strictEqual(scopeSmus, 'datazone:domain:access')
+ })
+ })
+
+ describe('createSmusProfile', function () {
+ it('should create profile with default scopes', function () {
+ const profile = createSmusProfile(testDomainUrl, testDomainId, testStartUrl, testRegion)
+
+ assert.strictEqual(profile.domainUrl, testDomainUrl)
+ assert.strictEqual(profile.domainId, testDomainId)
+ assert.strictEqual(profile.startUrl, testStartUrl)
+ assert.strictEqual(profile.ssoRegion, testRegion)
+ assert.strictEqual(profile.type, 'sso')
+ assert.deepStrictEqual(profile.scopes, [scopeSmus])
+ })
+
+ it('should create profile with custom scopes', function () {
+ const customScopes = ['custom:scope', 'another:scope']
+ const profile = createSmusProfile(testDomainUrl, testDomainId, testStartUrl, testRegion, customScopes)
+
+ assert.strictEqual(profile.domainUrl, testDomainUrl)
+ assert.strictEqual(profile.domainId, testDomainId)
+ assert.strictEqual(profile.startUrl, testStartUrl)
+ assert.strictEqual(profile.ssoRegion, testRegion)
+ assert.strictEqual(profile.type, 'sso')
+ assert.deepStrictEqual(profile.scopes, customScopes)
+ })
+
+ it('should create profile with all required properties', function () {
+ const profile = createSmusProfile(testDomainUrl, testDomainId, testStartUrl, testRegion)
+
+ // Check SsoProfile properties
+ assert.strictEqual(profile.type, 'sso')
+ assert.strictEqual(profile.startUrl, testStartUrl)
+ assert.strictEqual(profile.ssoRegion, testRegion)
+ assert.ok(Array.isArray(profile.scopes))
+
+ // Check SmusProfile properties
+ assert.strictEqual(profile.domainUrl, testDomainUrl)
+ assert.strictEqual(profile.domainId, testDomainId)
+ })
+ })
+
+ describe('isValidSmusConnection', function () {
+ it('should return true for valid SMUS connection', function () {
+ const validConnection = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ scopes: [scopeSmus],
+ label: 'Test SMUS Connection',
+ domainUrl: testDomainUrl,
+ domainId: testDomainId,
+ } as SmusConnection
+
+ assert.strictEqual(isValidSmusConnection(validConnection), true)
+ })
+
+ it('should return false for connection without SMUS scope', function () {
+ const connectionWithoutScope = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ scopes: ['sso:account:access'],
+ label: 'Test Connection',
+ domainUrl: testDomainUrl,
+ domainId: testDomainId,
+ } as any
+
+ assert.strictEqual(isValidSmusConnection(connectionWithoutScope), false)
+ })
+
+ it('should return false for connection without SMUS properties', function () {
+ const connectionWithoutSmusProps = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ scopes: [scopeSmus],
+ label: 'Test Connection',
+ } as SsoConnection
+
+ assert.strictEqual(isValidSmusConnection(connectionWithoutSmusProps), false)
+ })
+
+ it('should return false for non-SSO connection', function () {
+ const nonSsoConnection = {
+ id: 'test-connection-id',
+ type: 'iam',
+ label: 'Test IAM Connection',
+ domainUrl: testDomainUrl,
+ domainId: testDomainId,
+ scopes: [scopeSmus],
+ }
+
+ assert.strictEqual(isValidSmusConnection(nonSsoConnection), false)
+ })
+
+ it('should return false for undefined connection', function () {
+ assert.strictEqual(isValidSmusConnection(undefined), false)
+ })
+
+ it('should return false for null connection', function () {
+ assert.strictEqual(isValidSmusConnection(undefined), false)
+ })
+
+ it('should return false for connection without scopes', function () {
+ const connectionWithoutScopes = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ label: 'Test Connection',
+ domainUrl: testDomainUrl,
+ domainId: testDomainId,
+ }
+
+ assert.strictEqual(isValidSmusConnection(connectionWithoutScopes), false)
+ })
+
+ it('should return false for connection with empty scopes array', function () {
+ const connectionWithEmptyScopes = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ scopes: [],
+ label: 'Test Connection',
+ domainUrl: testDomainUrl,
+ domainId: testDomainId,
+ }
+
+ assert.strictEqual(isValidSmusConnection(connectionWithEmptyScopes), false)
+ })
+
+ it('should return true for connection with SMUS scope among other scopes', function () {
+ const connectionWithMultipleScopes = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ scopes: ['sso:account:access', scopeSmus, 'other:scope'],
+ label: 'Test SMUS Connection',
+ domainUrl: testDomainUrl,
+ domainId: testDomainId,
+ } as SmusConnection
+
+ assert.strictEqual(isValidSmusConnection(connectionWithMultipleScopes), true)
+ })
+
+ it('should return false for connection missing domainUrl', function () {
+ const connectionMissingDomainUrl = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ scopes: [scopeSmus],
+ label: 'Test Connection',
+ domainId: testDomainId,
+ }
+
+ assert.strictEqual(isValidSmusConnection(connectionMissingDomainUrl), false)
+ })
+
+ it('should return false for connection missing domainId', function () {
+ const connectionMissingDomainId = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ scopes: [scopeSmus],
+ label: 'Test Connection',
+ domainUrl: testDomainUrl,
+ }
+
+ assert.strictEqual(isValidSmusConnection(connectionMissingDomainId), false)
+ })
+ })
+
+ describe('SmusConnection interface', function () {
+ it('should extend both SmusProfile and SsoConnection', function () {
+ const connection = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: testStartUrl,
+ ssoRegion: testRegion,
+ scopes: [scopeSmus],
+ label: 'Test SMUS Connection',
+ domainUrl: testDomainUrl,
+ domainId: testDomainId,
+ } as SmusConnection
+
+ // Should have Connection properties
+ assert.strictEqual(connection.id, 'test-connection-id')
+ assert.strictEqual(connection.label, 'Test SMUS Connection')
+
+ // Should have SsoConnection properties
+ assert.strictEqual(connection.type, 'sso')
+ assert.strictEqual(connection.startUrl, testStartUrl)
+ assert.strictEqual(connection.ssoRegion, testRegion)
+ assert.ok(Array.isArray(connection.scopes))
+
+ // Should have SmusProfile properties
+ assert.strictEqual(connection.domainUrl, testDomainUrl)
+ assert.strictEqual(connection.domainId, testDomainId)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/auth/projectRoleCredentialsProvider.test.ts b/packages/core/src/test/sagemakerunifiedstudio/auth/projectRoleCredentialsProvider.test.ts
new file mode 100644
index 00000000000..6dd206593f8
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/auth/projectRoleCredentialsProvider.test.ts
@@ -0,0 +1,241 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import { ProjectRoleCredentialsProvider } from '../../../sagemakerunifiedstudio/auth/providers/projectRoleCredentialsProvider'
+import { DataZoneClient } from '../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { ToolkitError } from '../../../shared/errors'
+
+describe('ProjectRoleCredentialsProvider', function () {
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+ let mockSmusAuthProvider: any
+ let projectProvider: ProjectRoleCredentialsProvider
+ let dataZoneClientStub: sinon.SinonStub
+
+ const testProjectId = 'test-project-123'
+ const testDomainId = 'dzd_testdomain'
+ const testRegion = 'us-east-2'
+
+ const mockGetEnvironmentCredentialsResponse = {
+ accessKeyId: 'AKIA-PROJECT-KEY',
+ secretAccessKey: 'project-secret-key',
+ sessionToken: 'project-session-token',
+ expiration: new Date(Date.now() + 14 * 60 * 1000), // 14 minutes as Date object
+ $metadata: {
+ httpStatusCode: 200,
+ requestId: 'test-request-id',
+ },
+ }
+
+ beforeEach(function () {
+ // Mock SMUS auth provider
+ mockSmusAuthProvider = {
+ getDomainId: sinon.stub().returns(testDomainId),
+ getDomainRegion: sinon.stub().returns(testRegion),
+ isConnected: sinon.stub().returns(true),
+ } as any
+
+ // Mock DataZone client
+ mockDataZoneClient = {
+ getProjectDefaultEnvironmentCreds: sinon.stub().resolves(mockGetEnvironmentCredentialsResponse),
+ } as any
+
+ // Stub DataZoneClient.getInstance
+ dataZoneClientStub = sinon.stub(DataZoneClient, 'getInstance').resolves(mockDataZoneClient as any)
+
+ projectProvider = new ProjectRoleCredentialsProvider(mockSmusAuthProvider, testProjectId)
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('should initialize with DER provider and project ID', function () {
+ assert.strictEqual(projectProvider.getProjectId(), testProjectId)
+ })
+ })
+
+ describe('getCredentialsId', function () {
+ it('should return correct credentials ID', function () {
+ const credentialsId = projectProvider.getCredentialsId()
+ assert.strictEqual(credentialsId.credentialSource, 'temp')
+ assert.strictEqual(credentialsId.credentialTypeId, `${testDomainId}:${testProjectId}`)
+ })
+ })
+
+ describe('getProviderType', function () {
+ it('should return sso provider type', function () {
+ assert.strictEqual(projectProvider.getProviderType(), 'temp')
+ })
+ })
+
+ describe('getTelemetryType', function () {
+ it('should return smusProfile telemetry type', function () {
+ assert.strictEqual(projectProvider.getTelemetryType(), 'other')
+ })
+ })
+
+ describe('getDefaultRegion', function () {
+ it('should return DER provider default region', function () {
+ assert.strictEqual(projectProvider.getDefaultRegion(), testRegion)
+ })
+ })
+
+ describe('getHashCode', function () {
+ it('should return correct hash code', function () {
+ const hashCode = projectProvider.getHashCode()
+ assert.strictEqual(hashCode, `smus-project:${testDomainId}:${testProjectId}`)
+ })
+ })
+
+ describe('canAutoConnect', function () {
+ it('should return false', async function () {
+ const result = await projectProvider.canAutoConnect()
+ assert.strictEqual(result, false)
+ })
+ })
+
+ describe('isAvailable', function () {
+ it('should delegate to SMUS auth provider', async function () {
+ const result = await projectProvider.isAvailable()
+ assert.strictEqual(result, true)
+ assert.ok(mockSmusAuthProvider.isConnected.called)
+ })
+ })
+
+ describe('getCredentials', function () {
+ it('should fetch and cache project credentials', async function () {
+ const credentials = await projectProvider.getCredentials()
+
+ // Verify DataZone client getInstance was called
+ assert.ok(dataZoneClientStub.calledWith(mockSmusAuthProvider))
+
+ // Verify getProjectDefaultEnvironmentCreds was called
+ assert.ok(mockDataZoneClient.getProjectDefaultEnvironmentCreds.called)
+ assert.ok(mockDataZoneClient.getProjectDefaultEnvironmentCreds.calledWith(testProjectId))
+
+ // Verify returned credentials
+ assert.strictEqual(credentials.accessKeyId, mockGetEnvironmentCredentialsResponse.accessKeyId)
+ assert.strictEqual(credentials.secretAccessKey, mockGetEnvironmentCredentialsResponse.secretAccessKey)
+ assert.strictEqual(credentials.sessionToken, mockGetEnvironmentCredentialsResponse.sessionToken)
+ assert.ok(credentials.expiration)
+ })
+
+ it('should use cached credentials when available', async function () {
+ // First call should fetch credentials
+ const credentials1 = await projectProvider.getCredentials()
+
+ // Second call should use cache
+ const credentials2 = await projectProvider.getCredentials()
+
+ // DataZone client method should only be called once
+ assert.strictEqual(mockDataZoneClient.getProjectDefaultEnvironmentCreds.callCount, 1)
+
+ // Credentials should be the same
+ assert.strictEqual(credentials1, credentials2)
+ })
+
+ it('should handle DataZone client errors', async function () {
+ const error = new Error('DataZone client failed')
+ mockDataZoneClient.getProjectDefaultEnvironmentCreds.rejects(error)
+
+ await assert.rejects(
+ () => projectProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'ProjectCredentialsFetchFailed' && err.message.includes(testProjectId)
+ }
+ )
+ })
+
+ it('should handle GetEnvironmentCredentials API errors', async function () {
+ const error = new Error('API call failed')
+ mockDataZoneClient.getProjectDefaultEnvironmentCreds.rejects(error)
+
+ await assert.rejects(
+ () => projectProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'ProjectCredentialsFetchFailed'
+ }
+ )
+ })
+
+ it('should handle missing credentials in response', async function () {
+ mockDataZoneClient.getProjectDefaultEnvironmentCreds.resolves({
+ accessKeyId: undefined,
+ $metadata: {
+ httpStatusCode: 200,
+ requestId: 'test-request-id',
+ },
+ })
+
+ await assert.rejects(
+ () => projectProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'ProjectCredentialsFetchFailed'
+ }
+ )
+ })
+
+ it('should handle invalid credential fields', async function () {
+ const invalidResponse = {
+ accessKeyId: '', // Invalid empty string
+ secretAccessKey: 'valid-secret',
+ sessionToken: 'valid-token',
+ $metadata: {
+ httpStatusCode: 200,
+ requestId: 'test-request-id',
+ },
+ }
+ mockDataZoneClient.getProjectDefaultEnvironmentCreds.resolves(invalidResponse)
+
+ await assert.rejects(
+ () => projectProvider.getCredentials(),
+ (err: ToolkitError) => {
+ return err.code === 'ProjectCredentialsFetchFailed'
+ }
+ )
+ })
+
+ it('should use default expiration when not provided in response', async function () {
+ const responseWithoutExpiration = {
+ accessKeyId: 'AKIA-PROJECT-KEY',
+ secretAccessKey: 'project-secret-key',
+ sessionToken: 'project-session-token',
+ // No expiration field
+ $metadata: {
+ httpStatusCode: 200,
+ requestId: 'test-request-id',
+ },
+ }
+ mockDataZoneClient.getProjectDefaultEnvironmentCreds.resolves(responseWithoutExpiration)
+
+ const credentials = await projectProvider.getCredentials()
+
+ // Should have expiration set to ~10 minutes from now
+ assert.ok(credentials.expiration)
+ const expirationTime = credentials.expiration!.getTime()
+ const expectedTime = Date.now() + 10 * 60 * 1000
+ const timeDiff = Math.abs(expirationTime - expectedTime)
+ assert.ok(timeDiff < 5000, 'Expiration should be ~10 minutes from now')
+ })
+ })
+
+ describe('invalidate', function () {
+ it('should clear cache and force fresh fetch on next call', async function () {
+ // First call to populate cache
+ await projectProvider.getCredentials()
+ assert.strictEqual(mockDataZoneClient.getProjectDefaultEnvironmentCreds.callCount, 1)
+
+ // Invalidate should clear cache
+ projectProvider.invalidate()
+
+ // Next call should fetch fresh credentials
+ await projectProvider.getCredentials()
+ assert.strictEqual(mockDataZoneClient.getProjectDefaultEnvironmentCreds.callCount, 2)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/auth/smusAuthenticationProvider.test.ts b/packages/core/src/test/sagemakerunifiedstudio/auth/smusAuthenticationProvider.test.ts
new file mode 100644
index 00000000000..e00a050aedc
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/auth/smusAuthenticationProvider.test.ts
@@ -0,0 +1,410 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+
+// Mock the setContext function BEFORE importing modules that use it
+const setContextModule = require('../../../shared/vscode/setContext')
+const setContextStubGlobal = sinon.stub(setContextModule, 'setContext').resolves()
+
+import { SmusAuthenticationProvider } from '../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+import { SmusConnection } from '../../../sagemakerunifiedstudio/auth/model'
+import { DataZoneClient } from '../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { SmusUtils } from '../../../sagemakerunifiedstudio/shared/smusUtils'
+import { ToolkitError } from '../../../shared/errors'
+import * as messages from '../../../shared/utilities/messages'
+
+describe('SmusAuthenticationProvider', function () {
+ let mockAuth: any
+ let mockSecondaryAuth: any
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+ let smusAuthProvider: SmusAuthenticationProvider
+ let extractDomainInfoStub: sinon.SinonStub
+ let getSsoInstanceInfoStub: sinon.SinonStub
+ let isInSmusSpaceEnvironmentStub: sinon.SinonStub
+ let executeCommandStub: sinon.SinonStub
+ let mockSecondaryAuthState: {
+ activeConnection: SmusConnection | undefined
+ hasSavedConnection: boolean
+ isConnectionExpired: boolean
+ }
+
+ const testDomainUrl = 'https://dzd_domainId.sagemaker.us-east-2.on.aws'
+ const testDomainId = 'dzd_domainId'
+ const testRegion = 'us-east-2'
+ const testSsoInstanceInfo = {
+ issuerUrl: 'https://identitycenter.amazonaws.com/ssoins-testInstanceId',
+ ssoInstanceId: 'ssoins-testInstanceId',
+ clientId: 'arn:aws:sso::123456789:application/ssoins-testInstanceId/apl-testAppId',
+ region: testRegion,
+ }
+
+ const mockSmusConnection: SmusConnection = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: 'https://identitycenter.amazonaws.com/ssoins-testInstanceId',
+ ssoRegion: testRegion,
+ scopes: ['datazone:domain:access'],
+ label: 'Test SMUS Connection',
+ domainUrl: testDomainUrl,
+ domainId: testDomainId,
+ getToken: sinon.stub().resolves({ accessToken: 'mock-token', expiresAt: new Date() }),
+ getRegistration: sinon.stub().resolves({ clientId: 'mock-client', expiresAt: new Date() }),
+ }
+
+ beforeEach(function () {
+ mockAuth = {
+ createConnection: sinon.stub().resolves(mockSmusConnection),
+ listConnections: sinon.stub().resolves([]),
+ getConnectionState: sinon.stub().returns('valid'),
+ reauthenticate: sinon.stub().resolves(mockSmusConnection),
+ } as any
+
+ // Create a mock object with configurable properties
+ mockSecondaryAuthState = {
+ activeConnection: mockSmusConnection as SmusConnection | undefined,
+ hasSavedConnection: false,
+ isConnectionExpired: false,
+ }
+
+ mockSecondaryAuth = {
+ get activeConnection() {
+ return mockSecondaryAuthState.activeConnection
+ },
+ get hasSavedConnection() {
+ return mockSecondaryAuthState.hasSavedConnection
+ },
+ get isConnectionExpired() {
+ return mockSecondaryAuthState.isConnectionExpired
+ },
+ onDidChangeActiveConnection: sinon.stub().returns({ dispose: sinon.stub() }),
+ restoreConnection: sinon.stub().resolves(),
+ useNewConnection: sinon.stub().resolves(mockSmusConnection),
+ deleteConnection: sinon.stub().resolves(),
+ }
+
+ mockDataZoneClient = {
+ // Add any DataZoneClient methods that might be used
+ } as any
+
+ // Stub static methods
+ sinon.stub(DataZoneClient, 'getInstance').returns(mockDataZoneClient as any)
+ extractDomainInfoStub = sinon
+ .stub(SmusUtils, 'extractDomainInfoFromUrl')
+ .returns({ domainId: testDomainId, region: testRegion })
+ getSsoInstanceInfoStub = sinon.stub(SmusUtils, 'getSsoInstanceInfo').resolves(testSsoInstanceInfo)
+ isInSmusSpaceEnvironmentStub = sinon.stub(SmusUtils, 'isInSmusSpaceEnvironment').returns(false)
+ executeCommandStub = sinon.stub(vscode.commands, 'executeCommand').resolves()
+ sinon.stub(require('../../../auth/secondaryAuth'), 'getSecondaryAuth').returns(mockSecondaryAuth)
+
+ smusAuthProvider = new SmusAuthenticationProvider(mockAuth, mockSecondaryAuth)
+
+ // Reset the executeCommand stub for clean state
+ executeCommandStub.resetHistory()
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('should initialize with auth and secondary auth', function () {
+ assert.strictEqual(smusAuthProvider.auth, mockAuth)
+ assert.strictEqual(smusAuthProvider.secondaryAuth, mockSecondaryAuth)
+ })
+
+ it('should register event listeners', function () {
+ assert.ok(mockSecondaryAuth.onDidChangeActiveConnection.called)
+ })
+
+ it('should set initial context', async function () {
+ // Context should be set during construction (async call)
+ // Wait a bit for the async call to complete
+ await new Promise((resolve) => setTimeout(resolve, 0))
+ assert.ok(setContextStubGlobal.called)
+ })
+ })
+
+ describe('activeConnection', function () {
+ it('should return secondary auth active connection', function () {
+ assert.strictEqual(smusAuthProvider.activeConnection, mockSmusConnection)
+ })
+ })
+
+ describe('isUsingSavedConnection', function () {
+ it('should return secondary auth hasSavedConnection value', function () {
+ mockSecondaryAuthState.hasSavedConnection = true
+ assert.strictEqual(smusAuthProvider.isUsingSavedConnection, true)
+
+ mockSecondaryAuthState.hasSavedConnection = false
+ assert.strictEqual(smusAuthProvider.isUsingSavedConnection, false)
+ })
+ })
+
+ describe('isConnectionValid', function () {
+ it('should return true when connection exists and is not expired', function () {
+ mockSecondaryAuthState.activeConnection = mockSmusConnection
+ mockSecondaryAuthState.isConnectionExpired = false
+
+ assert.strictEqual(smusAuthProvider.isConnectionValid(), true)
+ })
+
+ it('should return false when no connection exists', function () {
+ mockSecondaryAuthState.activeConnection = undefined
+
+ assert.strictEqual(smusAuthProvider.isConnectionValid(), false)
+ })
+
+ it('should return false when connection is expired', function () {
+ mockSecondaryAuthState.activeConnection = mockSmusConnection
+ mockSecondaryAuthState.isConnectionExpired = true
+
+ assert.strictEqual(smusAuthProvider.isConnectionValid(), false)
+ })
+ })
+
+ describe('isConnected', function () {
+ it('should return true when active connection exists', function () {
+ mockSecondaryAuthState.activeConnection = mockSmusConnection
+ assert.strictEqual(smusAuthProvider.isConnected(), true)
+ })
+
+ it('should return false when no active connection', function () {
+ mockSecondaryAuthState.activeConnection = undefined
+ assert.strictEqual(smusAuthProvider.isConnected(), false)
+ })
+ })
+
+ describe('restore', function () {
+ it('should call secondary auth restoreConnection', async function () {
+ await smusAuthProvider.restore()
+ assert.ok(mockSecondaryAuth.restoreConnection.called)
+ })
+ })
+
+ describe('connectToSmus', function () {
+ it('should create new connection when none exists', async function () {
+ mockAuth.listConnections.resolves([])
+
+ const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+ assert.strictEqual(result, mockSmusConnection)
+ assert.ok(extractDomainInfoStub.calledWith(testDomainUrl))
+ assert.ok(getSsoInstanceInfoStub.calledWith(testDomainUrl))
+ assert.ok(mockAuth.createConnection.called)
+ assert.ok(mockSecondaryAuth.useNewConnection.called)
+ assert.ok(executeCommandStub.calledWith('aws.smus.switchProject'))
+ })
+
+ it('should reuse existing valid connection', async function () {
+ const existingConnection = { ...mockSmusConnection, domainUrl: testDomainUrl.toLowerCase() }
+ mockAuth.listConnections.resolves([existingConnection])
+ mockAuth.getConnectionState.returns('valid')
+
+ const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+ assert.strictEqual(result, mockSmusConnection)
+ assert.ok(mockAuth.createConnection.notCalled)
+ assert.ok(mockSecondaryAuth.useNewConnection.calledWith(existingConnection))
+ assert.ok(executeCommandStub.calledWith('aws.smus.switchProject'))
+ })
+
+ it('should reauthenticate existing invalid connection', async function () {
+ const existingConnection = { ...mockSmusConnection, domainUrl: testDomainUrl.toLowerCase() }
+ mockAuth.listConnections.resolves([existingConnection])
+ mockAuth.getConnectionState.returns('invalid')
+
+ const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+ assert.strictEqual(result, mockSmusConnection)
+ assert.ok(mockAuth.reauthenticate.calledWith(existingConnection))
+ assert.ok(mockSecondaryAuth.useNewConnection.called)
+ assert.ok(executeCommandStub.calledWith('aws.smus.switchProject'))
+ })
+
+ it('should throw error for invalid domain URL', async function () {
+ extractDomainInfoStub.returns({ domainId: undefined, region: testRegion })
+
+ await assert.rejects(
+ () => smusAuthProvider.connectToSmus('invalid-url'),
+ (err: ToolkitError) => {
+ // The error is wrapped with FailedToConnect, but the original error should be in the cause
+ return err.code === 'FailedToConnect' && (err.cause as any)?.code === 'InvalidDomainUrl'
+ }
+ )
+ // Should not trigger project selection on error
+ assert.ok(executeCommandStub.notCalled)
+ })
+
+ it('should handle SmusUtils errors', async function () {
+ const error = new Error('SmusUtils error')
+ getSsoInstanceInfoStub.rejects(error)
+
+ await assert.rejects(
+ () => smusAuthProvider.connectToSmus(testDomainUrl),
+ (err: ToolkitError) => err.code === 'FailedToConnect'
+ )
+ // Should not trigger project selection on error
+ assert.ok(executeCommandStub.notCalled)
+ })
+
+ it('should handle auth creation errors', async function () {
+ const error = new Error('Auth creation failed')
+ mockAuth.createConnection.rejects(error)
+
+ await assert.rejects(
+ () => smusAuthProvider.connectToSmus(testDomainUrl),
+ (err: ToolkitError) => err.code === 'FailedToConnect'
+ )
+ // Should not trigger project selection on error
+ assert.ok(executeCommandStub.notCalled)
+ })
+
+ it('should not trigger project selection in SMUS space environment', async function () {
+ isInSmusSpaceEnvironmentStub.returns(true)
+ mockAuth.listConnections.resolves([])
+
+ const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+ assert.strictEqual(result, mockSmusConnection)
+ assert.ok(mockAuth.createConnection.called)
+ assert.ok(mockSecondaryAuth.useNewConnection.called)
+ assert.ok(executeCommandStub.notCalled)
+ })
+
+ it('should not trigger project selection when reusing connection in SMUS space environment', async function () {
+ isInSmusSpaceEnvironmentStub.returns(true)
+ const existingConnection = { ...mockSmusConnection, domainUrl: testDomainUrl.toLowerCase() }
+ mockAuth.listConnections.resolves([existingConnection])
+ mockAuth.getConnectionState.returns('valid')
+
+ const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+ assert.strictEqual(result, mockSmusConnection)
+ assert.ok(mockSecondaryAuth.useNewConnection.calledWith(existingConnection))
+ assert.ok(executeCommandStub.notCalled)
+ })
+
+ it('should not trigger project selection when reauthenticating in SMUS space environment', async function () {
+ isInSmusSpaceEnvironmentStub.returns(true)
+ const existingConnection = { ...mockSmusConnection, domainUrl: testDomainUrl.toLowerCase() }
+ mockAuth.listConnections.resolves([existingConnection])
+ mockAuth.getConnectionState.returns('invalid')
+
+ const result = await smusAuthProvider.connectToSmus(testDomainUrl)
+
+ assert.strictEqual(result, mockSmusConnection)
+ assert.ok(mockAuth.reauthenticate.calledWith(existingConnection))
+ assert.ok(mockSecondaryAuth.useNewConnection.called)
+ assert.ok(executeCommandStub.notCalled)
+ })
+ })
+
+ describe('reauthenticate', function () {
+ it('should call auth reauthenticate', async function () {
+ const result = await smusAuthProvider.reauthenticate(mockSmusConnection)
+
+ assert.strictEqual(result, mockSmusConnection)
+ assert.ok(mockAuth.reauthenticate.calledWith(mockSmusConnection))
+ })
+
+ it('should wrap auth errors in ToolkitError', async function () {
+ const error = new Error('Reauthentication failed')
+ mockAuth.reauthenticate.rejects(error)
+
+ await assert.rejects(
+ () => smusAuthProvider.reauthenticate(mockSmusConnection),
+ (err: ToolkitError) => err.message.includes('Unable to reauthenticate')
+ )
+ })
+ })
+
+ describe('showReauthenticationPrompt', function () {
+ it('should show reauthentication message', async function () {
+ const showReauthenticateMessageStub = sinon.stub(messages, 'showReauthenticateMessage').resolves()
+
+ await smusAuthProvider.showReauthenticationPrompt(mockSmusConnection)
+
+ assert.ok(showReauthenticateMessageStub.called)
+ const callArgs = showReauthenticateMessageStub.firstCall.args[0]
+ assert.ok(callArgs.message.includes('SageMaker Unified Studio'))
+ assert.strictEqual(callArgs.suppressId, 'smusConnectionExpired')
+ })
+ })
+
+ describe('getAccessToken', function () {
+ beforeEach(function () {
+ mockSecondaryAuthState.activeConnection = mockSmusConnection
+ mockAuth.getSsoAccessToken = sinon.stub().resolves('mock-access-token')
+ mockAuth.invalidateConnection = sinon.stub()
+ })
+
+ it('should return access token when successful', async function () {
+ const token = await smusAuthProvider.getAccessToken()
+
+ assert.strictEqual(token, 'mock-access-token')
+ assert.ok(mockAuth.getSsoAccessToken.calledWith(mockSmusConnection))
+ })
+
+ it('should throw error when no active connection', async function () {
+ mockSecondaryAuthState.activeConnection = undefined
+
+ await assert.rejects(
+ () => smusAuthProvider.getAccessToken(),
+ (err: ToolkitError) => err.code === 'NoActiveConnection'
+ )
+ })
+
+ it('should handle InvalidGrantException and mark connection for reauthentication', async function () {
+ const invalidGrantError = new Error('UnknownError')
+ invalidGrantError.name = 'InvalidGrantException'
+ mockAuth.getSsoAccessToken.rejects(invalidGrantError)
+
+ await assert.rejects(
+ () => smusAuthProvider.getAccessToken(),
+ (err: ToolkitError) => {
+ return (
+ err.code === 'RedeemAccessTokenFailed' &&
+ err.message.includes('Failed to retrieve SSO access token for connection')
+ )
+ }
+ )
+
+ // Verify connection was NOT invalidated (current implementation doesn't handle InvalidGrantException specially)
+ assert.ok(mockAuth.invalidateConnection.notCalled)
+ })
+
+ it('should handle other errors normally', async function () {
+ const genericError = new Error('Network error')
+ mockAuth.getSsoAccessToken.rejects(genericError)
+
+ await assert.rejects(
+ () => smusAuthProvider.getAccessToken(),
+ (err: ToolkitError) =>
+ err.message.includes('Failed to retrieve SSO access token for connection') &&
+ err.code === 'RedeemAccessTokenFailed'
+ )
+
+ // Verify connection was NOT invalidated for generic errors
+ assert.ok(mockAuth.invalidateConnection.notCalled)
+ })
+ })
+
+ describe('fromContext', function () {
+ it('should return singleton instance', function () {
+ const instance1 = SmusAuthenticationProvider.fromContext()
+ const instance2 = SmusAuthenticationProvider.fromContext()
+
+ assert.strictEqual(instance1, instance2)
+ })
+
+ it('should return instance property', function () {
+ const instance = SmusAuthenticationProvider.fromContext()
+ assert.strictEqual(SmusAuthenticationProvider.instance, instance)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/connectionMagicsSelector/activation.test.ts b/packages/core/src/test/sagemakerunifiedstudio/connectionMagicsSelector/activation.test.ts
new file mode 100644
index 00000000000..86e37c76444
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/connectionMagicsSelector/activation.test.ts
@@ -0,0 +1,11 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import assert from 'assert'
+
+describe('Connection magic selector test', function () {
+ it('example test', function () {
+ assert.ok(true)
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/activation.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/activation.test.ts
new file mode 100644
index 00000000000..fcb76325293
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/activation.test.ts
@@ -0,0 +1,483 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import { activate } from '../../../sagemakerunifiedstudio/explorer/activation'
+import {
+ SmusAuthenticationProvider,
+ setSmusConnectedContext,
+} from '../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+import { DataZoneClient } from '../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { ResourceTreeDataProvider } from '../../../shared/treeview/resourceTreeDataProvider'
+import { SageMakerUnifiedStudioRootNode } from '../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode'
+import { getLogger } from '../../../shared/logger/logger'
+import { getTestWindow } from '../../shared/vscode/window'
+import { SeverityLevel } from '../../shared/vscode/message'
+import * as extensionUtilities from '../../../shared/extensionUtilities'
+
+describe('SMUS Explorer Activation', function () {
+ let mockExtensionContext: vscode.ExtensionContext
+ let mockSmusAuthProvider: sinon.SinonStubbedInstance
+ let mockTreeView: sinon.SinonStubbedInstance>
+ let mockTreeDataProvider: sinon.SinonStubbedInstance
+ let mockSmusRootNode: sinon.SinonStubbedInstance
+ let createTreeViewStub: sinon.SinonStub
+ let registerCommandStub: sinon.SinonStub
+ let dataZoneDisposeStub: sinon.SinonStub
+ let setupUserActivityMonitoringStub: sinon.SinonStub
+
+ beforeEach(function () {
+ mockExtensionContext = {
+ subscriptions: [],
+ } as any
+
+ mockSmusAuthProvider = {
+ restore: sinon.stub().resolves(),
+ isConnected: sinon.stub().returns(true),
+ reauthenticate: sinon.stub().resolves(),
+ onDidChange: sinon.stub().callsFake((_listener: () => void) => ({ dispose: sinon.stub() })),
+ activeConnection: {
+ id: 'test-connection',
+ domainId: 'test-domain',
+ ssoRegion: 'us-east-1',
+ },
+ } as any
+
+ mockTreeView = {
+ dispose: sinon.stub(),
+ } as any
+
+ mockTreeDataProvider = {
+ refresh: sinon.stub(),
+ } as any
+
+ mockSmusRootNode = {
+ getChildren: sinon.stub().resolves([]),
+ getProjectSelectNode: sinon.stub().returns({ refreshNode: sinon.stub().resolves() }),
+ } as any
+
+ // Stub vscode APIs
+ createTreeViewStub = sinon.stub(vscode.window, 'createTreeView').returns(mockTreeView as any)
+ registerCommandStub = sinon.stub(vscode.commands, 'registerCommand').returns({ dispose: sinon.stub() } as any)
+
+ // Stub SmusAuthenticationProvider
+ sinon.stub(SmusAuthenticationProvider, 'fromContext').returns(mockSmusAuthProvider as any)
+
+ // Stub DataZoneClient
+ dataZoneDisposeStub = sinon.stub(DataZoneClient, 'dispose')
+
+ // Stub SageMakerUnifiedStudioRootNode constructor
+ sinon.stub(SageMakerUnifiedStudioRootNode.prototype, 'getChildren').returns(mockSmusRootNode.getChildren())
+ sinon
+ .stub(SageMakerUnifiedStudioRootNode.prototype, 'getProjectSelectNode')
+ .returns(mockSmusRootNode.getProjectSelectNode())
+
+ // Stub ResourceTreeDataProvider constructor
+ sinon.stub(ResourceTreeDataProvider.prototype, 'refresh').value(mockTreeDataProvider.refresh)
+
+ // Stub logger
+ sinon.stub({ getLogger }, 'getLogger').returns({
+ debug: sinon.stub(),
+ info: sinon.stub(),
+ error: sinon.stub(),
+ } as any)
+
+ // Stub setSmusConnectedContext
+ sinon.stub({ setSmusConnectedContext }, 'setSmusConnectedContext').resolves()
+
+ // Stub setupUserActivityMonitoring
+ setupUserActivityMonitoringStub = sinon
+ .stub(require('../../../awsService/sagemaker/sagemakerSpace'), 'setupUserActivityMonitoring')
+ .resolves()
+
+ // Stub isSageMaker to return true for SMUS
+ sinon.stub(extensionUtilities, 'isSageMaker').returns(true)
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('activate', function () {
+ it('should initialize SMUS authentication provider and call restore', async function () {
+ await activate(mockExtensionContext)
+
+ assert.ok((SmusAuthenticationProvider.fromContext as sinon.SinonStub).called)
+ assert.ok(mockSmusAuthProvider.restore.called)
+ })
+
+ it('should create tree view with correct configuration', async function () {
+ await activate(mockExtensionContext)
+
+ assert.ok(createTreeViewStub.calledWith('aws.smus.rootView'))
+ const createTreeViewArgs = createTreeViewStub.firstCall.args[1]
+ assert.ok('treeDataProvider' in createTreeViewArgs)
+ })
+
+ it('should register all required commands', async function () {
+ await activate(mockExtensionContext)
+
+ // Check that commands are registered
+ const registeredCommands = registerCommandStub.getCalls().map((call) => call.args[0])
+
+ assert.ok(registeredCommands.includes('aws.smus.rootView.refresh'))
+ assert.ok(registeredCommands.includes('aws.smus.projectView'))
+ assert.ok(registeredCommands.includes('aws.smus.refreshProject'))
+ assert.ok(registeredCommands.includes('aws.smus.switchProject'))
+ assert.ok(registeredCommands.includes('aws.smus.stopSpace'))
+ assert.ok(registeredCommands.includes('aws.smus.openRemoteConnection'))
+ assert.ok(registeredCommands.includes('aws.smus.reauthenticate'))
+ })
+
+ it('should add all disposables to extension context subscriptions', async function () {
+ await activate(mockExtensionContext)
+
+ // Should have multiple subscriptions added
+ assert.ok(mockExtensionContext.subscriptions.length > 0)
+ })
+
+ it('should refresh tree data provider on initialization', async function () {
+ await activate(mockExtensionContext)
+
+ assert.ok(mockTreeDataProvider.refresh.called)
+ })
+
+ it('should register DataZone client disposal', async function () {
+ await activate(mockExtensionContext)
+
+ // Find the DataZone dispose subscription - it should be the last one added
+ const subscriptions = mockExtensionContext.subscriptions
+ assert.ok(subscriptions.length > 0)
+
+ // The DataZone dispose subscription should be among the subscriptions
+ let dataZoneDisposeFound = false
+ for (const subscription of subscriptions) {
+ if (subscription && typeof subscription.dispose === 'function') {
+ // Try calling dispose and see if it calls DataZoneClient.dispose
+ const callCountBefore = dataZoneDisposeStub.callCount
+ subscription.dispose()
+ if (dataZoneDisposeStub.callCount > callCountBefore) {
+ dataZoneDisposeFound = true
+ break
+ }
+ }
+ }
+
+ assert.ok(dataZoneDisposeFound, 'Should register DataZone client disposal')
+ })
+
+ describe('command handlers', function () {
+ beforeEach(async function () {
+ await activate(mockExtensionContext)
+ })
+
+ it('should handle aws.smus.rootView.refresh command', async function () {
+ const refreshCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.rootView.refresh')
+
+ assert.ok(refreshCommand)
+
+ // Execute the command handler
+ await refreshCommand.args[1]()
+
+ assert.ok(mockTreeDataProvider.refresh.called)
+ })
+
+ it('should handle aws.smus.reauthenticate command with connection', async function () {
+ const reauthCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.reauthenticate')
+
+ assert.ok(reauthCommand)
+
+ const mockConnection = {
+ id: 'test-connection',
+ type: 'sso',
+ startUrl: 'https://identitycenter.amazonaws.com/ssoins-testInstanceId',
+ ssoRegion: 'us-east-1',
+ scopes: ['datazone:domain:access'],
+ label: 'Test Connection',
+ } as any
+
+ const testWindow = getTestWindow()
+
+ // Execute the command handler with connection
+ await reauthCommand.args[1](mockConnection)
+
+ assert.ok(mockSmusAuthProvider.reauthenticate.calledWith(mockConnection))
+ assert.ok(mockTreeDataProvider.refresh.called)
+
+ // Check that an information message was shown
+ const infoMessages = testWindow.shownMessages.filter(
+ (msg) => msg.severity === SeverityLevel.Information
+ )
+ assert.ok(infoMessages.length > 0, 'Should show information message')
+ assert.ok(infoMessages.some((msg) => msg.message.includes('Successfully reauthenticated')))
+ })
+
+ it('should handle aws.smus.reauthenticate command without connection', async function () {
+ const reauthCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.reauthenticate')
+
+ assert.ok(reauthCommand)
+
+ // Execute the command handler without connection
+ await reauthCommand.args[1]()
+
+ assert.ok(mockSmusAuthProvider.reauthenticate.notCalled)
+ })
+
+ it('should handle reauthentication errors', async function () {
+ const reauthCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.reauthenticate')
+
+ assert.ok(reauthCommand)
+
+ const mockConnection = {
+ id: 'test-connection',
+ type: 'sso',
+ startUrl: 'https://identitycenter.amazonaws.com/ssoins-testInstanceId',
+ ssoRegion: 'us-east-1',
+ scopes: ['datazone:domain:access'],
+ label: 'Test Connection',
+ } as any
+ const error = new Error('Reauthentication failed')
+ mockSmusAuthProvider.reauthenticate.rejects(error)
+
+ const testWindow = getTestWindow()
+
+ // Execute the command handler
+ await reauthCommand.args[1](mockConnection)
+
+ // Check that an error message was shown
+ const errorMessages = testWindow.shownMessages.filter((msg) => msg.severity === SeverityLevel.Error)
+ assert.ok(errorMessages.length > 0, 'Should show error message')
+ assert.ok(errorMessages.some((msg) => msg.message.includes('Failed to reauthenticate')))
+ })
+
+ it('should handle aws.smus.refreshProject command', async function () {
+ const refreshProjectCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.refreshProject')
+
+ assert.ok(refreshProjectCommand)
+
+ // Execute the command handler
+ await refreshProjectCommand.args[1]()
+
+ // Verify that getProjectSelectNode was called and refreshNode was called on the returned node
+ assert.ok(mockSmusRootNode.getProjectSelectNode.called)
+ const projectNode = mockSmusRootNode.getProjectSelectNode()
+ assert.ok((projectNode.refreshNode as sinon.SinonStub).called)
+ })
+
+ it('should handle aws.smus.stopSpace command with valid node', async function () {
+ const stopSpaceCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.stopSpace')
+
+ assert.ok(stopSpaceCommand)
+
+ const mockSpaceNode = {
+ resource: {
+ sageMakerClient: {},
+ DomainSpaceKey: 'test-space-key',
+ regionCode: 'us-east-1',
+ getParent: sinon.stub().returns({
+ getAuthProvider: sinon.stub().returns({
+ activeConnection: { domainId: 'test-domain' },
+ }),
+ getProjectId: sinon.stub().returns('test-project'),
+ }),
+ },
+ getParent: sinon.stub().returns({
+ getAuthProvider: sinon.stub().returns({
+ activeConnection: { domainId: 'test-domain' },
+ }),
+ getProjectId: sinon.stub().returns('test-project'),
+ }),
+ } as any
+
+ // Mock the stopSpace function
+ const stopSpaceStub = sinon.stub()
+ sinon.stub(require('../../../awsService/sagemaker/commands'), 'stopSpace').value(stopSpaceStub)
+
+ // Execute the command handler
+ await stopSpaceCommand.args[1](mockSpaceNode)
+
+ assert.ok(
+ stopSpaceStub.calledWith(
+ mockSpaceNode.resource,
+ mockExtensionContext,
+ mockSpaceNode.resource.sageMakerClient
+ )
+ )
+ })
+
+ it('should handle aws.smus.stopSpace command with invalid node', async function () {
+ const stopSpaceCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.stopSpace')
+
+ assert.ok(stopSpaceCommand)
+
+ const testWindow = getTestWindow()
+
+ // Execute the command handler with undefined node
+ await stopSpaceCommand.args[1](undefined)
+
+ // Check that a warning message was shown
+ const warningMessages = testWindow.shownMessages.filter((msg) => msg.severity === SeverityLevel.Warning)
+ assert.ok(warningMessages.length > 0, 'Should show warning message')
+ assert.ok(warningMessages.some((msg) => msg.message.includes('Space information is being refreshed')))
+ })
+
+ it('should handle aws.smus.openRemoteConnection command with valid node', async function () {
+ const openRemoteCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.openRemoteConnection')
+
+ assert.ok(openRemoteCommand)
+
+ const mockSpaceNode = {
+ resource: {
+ sageMakerClient: {},
+ DomainSpaceKey: 'test-space-key',
+ regionCode: 'us-east-1',
+ getParent: sinon.stub().returns({
+ getAuthProvider: sinon.stub().returns({
+ activeConnection: { domainId: 'test-domain' },
+ }),
+ getProjectId: sinon.stub().returns('test-project'),
+ }),
+ },
+ getParent: sinon.stub().returns({
+ getAuthProvider: sinon.stub().returns({
+ activeConnection: { domainId: 'test-domain' },
+ }),
+ getProjectId: sinon.stub().returns('test-project'),
+ }),
+ } as any
+
+ // Mock the openRemoteConnect function
+ const openRemoteConnectStub = sinon.stub()
+ sinon
+ .stub(require('../../../awsService/sagemaker/commands'), 'openRemoteConnect')
+ .value(openRemoteConnectStub)
+
+ // Execute the command handler
+ await openRemoteCommand.args[1](mockSpaceNode)
+
+ assert.ok(
+ openRemoteConnectStub.calledWith(
+ mockSpaceNode.resource,
+ mockExtensionContext,
+ mockSpaceNode.resource.sageMakerClient
+ )
+ )
+ })
+
+ it('should handle aws.smus.openRemoteConnection command with invalid node', async function () {
+ const openRemoteCommand = registerCommandStub
+ .getCalls()
+ .find((call) => call.args[0] === 'aws.smus.openRemoteConnection')
+
+ assert.ok(openRemoteCommand)
+
+ const testWindow = getTestWindow()
+
+ // Execute the command handler with undefined node
+ await openRemoteCommand.args[1](undefined)
+
+ // Check that a warning message was shown
+ const warningMessages = testWindow.shownMessages.filter((msg) => msg.severity === SeverityLevel.Warning)
+ assert.ok(warningMessages.length > 0, 'Should show warning message')
+ assert.ok(warningMessages.some((msg) => msg.message.includes('Space information is being refreshed')))
+ })
+ })
+
+ it('should propagate auth provider initialization errors', async function () {
+ const error = new Error('Auth provider initialization failed')
+ mockSmusAuthProvider.restore.rejects(error)
+
+ // Should throw the error since there's no error handling in activate()
+ await assert.rejects(() => activate(mockExtensionContext), /Auth provider initialization failed/)
+ })
+
+ it('should create root node with auth provider', async function () {
+ await activate(mockExtensionContext)
+
+ // Verify that SageMakerUnifiedStudioRootNode was created with the auth provider
+ assert.ok(createTreeViewStub.called)
+ const treeDataProvider = createTreeViewStub.firstCall.args[1].treeDataProvider
+ assert.ok(treeDataProvider)
+ })
+
+ // TODO: Fix the activation test
+ it.skip('should setup user activity monitoring', async function () {
+ await activate(mockExtensionContext)
+
+ assert.ok(setupUserActivityMonitoringStub.called)
+ })
+ })
+
+ describe('command registration', function () {
+ it('should register commands with correct names', async function () {
+ await activate(mockExtensionContext)
+
+ const expectedCommands = [
+ 'aws.smus.rootView.refresh',
+ 'aws.smus.projectView',
+ 'aws.smus.refreshProject',
+ 'aws.smus.switchProject',
+ 'aws.smus.stopSpace',
+ 'aws.smus.openRemoteConnection',
+ 'aws.smus.reauthenticate',
+ ]
+
+ const registeredCommands = registerCommandStub.getCalls().map((call) => call.args[0])
+
+ for (const command of expectedCommands) {
+ assert.ok(registeredCommands.includes(command), `Command ${command} should be registered`)
+ }
+ })
+
+ it('should register commands that return disposables', async function () {
+ await activate(mockExtensionContext)
+
+ for (const call of registerCommandStub.getCalls()) {
+ const disposable = call.returnValue
+ assert.ok(disposable && typeof disposable.dispose === 'function')
+ }
+ })
+ })
+
+ describe('resource cleanup', function () {
+ it('should dispose DataZone client on extension deactivation', async function () {
+ await activate(mockExtensionContext)
+
+ // Find and execute the DataZone dispose subscription
+ const disposeSubscription = mockExtensionContext.subscriptions.find(
+ (sub) => sub.dispose && sub.dispose.toString().includes('DataZoneClient.dispose')
+ )
+
+ if (disposeSubscription) {
+ disposeSubscription.dispose()
+ assert.ok(dataZoneDisposeStub.called)
+ }
+ })
+
+ it('should add tree view to subscriptions for disposal', async function () {
+ await activate(mockExtensionContext)
+
+ assert.ok(mockExtensionContext.subscriptions.includes(mockTreeView))
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.test.ts
new file mode 100644
index 00000000000..6039e5e4e02
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.test.ts
@@ -0,0 +1,462 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import * as vscode from 'vscode'
+import {
+ LakehouseNode,
+ createLakehouseConnectionNode,
+} from '../../../../sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy'
+import { GlueCatalogClient } from '../../../../sagemakerunifiedstudio/shared/client/glueCatalogClient'
+import { GlueClient } from '../../../../sagemakerunifiedstudio/shared/client/glueClient'
+import { ConnectionClientStore } from '../../../../sagemakerunifiedstudio/shared/client/connectionClientStore'
+import { NodeType } from '../../../../sagemakerunifiedstudio/explorer/nodes/types'
+import { ConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+
+describe('LakehouseStrategy', function () {
+ let sandbox: sinon.SinonSandbox
+ let mockGlueCatalogClient: sinon.SinonStubbedInstance
+ let mockGlueClient: sinon.SinonStubbedInstance
+
+ const mockConnection = {
+ connectionId: 'lakehouse-conn-123',
+ name: 'test-lakehouse-connection',
+ type: 'ATHENA',
+ domainId: 'domain-123',
+ projectId: 'project-123',
+ }
+
+ const mockCredentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+
+ mockGlueCatalogClient = {
+ getCatalogs: sandbox.stub(),
+ } as any
+
+ mockGlueClient = {
+ getDatabases: sandbox.stub(),
+ getTables: sandbox.stub(),
+ getTable: sandbox.stub(),
+ } as any
+
+ sandbox.stub(GlueCatalogClient, 'createWithCredentials').returns(mockGlueCatalogClient as any)
+ sandbox.stub(GlueClient.prototype, 'getDatabases').callsFake(mockGlueClient.getDatabases)
+ sandbox.stub(GlueClient.prototype, 'getTables').callsFake(mockGlueClient.getTables)
+ sandbox.stub(GlueClient.prototype, 'getTable').callsFake(mockGlueClient.getTable)
+
+ const mockClientStore = {
+ getGlueClient: sandbox.stub().returns(mockGlueClient),
+ getGlueCatalogClient: sandbox.stub().returns(mockGlueCatalogClient),
+ }
+ sandbox.stub(ConnectionClientStore, 'getInstance').returns(mockClientStore as any)
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ describe('LakehouseNode', function () {
+ it('should initialize with correct properties', function () {
+ const nodeData = {
+ id: 'test-node',
+ nodeType: NodeType.CONNECTION,
+ value: { test: 'value' },
+ }
+
+ const node = new LakehouseNode(nodeData)
+
+ assert.strictEqual(node.id, 'test-node')
+ assert.deepStrictEqual(node.resource, { test: 'value' })
+ })
+
+ it('should return empty array for leaf nodes', async function () {
+ const nodeData = {
+ id: 'leaf-node',
+ nodeType: NodeType.REDSHIFT_COLUMN,
+ value: {},
+ }
+
+ const node = new LakehouseNode(nodeData)
+ const children = await node.getChildren()
+
+ assert.strictEqual(children.length, 0)
+ })
+
+ it('should return error node when children provider fails', async function () {
+ const nodeData = {
+ id: 'error-node',
+ nodeType: NodeType.CONNECTION,
+ value: {},
+ }
+
+ const failingProvider = async () => {
+ throw new Error('Provider failed')
+ }
+
+ const node = new LakehouseNode(nodeData, failingProvider)
+ const children = await node.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.ok(children[0].id.startsWith('error-node-error-getChildren-'))
+ })
+
+ it('should create correct tree item for column node', async function () {
+ const nodeData = {
+ id: 'column-node',
+ nodeType: NodeType.REDSHIFT_COLUMN,
+ value: { name: 'test_column', type: 'varchar' },
+ }
+
+ const node = new LakehouseNode(nodeData)
+ const treeItem = await node.getTreeItem()
+
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ assert.strictEqual(treeItem.description, 'varchar')
+ })
+
+ it('should cache children after first load', async function () {
+ const provider = sandbox
+ .stub()
+ .resolves([new LakehouseNode({ id: 'child', nodeType: NodeType.GLUE_DATABASE })])
+ const node = new LakehouseNode({ id: 'parent', nodeType: NodeType.CONNECTION }, provider)
+
+ await node.getChildren()
+ await node.getChildren()
+
+ assert.ok(provider.calledOnce)
+ })
+ })
+
+ describe('createLakehouseConnectionNode', function () {
+ it('should create connection node with correct structure', function () {
+ const node = createLakehouseConnectionNode(
+ mockConnection as any,
+ mockCredentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+
+ assert.strictEqual(node.id, 'lakehouse-conn-123')
+ assert.strictEqual(node.data.nodeType, NodeType.CONNECTION)
+ assert.strictEqual(node.data.path?.connection, 'test-lakehouse-connection')
+ })
+
+ it('should create AWS Data Catalog node for default connections', async function () {
+ const defaultConnection = {
+ ...mockConnection,
+ name: 'project.default_lakehouse',
+ }
+
+ mockGlueCatalogClient.getCatalogs.resolves({ catalogs: [], nextToken: undefined })
+ mockGlueClient.getDatabases.resolves({
+ databases: [{ Name: 'default-db' }],
+ nextToken: undefined,
+ })
+
+ const node = createLakehouseConnectionNode(
+ defaultConnection as any,
+ mockCredentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+ const children = await node.getChildren()
+
+ const awsDataCatalogNode = children.find((child) => child.id.includes('AwsDataCatalog')) as LakehouseNode
+ assert.ok(awsDataCatalogNode)
+ assert.strictEqual(awsDataCatalogNode.data.nodeType, NodeType.GLUE_CATALOG)
+ })
+
+ it('should not create AWS Data Catalog node for non-default connections', async function () {
+ mockGlueCatalogClient.getCatalogs.resolves({ catalogs: [], nextToken: undefined })
+
+ const node = createLakehouseConnectionNode(
+ mockConnection as any,
+ mockCredentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+ const children = await node.getChildren()
+
+ const awsDataCatalogNode = children.find((child) => child.id.includes('AwsDataCatalog'))
+ assert.strictEqual(awsDataCatalogNode, undefined)
+ })
+
+ it('should handle errors gracefully', async function () {
+ mockGlueCatalogClient.getCatalogs.rejects(new Error('Catalog error'))
+ mockGlueClient.getDatabases.rejects(new Error('Database error'))
+
+ const node = createLakehouseConnectionNode(
+ mockConnection as any,
+ mockCredentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+ const children = await node.getChildren()
+
+ assert.ok(children.length > 0)
+ assert.ok(children.some((child) => child.id.startsWith('lakehouse-conn-123-error-')))
+ })
+
+ it('should create placeholder when no catalogs found', async function () {
+ mockGlueCatalogClient.getCatalogs.resolves({ catalogs: [], nextToken: undefined })
+
+ const node = createLakehouseConnectionNode(
+ mockConnection as any,
+ mockCredentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+ const children = await node.getChildren()
+
+ assert.ok(children.some((child) => child.resource === '[No data found]'))
+ })
+ })
+
+ describe('Catalog nodes', function () {
+ it('should create catalog nodes from API', async function () {
+ mockGlueCatalogClient.getCatalogs.resolves({
+ catalogs: [{ CatalogId: 'test-catalog', CatalogType: 'HIVE' }],
+ })
+ mockGlueClient.getDatabases.resolves({
+ databases: [{ Name: 'test-db' }],
+ nextToken: undefined,
+ })
+
+ const node = createLakehouseConnectionNode(
+ mockConnection as any,
+ mockCredentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+ const children = await node.getChildren()
+
+ assert.ok(children.length > 0)
+ assert.ok(mockGlueCatalogClient.getCatalogs.called)
+ })
+
+ it('should handle catalog database pagination', async function () {
+ const catalogNode = new LakehouseNode(
+ {
+ id: 'catalog-node',
+ nodeType: NodeType.GLUE_CATALOG,
+ path: { catalog: 'test-catalog' },
+ },
+ async () => {
+ const allDatabases = []
+ let nextToken: string | undefined
+ do {
+ const { databases, nextToken: token } = await mockGlueClient.getDatabases(
+ 'test-catalog',
+ undefined,
+ undefined,
+ nextToken
+ )
+ allDatabases.push(...databases)
+ nextToken = token
+ } while (nextToken)
+ return allDatabases.map(
+ (db) => new LakehouseNode({ id: db.Name || '', nodeType: NodeType.GLUE_DATABASE })
+ )
+ }
+ )
+
+ mockGlueClient.getDatabases
+ .onFirstCall()
+ .resolves({ databases: [{ Name: 'db1' }], nextToken: 'token1' })
+ .onSecondCall()
+ .resolves({ databases: [{ Name: 'db2' }], nextToken: undefined })
+
+ const children = await catalogNode.getChildren()
+
+ assert.strictEqual(children.length, 2)
+ assert.ok(mockGlueClient.getDatabases.calledTwice)
+ })
+ })
+
+ describe('Database nodes', function () {
+ it('should handle table pagination', async function () {
+ const databaseNode = new LakehouseNode(
+ {
+ id: 'database-node',
+ nodeType: NodeType.GLUE_DATABASE,
+ path: { catalog: 'test-catalog', database: 'test-db' },
+ },
+ async () => {
+ const allTables = []
+ let nextToken: string | undefined
+ do {
+ const { tables, nextToken: token } = await mockGlueClient.getTables(
+ 'test-db',
+ 'test-catalog',
+ undefined,
+ nextToken
+ )
+ allTables.push(...tables)
+ nextToken = token
+ } while (nextToken)
+ return allTables.map(
+ (table) => new LakehouseNode({ id: table.Name || '', nodeType: NodeType.GLUE_TABLE })
+ )
+ }
+ )
+
+ mockGlueClient.getTables
+ .onFirstCall()
+ .resolves({ tables: [{ Name: 'table1' }], nextToken: 'token1' })
+ .onSecondCall()
+ .resolves({ tables: [{ Name: 'table2' }], nextToken: undefined })
+
+ const children = await databaseNode.getChildren()
+
+ assert.strictEqual(children.length, 2)
+ assert.ok(mockGlueClient.getTables.calledTwice)
+ })
+
+ it('should handle AWS Data Catalog database queries', async function () {
+ const databaseNode = new LakehouseNode(
+ {
+ id: 'database-node',
+ nodeType: NodeType.GLUE_DATABASE,
+ path: { catalog: 'aws-data-catalog', database: 'test-db' },
+ },
+ async () => {
+ const catalogId = undefined
+ const { tables } = await mockGlueClient.getTables('test-db', catalogId)
+ return tables.map(
+ (table) => new LakehouseNode({ id: table.Name || '', nodeType: NodeType.GLUE_TABLE })
+ )
+ }
+ )
+
+ mockGlueClient.getTables.resolves({ tables: [{ Name: 'aws-table' }], nextToken: undefined })
+
+ const children = await databaseNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.ok(mockGlueClient.getTables.calledWith('test-db', undefined))
+ })
+ })
+
+ describe('Table nodes', function () {
+ it('should create table node and load columns', async function () {
+ const tableNode = new LakehouseNode(
+ {
+ id: 'table-node',
+ nodeType: NodeType.GLUE_TABLE,
+ path: { database: 'test-db', table: 'test-table' },
+ },
+ async () => {
+ const tableDetails = await mockGlueClient.getTable('test-db', 'test-table')
+ const columns = tableDetails?.StorageDescriptor?.Columns || []
+ const partitions = tableDetails?.PartitionKeys || []
+ return [...columns, ...partitions].map(
+ (col) =>
+ new LakehouseNode({
+ id: `column-${col.Name}`,
+ nodeType: NodeType.REDSHIFT_COLUMN,
+ value: { name: col.Name, type: col.Type },
+ })
+ )
+ }
+ )
+
+ mockGlueClient.getTable.resolves({
+ StorageDescriptor: {
+ Columns: [{ Name: 'col1', Type: 'string' }],
+ },
+ PartitionKeys: [{ Name: 'partition_col', Type: 'date' }],
+ Name: undefined,
+ })
+
+ const children = await tableNode.getChildren()
+
+ assert.strictEqual(children.length, 2)
+ assert.ok(mockGlueClient.getTable.calledWith('test-db', 'test-table'))
+ })
+
+ it('should handle table with no columns', async function () {
+ const tableNode = new LakehouseNode(
+ {
+ id: 'empty-table-node',
+ nodeType: NodeType.GLUE_TABLE,
+ path: { database: 'test-db', table: 'empty-table' },
+ },
+ async () => {
+ const tableDetails = await mockGlueClient.getTable('test-db', 'empty-table')
+ const columns = tableDetails?.StorageDescriptor?.Columns || []
+ const partitions = tableDetails?.PartitionKeys || []
+ return [...columns, ...partitions].map(
+ (col) =>
+ new LakehouseNode({
+ id: `column-${col.Name}`,
+ nodeType: NodeType.REDSHIFT_COLUMN,
+ value: { name: col.Name, type: col.Type },
+ })
+ )
+ }
+ )
+
+ mockGlueClient.getTable.resolves({
+ StorageDescriptor: { Columns: [] },
+ PartitionKeys: [],
+ Name: undefined,
+ })
+
+ const children = await tableNode.getChildren()
+
+ assert.strictEqual(children.length, 0)
+ })
+
+ it('should handle table getTable errors gracefully', async function () {
+ const tableNode = new LakehouseNode(
+ {
+ id: 'error-table-node',
+ nodeType: NodeType.GLUE_TABLE,
+ path: { database: 'test-db', table: 'error-table' },
+ },
+ async () => {
+ try {
+ await mockGlueClient.getTable('test-db', 'error-table')
+ return []
+ } catch (err) {
+ return []
+ }
+ }
+ )
+
+ mockGlueClient.getTable.rejects(new Error('Table not found'))
+
+ const children = await tableNode.getChildren()
+
+ assert.strictEqual(children.length, 0)
+ })
+ })
+
+ describe('Column nodes', function () {
+ it('should create column node with correct properties', function () {
+ const parentNode = new LakehouseNode({
+ id: 'parent-table',
+ nodeType: NodeType.GLUE_TABLE,
+ path: { database: 'test-db', table: 'test-table' },
+ })
+
+ const columnNode = new LakehouseNode({
+ id: 'parent-table/test-column',
+ nodeType: NodeType.REDSHIFT_COLUMN,
+ value: { name: 'test-column', type: 'varchar' },
+ path: { database: 'test-db', table: 'test-table', column: 'test-column' },
+ parent: parentNode,
+ })
+
+ assert.strictEqual(columnNode.id, 'parent-table/test-column')
+ assert.strictEqual(columnNode.resource.name, 'test-column')
+ assert.strictEqual(columnNode.resource.type, 'varchar')
+ assert.strictEqual(columnNode.getParent(), parentNode)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/redshiftStrategy.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/redshiftStrategy.test.ts
new file mode 100644
index 00000000000..31ec0e8bb24
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/redshiftStrategy.test.ts
@@ -0,0 +1,358 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import * as vscode from 'vscode'
+import {
+ RedshiftNode,
+ createRedshiftConnectionNode,
+} from '../../../../sagemakerunifiedstudio/explorer/nodes/redshiftStrategy'
+import { SQLWorkbenchClient } from '../../../../sagemakerunifiedstudio/shared/client/sqlWorkbenchClient'
+import * as sqlWorkbenchClient from '../../../../sagemakerunifiedstudio/shared/client/sqlWorkbenchClient'
+import { ConnectionClientStore } from '../../../../sagemakerunifiedstudio/shared/client/connectionClientStore'
+import { NodeType } from '../../../../sagemakerunifiedstudio/explorer/nodes/types'
+import { ConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+
+describe('redshiftStrategy', function () {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ describe('RedshiftNode', function () {
+ describe('constructor', function () {
+ it('should create node with correct properties', function () {
+ const nodeData = {
+ id: 'test-id',
+ nodeType: NodeType.REDSHIFT_CLUSTER,
+ value: { clusterName: 'test-cluster' },
+ }
+
+ const node = new RedshiftNode(nodeData)
+
+ assert.strictEqual(node.id, 'test-id')
+ assert.strictEqual(node.data.nodeType, NodeType.REDSHIFT_CLUSTER)
+ assert.deepStrictEqual(node.resource, { clusterName: 'test-cluster' })
+ })
+ })
+
+ describe('getChildren', function () {
+ it('should return cached children if available', async function () {
+ const nodeData = {
+ id: 'test-id',
+ nodeType: NodeType.REDSHIFT_CLUSTER,
+ }
+
+ const node = new RedshiftNode(nodeData)
+ // Simulate cached children
+ ;(node as any).childrenNodes = [{ id: 'cached-child' }]
+
+ const children = await node.getChildren()
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual((children[0] as any).id, 'cached-child')
+ })
+
+ it('should return empty array for leaf nodes', async function () {
+ const nodeData = {
+ id: 'test-id',
+ nodeType: NodeType.REDSHIFT_COLUMN,
+ }
+
+ const node = new RedshiftNode(nodeData)
+ const children = await node.getChildren()
+ assert.strictEqual(children.length, 0)
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('should return correct tree item for regular nodes', async function () {
+ const nodeData = {
+ id: 'test-cluster',
+ nodeType: NodeType.REDSHIFT_CLUSTER,
+ value: { clusterName: 'test-cluster' },
+ }
+
+ const node = new RedshiftNode(nodeData)
+ const treeItem = await node.getTreeItem()
+
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Collapsed)
+ assert.strictEqual(treeItem.contextValue, NodeType.REDSHIFT_CLUSTER)
+ })
+
+ it('should return column tree item for column nodes', async function () {
+ const nodeData = {
+ id: 'test-column',
+ nodeType: NodeType.REDSHIFT_COLUMN,
+ value: { type: 'VARCHAR(255)' },
+ }
+
+ const node = new RedshiftNode(nodeData)
+ const treeItem = await node.getTreeItem()
+
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ assert.strictEqual(treeItem.description, 'VARCHAR(255)')
+ })
+
+ it('should return leaf tree item for leaf nodes', async function () {
+ const nodeData = {
+ id: 'test-column',
+ nodeType: NodeType.REDSHIFT_COLUMN,
+ }
+
+ const node = new RedshiftNode(nodeData)
+ const treeItem = await node.getTreeItem()
+
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ })
+ })
+
+ describe('getParent', function () {
+ it('should return parent node', function () {
+ const parentData = { id: 'parent', nodeType: NodeType.REDSHIFT_CLUSTER }
+ const parent = new RedshiftNode(parentData)
+
+ const nodeData = {
+ id: 'child',
+ nodeType: NodeType.REDSHIFT_DATABASE,
+ parent: parent,
+ }
+
+ const node = new RedshiftNode(nodeData)
+ assert.strictEqual(node.getParent(), parent)
+ })
+ })
+ })
+
+ describe('createRedshiftConnectionNode', function () {
+ let mockSQLClient: sinon.SinonStubbedInstance
+
+ beforeEach(function () {
+ mockSQLClient = {
+ executeQuery: sandbox.stub(),
+ getResources: sandbox.stub(),
+ } as any
+
+ sandbox.stub(SQLWorkbenchClient, 'createWithCredentials').returns(mockSQLClient as any)
+ sandbox.stub(sqlWorkbenchClient, 'createRedshiftConnectionConfig').resolves({
+ id: 'test-connection-id',
+ type: '4',
+ databaseType: 'REDSHIFT',
+ connectableResourceIdentifier: 'test-cluster',
+ connectableResourceType: 'CLUSTER',
+ database: 'test-db',
+ })
+
+ const mockClientStore = {
+ getSQLWorkbenchClient: sandbox.stub().returns(mockSQLClient),
+ }
+ sandbox.stub(ConnectionClientStore, 'getInstance').returns(mockClientStore as any)
+ })
+
+ it.skip('should create Redshift connection node with JDBC URL', async function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'Test Redshift Connection',
+ type: 'RedshiftConnection',
+ props: {
+ jdbcConnection: {
+ jdbcUrl: 'jdbc:redshift://test-cluster.123456789012.us-east-1.redshift.amazonaws.com:5439/dev',
+ dbname: 'test-db',
+ },
+ redshiftProperties: {},
+ },
+ location: {
+ awsAccountId: '',
+ awsRegion: 'us-east-1',
+ },
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ mockSQLClient.executeQuery.resolves('query-id')
+ mockSQLClient.getResources.resolves({
+ resources: [
+ {
+ displayName: 'test-db',
+ type: 'DATABASE',
+ identifier: '',
+ childObjectTypes: [],
+ },
+ ],
+ })
+
+ const node = createRedshiftConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider
+ )
+
+ assert.strictEqual(node.data.nodeType, NodeType.CONNECTION)
+ assert.strictEqual(node.data.value.connection.name, 'Test Redshift Connection')
+
+ // Test children provider - now creates database nodes directly
+ const children = await node.getChildren()
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual((children[0] as RedshiftNode).data.nodeType, NodeType.REDSHIFT_DATABASE)
+ })
+
+ it.skip('should create connection node with host from jdbcConnection', async function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'Test Connection',
+ type: 'RedshiftConnection',
+ props: {
+ jdbcConnection: {
+ host: 'test-host.redshift.amazonaws.com',
+ dbname: 'test-db',
+ },
+ redshiftProperties: {},
+ },
+ location: {
+ awsAccountId: '',
+ awsRegion: 'us-east-1',
+ },
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ mockSQLClient.executeQuery.resolves('query-id')
+ mockSQLClient.getResources.resolves({ resources: [] })
+
+ const node = createRedshiftConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider
+ )
+ const children = await node.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual((children[0] as RedshiftNode).data.nodeType, NodeType.REDSHIFT_DATABASE)
+ })
+
+ it('should return placeholder when connection params are missing', async function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'Test Connection',
+ type: 'RedshiftConnection',
+ props: {
+ jdbcConnection: {},
+ redshiftProperties: {},
+ },
+ location: {},
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ const node = createRedshiftConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider
+ )
+ const children = await node.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual(children[0].resource, '[No data found]')
+ })
+
+ it.skip('should handle workgroup name in host', async function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'Test Connection',
+ type: 'RedshiftConnection',
+ props: {
+ jdbcConnection: {
+ host: 'test-host.redshift-serverless.amazonaws.com',
+ dbname: 'test-db',
+ },
+ redshiftProperties: {
+ storage: {
+ workgroupName: 'test-workgroup',
+ },
+ },
+ },
+ location: {
+ awsAccountId: '',
+ awsRegion: 'us-east-1',
+ },
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ mockSQLClient.executeQuery.resolves('query-id')
+ mockSQLClient.getResources.resolves({ resources: [] })
+
+ const node = createRedshiftConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider
+ )
+ const children = await node.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ })
+
+ it.skip('should handle connection errors gracefully', async function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'Test Connection',
+ type: 'RedshiftConnection',
+ props: {
+ jdbcConnection: {
+ host: 'test-host.redshift.amazonaws.com',
+ dbname: 'test-db',
+ },
+ },
+ location: {
+ awsAccountId: '',
+ awsRegion: 'us-east-1',
+ },
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ // Make createRedshiftConnectionConfig throw an error
+ ;(sqlWorkbenchClient.createRedshiftConnectionConfig as sinon.SinonStub).rejects(
+ new Error('Connection config failed')
+ )
+
+ const node = createRedshiftConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider
+ )
+
+ // The error should be handled gracefully and return an error node
+ const children = await node.getChildren()
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual((children[0] as any).id.includes('error'), true)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/s3Strategy.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/s3Strategy.test.ts
new file mode 100644
index 00000000000..9b193ea0106
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/s3Strategy.test.ts
@@ -0,0 +1,287 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import * as vscode from 'vscode'
+import { S3Node, createS3ConnectionNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/s3Strategy'
+import { S3Client } from '../../../../sagemakerunifiedstudio/shared/client/s3Client'
+import { ConnectionClientStore } from '../../../../sagemakerunifiedstudio/shared/client/connectionClientStore'
+import { NodeType, ConnectionType } from '../../../../sagemakerunifiedstudio/explorer/nodes/types'
+import { ConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+
+describe('s3Strategy', function () {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ describe('S3Node', function () {
+ describe('constructor', function () {
+ it('should create node with correct properties', function () {
+ const node = new S3Node({
+ id: 'test-id',
+ nodeType: NodeType.S3_BUCKET,
+ connectionType: ConnectionType.S3,
+ value: { bucket: 'test-bucket' },
+ path: { bucket: 'test-bucket' },
+ })
+
+ assert.strictEqual(node.id, 'test-id')
+ assert.strictEqual(node.data.nodeType, NodeType.S3_BUCKET)
+ assert.strictEqual(node.data.connectionType, ConnectionType.S3)
+ })
+ })
+
+ describe('getChildren', function () {
+ it('should return empty array for leaf nodes', async function () {
+ const node = new S3Node({
+ id: 'test-id',
+ nodeType: NodeType.S3_FILE,
+ connectionType: ConnectionType.S3,
+ })
+
+ const children = await node.getChildren()
+ assert.strictEqual(children.length, 0)
+ })
+
+ it('should handle children provider errors', async function () {
+ const errorProvider = async () => {
+ throw new Error('Provider error')
+ }
+
+ const node = new S3Node(
+ {
+ id: 'test-id',
+ nodeType: NodeType.S3_BUCKET,
+ connectionType: ConnectionType.S3,
+ },
+ errorProvider
+ )
+
+ const children = await node.getChildren()
+ assert.strictEqual(children.length, 1)
+ assert.ok(children[0].id.startsWith('test-id-error-getChildren-'))
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('should return correct tree item for non-leaf node', async function () {
+ const node = new S3Node({
+ id: 'test-id',
+ nodeType: NodeType.S3_BUCKET,
+ connectionType: ConnectionType.S3,
+ })
+
+ const treeItem = await node.getTreeItem()
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Collapsed)
+ assert.strictEqual(treeItem.contextValue, NodeType.S3_BUCKET)
+ })
+
+ it('should return correct tree item for leaf node', async function () {
+ const node = new S3Node({
+ id: 'test-id',
+ nodeType: NodeType.S3_FILE,
+ connectionType: ConnectionType.S3,
+ })
+
+ const treeItem = await node.getTreeItem()
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ })
+ })
+ })
+
+ describe('createS3ConnectionNode', function () {
+ let mockS3Client: sinon.SinonStubbedInstance
+
+ beforeEach(function () {
+ mockS3Client = {
+ listPaths: sandbox.stub(),
+ } as any
+
+ sandbox.stub(S3Client.prototype, 'constructor' as any)
+ sandbox.stub(S3Client.prototype, 'listPaths').callsFake(mockS3Client.listPaths)
+
+ const mockClientStore = {
+ getS3Client: sandbox.stub().returns(mockS3Client),
+ }
+ sandbox.stub(ConnectionClientStore, 'getInstance').returns(mockClientStore as any)
+ })
+
+ it('should create S3 connection node successfully for non-default connection', function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'Test S3 Connection',
+ type: 'S3Connection',
+ props: {
+ s3Properties: {
+ s3Uri: 's3://test-bucket/prefix/',
+ },
+ },
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ const node = createS3ConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+
+ assert.strictEqual(node.data.nodeType, NodeType.CONNECTION)
+ assert.strictEqual(node.data.connectionType, ConnectionType.S3)
+ })
+
+ it('should create S3 connection node for default connection with full path', function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'project.s3_default_folder',
+ type: 'S3Connection',
+ props: {
+ s3Properties: {
+ s3Uri: 's3://test-bucket/domain/project/',
+ },
+ },
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ const node = createS3ConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+
+ assert.strictEqual(node.data.nodeType, NodeType.CONNECTION)
+ assert.strictEqual(node.data.connectionType, ConnectionType.S3)
+ })
+
+ it('should return error node when no S3 URI found', function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'Test S3 Connection',
+ type: 'S3Connection',
+ props: {},
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ const node = createS3ConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+
+ assert.ok(node.id.startsWith('conn-123-error-connection-'))
+ })
+
+ it('should handle bucket listing for non-default connection', async function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'Test S3 Connection',
+ type: 'S3Connection',
+ props: {
+ s3Properties: {
+ s3Uri: 's3://test-bucket/',
+ },
+ },
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ mockS3Client.listPaths.resolves({
+ paths: [
+ {
+ bucket: 'test-bucket',
+ prefix: 'file.txt',
+ displayName: 'file.txt',
+ isFolder: false,
+ },
+ ],
+ nextToken: undefined,
+ })
+
+ const node = createS3ConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+ const children = await node.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual((children[0] as S3Node).data.nodeType, NodeType.S3_BUCKET)
+ })
+
+ it('should handle bucket listing for default connection with full path display', async function () {
+ const connection = {
+ connectionId: 'conn-123',
+ name: 'project.s3_default_folder',
+ type: 'S3Connection',
+ props: {
+ s3Properties: {
+ s3Uri: 's3://test-bucket/domain/project/',
+ },
+ },
+ }
+
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ }),
+ }
+
+ mockS3Client.listPaths.resolves({
+ paths: [
+ {
+ bucket: 'test-bucket',
+ prefix: 'domain/project/dev/',
+ displayName: 'dev',
+ isFolder: true,
+ },
+ ],
+ nextToken: undefined,
+ })
+
+ const node = createS3ConnectionNode(
+ connection as any,
+ credentialsProvider as ConnectionCredentialsProvider,
+ 'us-east-1'
+ )
+ const children = await node.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ const bucketNode = children[0] as S3Node
+ assert.strictEqual(bucketNode.data.nodeType, NodeType.S3_BUCKET)
+ // For default connection, should show full path
+ assert.strictEqual(bucketNode.data.path?.label, 'test-bucket/domain/project/')
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioAuthInfoNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioAuthInfoNode.test.ts
new file mode 100644
index 00000000000..ebf2eae2cb0
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioAuthInfoNode.test.ts
@@ -0,0 +1,291 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import { SageMakerUnifiedStudioAuthInfoNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioAuthInfoNode'
+import { SmusAuthenticationProvider } from '../../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+import { SmusConnection } from '../../../../sagemakerunifiedstudio/auth/model'
+
+describe('SageMakerUnifiedStudioAuthInfoNode', function () {
+ let authInfoNode: SageMakerUnifiedStudioAuthInfoNode
+ let mockAuthProvider: any
+ let mockConnection: SmusConnection
+ let currentActiveConnection: SmusConnection | undefined
+
+ beforeEach(function () {
+ mockConnection = {
+ id: 'test-connection-id',
+ type: 'sso',
+ startUrl: 'https://identitycenter.amazonaws.com/ssoins-testInstanceId',
+ ssoRegion: 'us-east-2',
+ scopes: ['datazone:domain:access'],
+ label: 'Test SMUS Connection',
+ domainUrl: 'https://dzd_domainId.sagemaker.us-east-2.on.aws',
+ domainId: 'dzd_domainId',
+ // Mock the required methods from SsoConnection
+ getToken: sinon.stub().resolves(),
+ getRegistration: sinon.stub().resolves(),
+ } as any
+
+ // Initialize the current active connection
+ currentActiveConnection = mockConnection
+
+ // Create mock auth provider with getter for activeConnection
+ mockAuthProvider = {
+ isConnected: sinon.stub().returns(true),
+ isConnectionValid: sinon.stub().returns(true),
+ onDidChange: sinon.stub().callsFake((listener: () => void) => ({ dispose: sinon.stub() })),
+ get activeConnection() {
+ return currentActiveConnection
+ },
+ set activeConnection(value: SmusConnection | undefined) {
+ currentActiveConnection = value
+ },
+ }
+
+ // Stub SmusAuthenticationProvider.fromContext
+ sinon.stub(SmusAuthenticationProvider, 'fromContext').returns(mockAuthProvider as any)
+
+ authInfoNode = new SageMakerUnifiedStudioAuthInfoNode()
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('should initialize with correct properties', function () {
+ assert.strictEqual(authInfoNode.id, 'smusAuthInfoNode')
+ assert.strictEqual(authInfoNode.resource, authInfoNode)
+ })
+
+ it('should register for auth provider changes', function () {
+ assert.ok(mockAuthProvider.onDidChange.called)
+ })
+
+ it('should have onDidChangeTreeItem event', function () {
+ assert.ok(typeof authInfoNode.onDidChangeTreeItem === 'function')
+ })
+ })
+
+ describe('getTreeItem', function () {
+ describe('when connected and valid', function () {
+ beforeEach(function () {
+ mockAuthProvider.isConnected.returns(true)
+ mockAuthProvider.isConnectionValid.returns(true)
+ mockAuthProvider.activeConnection = mockConnection
+ })
+
+ it('should return connected tree item', function () {
+ const treeItem = authInfoNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Domain: dzd_domainId')
+ assert.strictEqual(treeItem.description, 'us-east-2')
+ assert.strictEqual(treeItem.contextValue, 'smusAuthInfo')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+
+ // Check icon
+ assert.ok(treeItem.iconPath instanceof vscode.ThemeIcon)
+ assert.strictEqual((treeItem.iconPath as vscode.ThemeIcon).id, 'key')
+
+ // Check tooltip
+ const tooltip = treeItem.tooltip as string
+ assert.ok(tooltip?.includes('Connected to SageMaker Unified Studio'))
+ assert.ok(tooltip?.includes('dzd_domainId'))
+ assert.ok(tooltip?.includes('us-east-2'))
+ assert.ok(tooltip?.includes('Status: Connected'))
+
+ // Should not have command when valid
+ assert.strictEqual(treeItem.command, undefined)
+ })
+ })
+
+ describe('when connected but expired', function () {
+ beforeEach(function () {
+ mockAuthProvider.isConnected.returns(true)
+ mockAuthProvider.isConnectionValid.returns(false)
+ mockAuthProvider.activeConnection = mockConnection
+ })
+
+ it('should return expired tree item with reauthenticate command', function () {
+ const treeItem = authInfoNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Domain: dzd_domainId (Expired) - Click to reauthenticate')
+ assert.strictEqual(treeItem.description, 'us-east-2')
+ assert.strictEqual(treeItem.contextValue, 'smusAuthInfo')
+
+ // Check icon
+ assert.ok(treeItem.iconPath instanceof vscode.ThemeIcon)
+ assert.strictEqual((treeItem.iconPath as vscode.ThemeIcon).id, 'warning')
+
+ // Check tooltip
+ const tooltip = treeItem.tooltip as string
+ assert.ok(tooltip?.includes('Connection to SageMaker Unified Studio has expired'))
+ assert.ok(tooltip?.includes('Status: Expired - Click to reauthenticate'))
+
+ // Should have reauthenticate command
+ assert.ok(treeItem.command)
+ assert.strictEqual(treeItem.command.command, 'aws.smus.reauthenticate')
+ assert.strictEqual(treeItem.command.title, 'Reauthenticate')
+ assert.deepStrictEqual(treeItem.command.arguments, [mockConnection])
+ })
+ })
+
+ describe('when not connected', function () {
+ beforeEach(function () {
+ mockAuthProvider.isConnected.returns(false)
+ mockAuthProvider.isConnectionValid.returns(false)
+ mockAuthProvider.activeConnection = undefined
+ })
+
+ it('should return not connected tree item', function () {
+ const treeItem = authInfoNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Not Connected')
+ assert.strictEqual(treeItem.description, undefined)
+ assert.strictEqual(treeItem.contextValue, 'smusAuthInfo')
+
+ // Check icon
+ assert.ok(treeItem.iconPath instanceof vscode.ThemeIcon)
+ assert.strictEqual((treeItem.iconPath as vscode.ThemeIcon).id, 'circle-slash')
+
+ // Check tooltip
+ const tooltip = treeItem.tooltip as string
+ assert.ok(tooltip?.includes('Not connected to SageMaker Unified Studio'))
+ assert.ok(tooltip?.includes('Please sign in to access your projects'))
+
+ // Should not have command when not connected
+ assert.strictEqual(treeItem.command, undefined)
+ })
+ })
+
+ describe('with missing connection details', function () {
+ beforeEach(function () {
+ const incompleteConnection = {
+ ...mockConnection,
+ domainId: undefined,
+ ssoRegion: undefined,
+ } as any
+
+ mockAuthProvider.isConnected.returns(true)
+ mockAuthProvider.isConnectionValid.returns(true)
+ mockAuthProvider.activeConnection = incompleteConnection
+ })
+
+ it('should handle missing domain ID and region gracefully', function () {
+ const treeItem = authInfoNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Domain: Unknown')
+ assert.strictEqual(treeItem.description, 'Unknown')
+
+ const tooltip = treeItem.tooltip as string
+ assert.ok(tooltip?.includes('Domain ID: Unknown'))
+ assert.ok(tooltip?.includes('Region: Unknown'))
+ })
+ })
+ })
+
+ describe('getParent', function () {
+ it('should return undefined', function () {
+ assert.strictEqual(authInfoNode.getParent(), undefined)
+ })
+ })
+
+ describe('event handling', function () {
+ it('should fire onDidChangeTreeItem when auth provider changes', function () {
+ const eventSpy = sinon.spy()
+ authInfoNode.onDidChangeTreeItem(eventSpy)
+
+ // Simulate auth provider change
+ const onDidChangeCallback = mockAuthProvider.onDidChange.firstCall.args[0]
+ onDidChangeCallback()
+
+ assert.ok(eventSpy.called)
+ })
+
+ it('should dispose event listener properly', function () {
+ const disposeSpy = sinon.spy()
+ mockAuthProvider.onDidChange.returns({ dispose: disposeSpy })
+
+ // Create new node to trigger event listener registration
+ new SageMakerUnifiedStudioAuthInfoNode()
+
+ // The dispose should be available for cleanup
+ assert.ok(mockAuthProvider.onDidChange.called)
+ })
+ })
+
+ describe('theme icon colors', function () {
+ it('should use green color for connected state', function () {
+ mockAuthProvider.isConnected.returns(true)
+ mockAuthProvider.isConnectionValid.returns(true)
+
+ const treeItem = authInfoNode.getTreeItem()
+ const icon = treeItem.iconPath as vscode.ThemeIcon
+
+ assert.ok(icon.color instanceof vscode.ThemeColor)
+ assert.strictEqual((icon.color as any).id, 'charts.green')
+ })
+
+ it('should use yellow color for expired state', function () {
+ mockAuthProvider.isConnected.returns(true)
+ mockAuthProvider.isConnectionValid.returns(false)
+
+ const treeItem = authInfoNode.getTreeItem()
+ const icon = treeItem.iconPath as vscode.ThemeIcon
+
+ assert.ok(icon.color instanceof vscode.ThemeColor)
+ assert.strictEqual((icon.color as any).id, 'charts.yellow')
+ })
+
+ it('should use red color for not connected state', function () {
+ mockAuthProvider.isConnected.returns(false)
+
+ const treeItem = authInfoNode.getTreeItem()
+ const icon = treeItem.iconPath as vscode.ThemeIcon
+
+ assert.ok(icon.color instanceof vscode.ThemeColor)
+ assert.strictEqual((icon.color as any).id, 'charts.red')
+ })
+ })
+
+ describe('tooltip content', function () {
+ it('should include all relevant information for connected state', function () {
+ mockAuthProvider.isConnected.returns(true)
+ mockAuthProvider.isConnectionValid.returns(true)
+
+ const treeItem = authInfoNode.getTreeItem()
+ const tooltip = treeItem.tooltip as string
+
+ assert.ok(tooltip.includes('Connected to SageMaker Unified Studio'))
+ assert.ok(tooltip.includes(`Domain ID: ${mockConnection.domainId}`))
+ assert.ok(tooltip.includes(`Region: ${mockConnection.ssoRegion}`))
+ assert.ok(tooltip.includes('Status: Connected'))
+ })
+
+ it('should include expiration information for expired state', function () {
+ mockAuthProvider.isConnected.returns(true)
+ mockAuthProvider.isConnectionValid.returns(false)
+
+ const treeItem = authInfoNode.getTreeItem()
+ const tooltip = treeItem.tooltip as string
+
+ assert.ok(tooltip.includes('Connection to SageMaker Unified Studio has expired'))
+ assert.ok(tooltip.includes('Status: Expired - Click to reauthenticate'))
+ })
+
+ it('should include sign-in prompt for not connected state', function () {
+ mockAuthProvider.isConnected.returns(false)
+
+ const treeItem = authInfoNode.getTreeItem()
+ const tooltip = treeItem.tooltip as string
+
+ assert.ok(tooltip.includes('Not connected to SageMaker Unified Studio'))
+ assert.ok(tooltip.includes('Please sign in to access your projects'))
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioComputeNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioComputeNode.test.ts
new file mode 100644
index 00000000000..fc74eeab435
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioComputeNode.test.ts
@@ -0,0 +1,93 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import { SageMakerUnifiedStudioComputeNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioComputeNode'
+import { SageMakerUnifiedStudioProjectNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode'
+import { SageMakerUnifiedStudioSpacesParentNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode'
+import { SagemakerClient } from '../../../../shared/clients/sagemaker'
+import { SmusAuthenticationProvider } from '../../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+
+describe('SageMakerUnifiedStudioComputeNode', function () {
+ let computeNode: SageMakerUnifiedStudioComputeNode
+ let mockParent: SageMakerUnifiedStudioProjectNode
+ let mockExtensionContext: vscode.ExtensionContext
+ let mockAuthProvider: SmusAuthenticationProvider
+ let mockSagemakerClient: SagemakerClient
+
+ beforeEach(function () {
+ mockParent = {
+ getProject: sinon.stub(),
+ } as any
+
+ mockExtensionContext = {
+ subscriptions: [],
+ extensionUri: vscode.Uri.file('/test'),
+ } as any
+
+ mockAuthProvider = {} as any
+ mockSagemakerClient = {} as any
+
+ computeNode = new SageMakerUnifiedStudioComputeNode(
+ mockParent,
+ mockExtensionContext,
+ mockAuthProvider,
+ mockSagemakerClient
+ )
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('creates instance with correct properties', function () {
+ assert.strictEqual(computeNode.id, 'smusComputeNode')
+ assert.strictEqual(computeNode.resource, computeNode)
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('returns correct tree item', async function () {
+ const treeItem = await computeNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Compute')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Expanded)
+ assert.strictEqual(treeItem.contextValue, 'smusComputeNode')
+ assert.ok(treeItem.iconPath)
+ })
+ })
+
+ describe('getChildren', function () {
+ it('returns empty array when no project is selected', async function () {
+ ;(mockParent.getProject as sinon.SinonStub).returns(undefined)
+
+ const children = await computeNode.getChildren()
+
+ assert.deepStrictEqual(children, [])
+ })
+
+ it('returns connection nodes and spaces node when project is selected', async function () {
+ const mockProject = { id: 'project-123', name: 'Test Project' }
+ ;(mockParent.getProject as sinon.SinonStub).returns(mockProject)
+
+ const children = await computeNode.getChildren()
+
+ assert.strictEqual(children.length, 3)
+ assert.strictEqual(children[0].id, 'Data warehouse')
+ assert.strictEqual(children[1].id, 'Data processing')
+ assert.ok(children[2] instanceof SageMakerUnifiedStudioSpacesParentNode)
+ })
+ })
+
+ describe('getParent', function () {
+ it('returns parent node', function () {
+ const parent = computeNode.getParent()
+ assert.strictEqual(parent, mockParent)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionNode.test.ts
new file mode 100644
index 00000000000..a85d63302a6
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionNode.test.ts
@@ -0,0 +1,144 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import { SageMakerUnifiedStudioConnectionNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionNode'
+import { SageMakerUnifiedStudioConnectionParentNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionParentNode'
+import { ConnectionType, ConnectionSummary } from '@aws-sdk/client-datazone'
+import { getLogger } from '../../../../shared/logger/logger'
+
+describe('SageMakerUnifiedStudioConnectionNode', function () {
+ let connectionNode: SageMakerUnifiedStudioConnectionNode
+ let mockParent: sinon.SinonStubbedInstance
+
+ const mockRedshiftConnection: ConnectionSummary = {
+ connectionId: 'conn-1',
+ name: 'Test Redshift Connection',
+ type: ConnectionType.REDSHIFT,
+ environmentId: 'env-1',
+ domainId: 'domain-1',
+ domainUnitId: 'unit-1',
+ physicalEndpoints: [],
+ props: {
+ redshiftProperties: {
+ jdbcUrl: 'jdbc:redshift://test-cluster:5439/testdb',
+ },
+ },
+ }
+
+ const mockSparkConnection: ConnectionSummary = {
+ connectionId: 'conn-2',
+ name: 'Test Spark Connection',
+ type: ConnectionType.SPARK,
+ environmentId: 'env-2',
+ domainId: 'domain-2',
+ domainUnitId: 'unit-2',
+ physicalEndpoints: [],
+ props: {
+ sparkGlueProperties: {
+ glueVersion: '4.0',
+ workerType: 'G.1X',
+ numberOfWorkers: 2,
+ idleTimeout: 30,
+ },
+ },
+ }
+
+ beforeEach(function () {
+ mockParent = {} as any
+ sinon.stub(getLogger(), 'debug')
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('creates instance with correct properties for Redshift connection', function () {
+ connectionNode = new SageMakerUnifiedStudioConnectionNode(mockParent as any, mockRedshiftConnection)
+
+ assert.strictEqual(connectionNode.id, 'Test Redshift Connection')
+ assert.strictEqual(connectionNode.resource, connectionNode)
+ assert.strictEqual(connectionNode.contextValue, 'SageMakerUnifiedStudioConnectionNode')
+ })
+
+ it('creates instance with empty id when connection name is undefined', function () {
+ const connectionWithoutName = { ...mockRedshiftConnection, name: undefined }
+ connectionNode = new SageMakerUnifiedStudioConnectionNode(mockParent as any, connectionWithoutName)
+
+ assert.strictEqual(connectionNode.id, '')
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('returns correct tree item for Redshift connection', async function () {
+ connectionNode = new SageMakerUnifiedStudioConnectionNode(mockParent as any, mockRedshiftConnection)
+
+ const treeItem = await connectionNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Test Redshift Connection')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ assert.strictEqual(treeItem.contextValue, 'SageMakerUnifiedStudioConnectionNode')
+ assert.ok(treeItem.tooltip instanceof vscode.MarkdownString)
+ })
+
+ it('returns correct tree item for Spark connection', async function () {
+ connectionNode = new SageMakerUnifiedStudioConnectionNode(mockParent as any, mockSparkConnection)
+
+ const treeItem = await connectionNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Test Spark Connection')
+ assert.ok(treeItem.tooltip instanceof vscode.MarkdownString)
+ })
+ })
+
+ describe('tooltip generation', function () {
+ it('generates correct tooltip for Redshift connection', async function () {
+ connectionNode = new SageMakerUnifiedStudioConnectionNode(mockParent as any, mockRedshiftConnection)
+
+ const treeItem = await connectionNode.getTreeItem()
+ const tooltip = (treeItem.tooltip as vscode.MarkdownString).value
+
+ assert(tooltip.includes('REDSHIFT'))
+ assert(tooltip.includes('env-1'))
+ assert(tooltip.includes('jdbc:redshift://test-cluster:5439/testdb'))
+ })
+
+ it('generates correct tooltip for Spark connection', async function () {
+ connectionNode = new SageMakerUnifiedStudioConnectionNode(mockParent as any, mockSparkConnection)
+
+ const treeItem = await connectionNode.getTreeItem()
+ const tooltip = (treeItem.tooltip as vscode.MarkdownString).value
+
+ assert(tooltip.includes('SPARK'))
+ assert(tooltip.includes('4.0'))
+ assert(tooltip.includes('G.1X'))
+ assert(tooltip.includes('2'))
+ assert(tooltip.includes('30'))
+ })
+
+ it('generates empty tooltip for unknown connection type', async function () {
+ const unknownConnection = { ...mockRedshiftConnection, type: 'UNKNOWN' as ConnectionType }
+ connectionNode = new SageMakerUnifiedStudioConnectionNode(mockParent as any, unknownConnection)
+
+ const treeItem = await connectionNode.getTreeItem()
+ const tooltip = (treeItem.tooltip as vscode.MarkdownString).value
+
+ assert.strictEqual(tooltip, '')
+ })
+ })
+
+ describe('getParent', function () {
+ it('returns the parent node', function () {
+ connectionNode = new SageMakerUnifiedStudioConnectionNode(mockParent as any, mockRedshiftConnection)
+
+ const parent = connectionNode.getParent()
+
+ assert.strictEqual(parent, mockParent)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionParentNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionParentNode.test.ts
new file mode 100644
index 00000000000..686c85a0055
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionParentNode.test.ts
@@ -0,0 +1,234 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import { SageMakerUnifiedStudioConnectionParentNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionParentNode'
+import { SageMakerUnifiedStudioComputeNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioComputeNode'
+import { SageMakerUnifiedStudioConnectionNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioConnectionNode'
+import { DataZoneClient } from '../../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+
+import { ConnectionType, ListConnectionsCommandOutput, ConnectionSummary } from '@aws-sdk/client-datazone'
+import { getLogger } from '../../../../shared/logger/logger'
+
+describe('SageMakerUnifiedStudioConnectionParentNode', function () {
+ let connectionParentNode: SageMakerUnifiedStudioConnectionParentNode
+ let mockComputeNode: sinon.SinonStubbedInstance
+
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+
+ const mockProject = {
+ id: 'project-123',
+ domainId: 'domain-123',
+ }
+
+ const mockConnectionsOutput: ListConnectionsCommandOutput = {
+ items: [
+ {
+ connectionId: 'conn-1',
+ name: 'Test Connection 1',
+ type: ConnectionType.REDSHIFT,
+ environmentId: 'env-1',
+ } as ConnectionSummary,
+ {
+ connectionId: 'conn-2',
+ name: 'Test Connection 2',
+ type: ConnectionType.REDSHIFT,
+ environmentId: 'env-2',
+ } as ConnectionSummary,
+ ],
+ $metadata: {},
+ }
+
+ beforeEach(function () {
+ // Create mock objects
+ mockDataZoneClient = {
+ fetchConnections: sinon.stub(),
+ } as any
+
+ mockComputeNode = {
+ authProvider: {} as any,
+ parent: {
+ project: mockProject,
+ } as any,
+ } as any
+
+ // Stub static methods
+ sinon.stub(DataZoneClient, 'getInstance').resolves(mockDataZoneClient as any)
+ sinon.stub(getLogger(), 'debug')
+
+ connectionParentNode = new SageMakerUnifiedStudioConnectionParentNode(
+ mockComputeNode as any,
+ ConnectionType.REDSHIFT,
+ 'Data warehouse'
+ )
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('creates instance with correct properties', function () {
+ assert.strictEqual(connectionParentNode.id, 'Data warehouse')
+ assert.strictEqual(connectionParentNode.resource, connectionParentNode)
+ assert.strictEqual(connectionParentNode.contextValue, 'SageMakerUnifiedStudioConnectionParentNode')
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('returns correct tree item', async function () {
+ const treeItem = await connectionParentNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Data warehouse')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Collapsed)
+ assert.strictEqual(treeItem.contextValue, 'SageMakerUnifiedStudioConnectionParentNode')
+ })
+ })
+
+ describe('getChildren', function () {
+ it('returns connection nodes when connections exist', async function () {
+ mockDataZoneClient.fetchConnections.resolves(mockConnectionsOutput)
+
+ const children = await connectionParentNode.getChildren()
+
+ assert.strictEqual(children.length, 2)
+ assert(children[0] instanceof SageMakerUnifiedStudioConnectionNode)
+ assert(children[1] instanceof SageMakerUnifiedStudioConnectionNode)
+
+ // Verify fetchConnections was called with correct parameters
+ assert(
+ mockDataZoneClient.fetchConnections.calledOnceWith(
+ mockProject.domainId,
+ mockProject.id,
+ ConnectionType.REDSHIFT
+ )
+ )
+ })
+
+ it('returns no connections node when no connections exist', async function () {
+ const emptyOutput: ListConnectionsCommandOutput = { items: [], $metadata: {} }
+ mockDataZoneClient.fetchConnections.resolves(emptyOutput)
+
+ const children = await connectionParentNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual(children[0].id, 'smusNoConnections')
+ const treeItem = await children[0].getTreeItem()
+ assert.strictEqual(treeItem.label, '[No connections found]')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ })
+
+ it('returns no connections node when connections items is undefined', async function () {
+ const undefinedOutput: ListConnectionsCommandOutput = { items: undefined, $metadata: {} }
+ mockDataZoneClient.fetchConnections.resolves(undefinedOutput)
+
+ const children = await connectionParentNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual(children[0].id, 'smusNoConnections')
+ })
+
+ it('handles missing project information gracefully', async function () {
+ const nodeWithoutProject = new SageMakerUnifiedStudioConnectionParentNode(
+ {
+ authProvider: {} as any,
+ parent: {
+ project: undefined,
+ } as any,
+ } as any,
+ ConnectionType.SPARK,
+ 'Data processing'
+ )
+
+ mockDataZoneClient.fetchConnections.resolves({ items: [], $metadata: {} })
+
+ const children = await nodeWithoutProject.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual(children[0].id, 'smusNoConnections')
+ assert(mockDataZoneClient.fetchConnections.calledOnceWith(undefined, undefined, ConnectionType.SPARK))
+ })
+ })
+
+ describe('getParent', function () {
+ it('returns the parent compute node', function () {
+ const parent = connectionParentNode.getParent()
+ assert.strictEqual(parent, mockComputeNode)
+ })
+ })
+
+ describe('error handling', function () {
+ it('handles DataZoneClient.getInstance error', async function () {
+ sinon.restore()
+ sinon.stub(DataZoneClient, 'getInstance').rejects(new Error('Client error'))
+ sinon.stub(getLogger(), 'debug')
+
+ try {
+ await connectionParentNode.getChildren()
+ assert.fail('Expected error to be thrown')
+ } catch (error) {
+ assert.strictEqual((error as Error).message, 'Client error')
+ }
+ })
+
+ it('handles fetchConnections error', async function () {
+ mockDataZoneClient.fetchConnections.rejects(new Error('Fetch error'))
+
+ try {
+ await connectionParentNode.getChildren()
+ assert.fail('Expected error to be thrown')
+ } catch (error) {
+ assert.strictEqual((error as Error).message, 'Fetch error')
+ }
+ })
+ })
+
+ describe('connections property', function () {
+ it('sets connections property after getChildren call', async function () {
+ mockDataZoneClient.fetchConnections.resolves(mockConnectionsOutput)
+
+ await connectionParentNode.getChildren()
+
+ assert.strictEqual(connectionParentNode.connections, mockConnectionsOutput)
+ })
+ })
+
+ describe('different connection types', function () {
+ it('works with SPARK connection type', async function () {
+ const sparkNode = new SageMakerUnifiedStudioConnectionParentNode(
+ mockComputeNode as any,
+ ConnectionType.SPARK,
+ 'Spark connections'
+ )
+
+ const sparkOutput = {
+ items: [
+ {
+ connectionId: 'spark-1',
+ name: 'Spark Connection',
+ type: ConnectionType.SPARK,
+ environmentId: 'env-spark',
+ } as ConnectionSummary,
+ ],
+ $metadata: {},
+ }
+
+ mockDataZoneClient.fetchConnections.resolves(sparkOutput)
+
+ const children = await sparkNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert(
+ mockDataZoneClient.fetchConnections.calledWith(
+ mockProject.domainId,
+ mockProject.id,
+ ConnectionType.SPARK
+ )
+ )
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioDataNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioDataNode.test.ts
new file mode 100644
index 00000000000..991e5955989
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioDataNode.test.ts
@@ -0,0 +1,235 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import * as vscode from 'vscode'
+import { SageMakerUnifiedStudioDataNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioDataNode'
+import { SageMakerUnifiedStudioProjectNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode'
+import { DataZoneClient, DataZoneProject } from '../../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { SmusAuthenticationProvider } from '../../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+import * as s3Strategy from '../../../../sagemakerunifiedstudio/explorer/nodes/s3Strategy'
+import * as redshiftStrategy from '../../../../sagemakerunifiedstudio/explorer/nodes/redshiftStrategy'
+import * as lakehouseStrategy from '../../../../sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy'
+
+describe('SageMakerUnifiedStudioDataNode', function () {
+ let sandbox: sinon.SinonSandbox
+ let dataNode: SageMakerUnifiedStudioDataNode
+ let mockParent: sinon.SinonStubbedInstance
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+ let mockAuthProvider: sinon.SinonStubbedInstance
+ let mockProjectCredentialProvider: any
+
+ const mockProject: DataZoneProject = {
+ id: 'project-123',
+ name: 'Test Project',
+ domainId: 'domain-123',
+ }
+
+ const mockCredentials = {
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ $metadata: {},
+ }
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+
+ mockParent = {
+ getProject: sandbox.stub().returns(mockProject),
+ } as any
+
+ mockProjectCredentialProvider = {
+ getCredentials: sandbox.stub().resolves(mockCredentials),
+ }
+
+ mockAuthProvider = {
+ getProjectCredentialProvider: sandbox.stub().resolves(mockProjectCredentialProvider),
+ getConnectionCredentialsProvider: sandbox.stub().resolves(mockProjectCredentialProvider),
+ getDomainRegion: sandbox.stub().returns('us-east-1'),
+ } as any
+
+ mockDataZoneClient = {
+ getInstance: sandbox.stub(),
+ getProjectDefaultEnvironmentCreds: sandbox.stub(),
+ listConnections: sandbox.stub(),
+ getConnection: sandbox.stub(),
+ getRegion: sandbox.stub().returns('us-east-1'),
+ } as any
+
+ sandbox.stub(DataZoneClient, 'getInstance').returns(mockDataZoneClient as any)
+ sandbox.stub(SmusAuthenticationProvider, 'fromContext').returns(mockAuthProvider as any)
+ sandbox.stub(s3Strategy, 'createS3ConnectionNode').returns({
+ id: 's3-node',
+ getChildren: () => Promise.resolve([]),
+ getTreeItem: () => ({}) as any,
+ getParent: () => undefined,
+ } as any)
+ sandbox.stub(s3Strategy, 'createS3AccessGrantNodes').resolves([])
+ sandbox.stub(redshiftStrategy, 'createRedshiftConnectionNode').returns({
+ id: 'redshift-node',
+ getChildren: () => Promise.resolve([]),
+ getTreeItem: () => ({}) as any,
+ getParent: () => undefined,
+ } as any)
+ sandbox.stub(lakehouseStrategy, 'createLakehouseConnectionNode').returns({
+ id: 'lakehouse-node',
+ getChildren: () => Promise.resolve([]),
+ getTreeItem: () => ({}) as any,
+ getParent: () => undefined,
+ } as any)
+
+ dataNode = new SageMakerUnifiedStudioDataNode(mockParent as any)
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ describe('constructor', function () {
+ it('should initialize with correct properties', function () {
+ assert.strictEqual(dataNode.id, 'smusDataExplorer')
+ assert.deepStrictEqual(dataNode.resource, {})
+ })
+
+ it('should initialize with provided children', function () {
+ const initialChildren = [{ id: 'child1' } as any]
+ const nodeWithChildren = new SageMakerUnifiedStudioDataNode(mockParent as any, initialChildren)
+ // Children should be cached
+ assert.ok(nodeWithChildren)
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('should return correct tree item', function () {
+ const treeItem = dataNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Data')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Collapsed)
+ assert.strictEqual(treeItem.contextValue, 'dataFolder')
+ })
+ })
+
+ describe('getParent', function () {
+ it('should return parent node', function () {
+ assert.strictEqual(dataNode.getParent(), mockParent)
+ })
+ })
+
+ describe('getChildren', function () {
+ it('should return cached children if available', async function () {
+ const initialChildren = [{ id: 'cached' } as any]
+ const nodeWithCache = new SageMakerUnifiedStudioDataNode(mockParent as any, initialChildren)
+
+ const children = await nodeWithCache.getChildren()
+ assert.strictEqual(children, initialChildren)
+ })
+
+ it('should return error node when no project available', async function () {
+ mockParent.getProject.returns(undefined)
+
+ const children = await dataNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.ok(children[0].id.startsWith('smusDataExplorer-error-project-'))
+ })
+
+ it('should return error node when credentials are missing', async function () {
+ mockProjectCredentialProvider.getCredentials.resolves(undefined)
+
+ const children = await dataNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.ok(children[0].id.startsWith('smusDataExplorer-error-connections-'))
+ })
+
+ it('should return placeholder when no connections found', async function () {
+ mockDataZoneClient.listConnections.resolves([])
+
+ const children = await dataNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual(children[0].resource, '[No data found]')
+ })
+
+ it('should create Bucket parent node and Redshift nodes for connections', async function () {
+ const mockConnections = [
+ { connectionId: 's3-conn', type: 'S3', name: 's3-connection' },
+ { connectionId: 'redshift-conn', type: 'REDSHIFT', name: 'redshift-connection' },
+ ]
+
+ mockDataZoneClient.listConnections.resolves(mockConnections as any)
+ mockDataZoneClient.getConnection
+ .onFirstCall()
+ .resolves({
+ location: { awsRegion: 'us-east-1', awsAccountId: '' },
+ connectionCredentials: mockCredentials,
+ connectionId: '',
+ name: '',
+ type: '',
+ domainId: '',
+ projectId: '',
+ })
+ .onSecondCall()
+ .resolves({
+ location: { awsRegion: 'us-east-1', awsAccountId: '' },
+ connectionCredentials: mockCredentials,
+ connectionId: '',
+ name: '',
+ type: '',
+ domainId: '',
+ projectId: '',
+ })
+
+ const children = await dataNode.getChildren()
+
+ // Should have Bucket parent node and Redshift node
+ assert.strictEqual(children.length, 2)
+
+ // Check for Bucket parent node
+ const bucketNode = children.find((child) => child.id === 'bucket-parent')
+ assert.ok(bucketNode, 'Should have bucket parent node')
+
+ // Verify Bucket node has correct tree item
+ const bucketTreeItem = await bucketNode!.getTreeItem()
+ assert.strictEqual(bucketTreeItem.label, 'Buckets')
+ assert.strictEqual(bucketTreeItem.contextValue, 'bucketFolder')
+
+ // Verify S3 nodes are created when Bucket node is expanded
+ await bucketNode!.getChildren!()
+ assert.ok((s3Strategy.createS3ConnectionNode as sinon.SinonStub).calledOnce)
+
+ assert.ok((redshiftStrategy.createRedshiftConnectionNode as sinon.SinonStub).calledOnce)
+ })
+
+ it('should handle connection detail errors gracefully', async function () {
+ const mockConnections = [{ connectionId: 's3-conn', type: 'S3', name: 's3-connection' }]
+
+ mockDataZoneClient.listConnections.resolves(mockConnections as any)
+ mockDataZoneClient.getConnection.rejects(new Error('Connection error'))
+
+ const children = await dataNode.getChildren()
+
+ // Should have Bucket parent node even with connection errors
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual(children[0].id, 'bucket-parent')
+
+ // Error should occur when expanding the Bucket node
+ const bucketChildren = await children[0].getChildren!()
+ assert.strictEqual(bucketChildren.length, 1)
+ assert.ok(bucketChildren[0].id.startsWith('smusDataExplorer-error-s3-'))
+ })
+
+ it('should return error node when general error occurs', async function () {
+ mockAuthProvider.getProjectCredentialProvider.rejects(new Error('General error'))
+
+ const children = await dataNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.ok(children[0].id.startsWith('smusDataExplorer-error-connections-'))
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode.test.ts
new file mode 100644
index 00000000000..02a0c1078d9
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode.test.ts
@@ -0,0 +1,343 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import { SageMakerUnifiedStudioProjectNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode'
+import { DataZoneClient, DataZoneProject } from '../../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { getLogger } from '../../../../shared/logger/logger'
+import { telemetry } from '../../../../shared/telemetry/telemetry'
+import { SagemakerClient } from '../../../../shared/clients/sagemaker'
+import { SageMakerUnifiedStudioDataNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioDataNode'
+import { SageMakerUnifiedStudioComputeNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioComputeNode'
+import * as vscodeUtils from '../../../../shared/vscode/setContext'
+
+describe('SageMakerUnifiedStudioProjectNode', function () {
+ let projectNode: SageMakerUnifiedStudioProjectNode
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+
+ const mockProject: DataZoneProject = {
+ id: 'project-123',
+ name: 'Test Project',
+ description: 'Test Description',
+ domainId: 'domain-123',
+ }
+
+ beforeEach(function () {
+ // Create mock parent
+ const mockParent = {} as any
+
+ // Create mock auth provider
+ const mockAuthProvider = {
+ activeConnection: { domainId: 'test-domain', ssoRegion: 'us-west-2' },
+ invalidateAllProjectCredentialsInCache: sinon.stub(),
+ getProjectCredentialProvider: sinon.stub(),
+ getDomainRegion: sinon.stub().returns('us-west-2'),
+ } as any
+
+ // Create mock extension context
+ const mockExtensionContext = {
+ subscriptions: [],
+ workspaceState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ globalState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ } as any
+
+ projectNode = new SageMakerUnifiedStudioProjectNode(mockParent, mockAuthProvider, mockExtensionContext)
+
+ sinon.stub(getLogger(), 'info')
+ sinon.stub(getLogger(), 'warn')
+
+ // Stub telemetry
+ sinon.stub(telemetry, 'record')
+
+ // Create mock DataZone client
+ mockDataZoneClient = {
+ getProjectDefaultEnvironmentCreds: sinon.stub(),
+ getUserId: sinon.stub(),
+ fetchAllProjectMemberships: sinon.stub(),
+ getDomainId: sinon.stub().returns('test-domain-id'),
+ getToolingEnvironmentId: sinon.stub(),
+ getEnvironmentDetails: sinon.stub(),
+ getToolingEnvironment: sinon.stub(),
+ } as any
+
+ // Stub DataZoneClient static methods
+ sinon.stub(DataZoneClient, 'getInstance').returns(mockDataZoneClient as any)
+
+ // Stub SagemakerClient constructor
+ sinon.stub(SagemakerClient.prototype, 'dispose')
+
+ // Stub child node constructors to prevent actual instantiation
+ sinon.stub(SageMakerUnifiedStudioDataNode.prototype, 'constructor' as any).returns({})
+ sinon.stub(SageMakerUnifiedStudioComputeNode.prototype, 'constructor' as any).returns({})
+
+ // Stub getContext to return false for SMUS space environment
+ sinon.stub(vscodeUtils, 'getContext').returns(false)
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('creates instance with correct properties', function () {
+ assert.strictEqual(projectNode.id, 'smusProjectNode')
+ assert.strictEqual(projectNode.resource, projectNode)
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('returns correct tree item when no project is selected', async function () {
+ const treeItem = await projectNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Select a project')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Expanded)
+ assert.strictEqual(treeItem.contextValue, 'smusProjectSelectPicker')
+ assert.ok(treeItem.command)
+ assert.strictEqual(treeItem.command?.command, 'aws.smus.projectView')
+ })
+
+ it('returns correct tree item when project is selected', async function () {
+ await projectNode.setProject(mockProject)
+ const treeItem = await projectNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Project: ' + mockProject.name)
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Expanded)
+ assert.strictEqual(treeItem.contextValue, 'smusSelectedProject')
+ assert.strictEqual(treeItem.tooltip, `Project: ${mockProject.name}\nID: ${mockProject.id}`)
+ })
+ })
+
+ describe('getParent', function () {
+ it('returns parent node', function () {
+ const parent = projectNode.getParent()
+ assert.ok(parent)
+ })
+ })
+
+ describe('setProject', function () {
+ it('updates the project and calls cleanupProjectResources', async function () {
+ const cleanupSpy = sinon.spy(projectNode as any, 'cleanupProjectResources')
+ await projectNode.setProject(mockProject)
+ assert.strictEqual(projectNode['project'], mockProject)
+ assert(cleanupSpy.calledOnce)
+ })
+ })
+
+ describe('clearProject', function () {
+ it('clears the project, calls cleanupProjectResources and fires change event', async function () {
+ await projectNode.setProject(mockProject)
+ const cleanupSpy = sinon.spy(projectNode as any, 'cleanupProjectResources')
+ const emitterSpy = sinon.spy(projectNode['onDidChangeEmitter'], 'fire')
+
+ await projectNode.clearProject()
+
+ assert.strictEqual(projectNode['project'], undefined)
+ assert(cleanupSpy.calledOnce)
+ assert(emitterSpy.calledOnce)
+ })
+ })
+
+ describe('getProject', function () {
+ it('returns undefined when no project is set', function () {
+ assert.strictEqual(projectNode.getProject(), undefined)
+ })
+
+ it('returns project when set', async function () {
+ await projectNode.setProject(mockProject)
+ assert.strictEqual(projectNode.getProject(), mockProject)
+ })
+ })
+
+ describe('refreshNode', function () {
+ it('fires change event', async function () {
+ const emitterSpy = sinon.spy(projectNode['onDidChangeEmitter'], 'fire')
+ await projectNode.refreshNode()
+ assert(emitterSpy.calledOnce)
+ })
+ })
+
+ describe('getChildren', function () {
+ it('returns empty array when no project is selected', async function () {
+ const children = await projectNode.getChildren()
+ assert.deepStrictEqual(children, [])
+ })
+
+ it('returns data and compute nodes when project is selected and user has access', async function () {
+ await projectNode.setProject(mockProject)
+ const mockCredProvider = {
+ getCredentials: sinon.stub().resolves({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+ projectNode['authProvider'].getProjectCredentialProvider = sinon.stub().resolves(mockCredProvider)
+
+ // Mock getToolingEnvironment method
+ mockDataZoneClient.getToolingEnvironment.resolves({
+ id: 'env-123',
+ awsAccountRegion: 'us-east-1',
+ projectId: undefined,
+ domainId: undefined,
+ createdBy: undefined,
+ name: undefined,
+ provider: undefined,
+ $metadata: {},
+ })
+
+ const children = await projectNode.getChildren()
+ assert.strictEqual(children.length, 2)
+ })
+
+ it('returns access denied message when user does not have project access', async function () {
+ await projectNode.setProject(mockProject)
+
+ // Mock access check to return false by throwing AccessDeniedException
+ const accessError = new Error('Access denied')
+ accessError.name = 'AccessDeniedException'
+ projectNode['authProvider'].getProjectCredentialProvider = sinon.stub().rejects(accessError)
+
+ const children = await projectNode.getChildren()
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual(children[0].id, 'smusProjectAccessDenied')
+
+ const treeItem = await children[0].getTreeItem()
+ assert.strictEqual(treeItem.label, 'You do not have access to this project. Contact your administrator.')
+ })
+
+ it('throws error when initializeSagemakerClient fails', async function () {
+ await projectNode.setProject(mockProject)
+ const credError = new Error('Failed to initialize SageMaker client')
+
+ // First call succeeds for access check, second call fails for initializeSagemakerClient
+ const mockCredProvider = {
+ getCredentials: sinon.stub().resolves({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+ projectNode['authProvider'].getProjectCredentialProvider = sinon
+ .stub()
+ .onFirstCall()
+ .resolves(mockCredProvider)
+ .onSecondCall()
+ .rejects(credError)
+
+ // Mock getToolingEnvironment method
+ mockDataZoneClient.getToolingEnvironment.resolves({
+ id: 'env-123',
+ awsAccountRegion: 'us-east-1',
+ projectId: undefined,
+ domainId: undefined,
+ createdBy: undefined,
+ name: undefined,
+ provider: undefined,
+ $metadata: {},
+ })
+
+ await assert.rejects(async () => await projectNode.getChildren(), credError)
+ })
+ })
+
+ describe('initializeSagemakerClient', function () {
+ it('throws error when no project is selected', async function () {
+ await assert.rejects(
+ async () => await projectNode['initializeSagemakerClient']('us-east-1'),
+ /No project selected for initializing SageMaker client/
+ )
+ })
+
+ it('creates SagemakerClient with project credentials', async function () {
+ await projectNode.setProject(mockProject)
+ const mockCredProvider = {
+ getCredentials: sinon.stub().resolves({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+ projectNode['authProvider'].getProjectCredentialProvider = sinon.stub().resolves(mockCredProvider)
+
+ const client = await projectNode['initializeSagemakerClient']('us-east-1')
+ assert.ok(client instanceof SagemakerClient)
+ assert(
+ (projectNode['authProvider'].getProjectCredentialProvider as sinon.SinonStub).calledWith(mockProject.id)
+ )
+ })
+ })
+
+ describe('checkProjectCredsAccess', function () {
+ it('returns true when user has project access', async function () {
+ const mockCredProvider = {
+ getCredentials: sinon.stub().resolves({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+ projectNode['authProvider'].getProjectCredentialProvider = sinon.stub().resolves(mockCredProvider)
+
+ const hasAccess = await projectNode['checkProjectCredsAccess']('project-123')
+ assert.strictEqual(hasAccess, true)
+ })
+
+ it('returns false when user does not have project access', async function () {
+ const accessError = new Error('Access denied')
+ accessError.name = 'AccessDeniedException'
+ projectNode['authProvider'].getProjectCredentialProvider = sinon.stub().rejects(accessError)
+
+ const hasAccess = await projectNode['checkProjectCredsAccess']('project-123')
+ assert.strictEqual(hasAccess, false)
+ })
+
+ it('returns false when getCredentials fails', async function () {
+ const mockCredProvider = {
+ getCredentials: sinon.stub().rejects(new Error('Credentials error')),
+ }
+ projectNode['authProvider'].getProjectCredentialProvider = sinon.stub().resolves(mockCredProvider)
+
+ const hasAccess = await projectNode['checkProjectCredsAccess']('project-123')
+ assert.strictEqual(hasAccess, false)
+ })
+
+ it('returns false when access check throws non-AccessDeniedException error', async function () {
+ projectNode['authProvider'].getProjectCredentialProvider = sinon.stub().rejects(new Error('Other error'))
+
+ const hasAccess = await projectNode['checkProjectCredsAccess']('project-123')
+ assert.strictEqual(hasAccess, false)
+ })
+ })
+
+ describe('cleanupProjectResources', function () {
+ it('invalidates credentials and disposes existing sagemaker client', async function () {
+ // Set up existing sagemaker client with mock
+ const mockClient = { dispose: sinon.stub() } as any
+ projectNode['sagemakerClient'] = mockClient
+
+ await projectNode['cleanupProjectResources']()
+
+ assert((projectNode['authProvider'].invalidateAllProjectCredentialsInCache as sinon.SinonStub).calledOnce)
+ assert(mockClient.dispose.calledOnce)
+ assert.strictEqual(projectNode['sagemakerClient'], undefined)
+ })
+
+ it('handles case when no sagemaker client exists', async function () {
+ projectNode['sagemakerClient'] = undefined
+
+ await projectNode['cleanupProjectResources']()
+
+ assert((projectNode['authProvider'].invalidateAllProjectCredentialsInCache as sinon.SinonStub).calledOnce)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.test.ts
new file mode 100644
index 00000000000..f89413e9528
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode.test.ts
@@ -0,0 +1,603 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import {
+ SageMakerUnifiedStudioRootNode,
+ selectSMUSProject,
+} from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioRootNode'
+import { SageMakerUnifiedStudioProjectNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioProjectNode'
+import { DataZoneClient, DataZoneProject } from '../../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { SageMakerUnifiedStudioAuthInfoNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioAuthInfoNode'
+import { SmusAuthenticationProvider } from '../../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+import * as pickerPrompter from '../../../../shared/ui/pickerPrompter'
+import { getTestWindow } from '../../../shared/vscode/window'
+import { assertTelemetry } from '../../../../../src/test/testUtil'
+
+describe('SmusRootNode', function () {
+ let rootNode: SageMakerUnifiedStudioRootNode
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+
+ const testDomainId = 'test-domain-123'
+ const mockProject: DataZoneProject = {
+ id: 'project-123',
+ name: 'Test Project',
+ description: 'Test Description',
+ domainId: testDomainId,
+ }
+
+ beforeEach(function () {
+ // Create mock extension context
+ const mockExtensionContext = {
+ subscriptions: [],
+ workspaceState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ globalState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ } as any
+
+ // Create a mock auth provider
+ const mockAuthProvider = {
+ isConnected: sinon.stub().returns(true),
+ isConnectionValid: sinon.stub().returns(true),
+ activeConnection: { domainId: testDomainId, ssoRegion: 'us-west-2' },
+ onDidChange: sinon.stub().returns({ dispose: sinon.stub() }),
+ } as any
+
+ rootNode = new SageMakerUnifiedStudioRootNode(mockAuthProvider, mockExtensionContext)
+
+ // Mock domain ID is handled by the mock auth provider
+
+ // Create mock DataZone client
+ mockDataZoneClient = {
+ getDomainId: sinon.stub().returns(testDomainId),
+ listProjects: sinon.stub(),
+ } as any
+
+ // Stub DataZoneClient static methods
+ sinon.stub(DataZoneClient, 'getInstance').returns(mockDataZoneClient as any)
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('should initialize id and resource properties', function () {
+ // Create a mock auth provider
+ const mockAuthProvider = {
+ isConnected: sinon.stub().returns(true),
+ isConnectionValid: sinon.stub().returns(true),
+ activeConnection: { domainId: testDomainId, ssoRegion: 'us-west-2' },
+ onDidChange: sinon.stub().returns({ dispose: sinon.stub() }),
+ } as any
+
+ const mockExtensionContext = {
+ subscriptions: [],
+ workspaceState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ globalState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ } as any
+
+ const node = new SageMakerUnifiedStudioRootNode(mockAuthProvider, mockExtensionContext)
+ assert.strictEqual(node.id, 'smusRootNode')
+ assert.strictEqual(node.resource, node)
+ assert.ok(node.getAuthInfoNode() instanceof SageMakerUnifiedStudioAuthInfoNode)
+ assert.ok(node.getProjectSelectNode() instanceof SageMakerUnifiedStudioProjectNode)
+ assert.strictEqual(typeof node.onDidChangeTreeItem, 'function')
+ assert.strictEqual(typeof node.onDidChangeChildren, 'function')
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('returns correct tree item when authenticated', async function () {
+ const treeItem = rootNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'SageMaker Unified Studio')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Expanded)
+ assert.strictEqual(treeItem.contextValue, 'sageMakerUnifiedStudioRoot')
+ assert.strictEqual(treeItem.description, 'Connected')
+ assert.ok(treeItem.iconPath)
+ })
+
+ it('returns correct tree item when not authenticated', async function () {
+ // Create a mock auth provider for unauthenticated state
+ const mockAuthProvider = {
+ isConnected: sinon.stub().returns(false),
+ isConnectionValid: sinon.stub().returns(false),
+ activeConnection: undefined,
+ onDidChange: sinon.stub().returns({ dispose: sinon.stub() }),
+ } as any
+
+ const mockExtensionContext = {
+ subscriptions: [],
+ workspaceState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ globalState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ } as any
+
+ const unauthenticatedNode = new SageMakerUnifiedStudioRootNode(mockAuthProvider, mockExtensionContext)
+ const treeItem = unauthenticatedNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'SageMaker Unified Studio')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Expanded)
+ assert.strictEqual(treeItem.contextValue, 'sageMakerUnifiedStudioRoot')
+ assert.strictEqual(treeItem.description, 'Not authenticated')
+ assert.ok(treeItem.iconPath)
+ })
+ })
+
+ describe('getChildren', function () {
+ it('returns login node when not authenticated (empty domain ID)', async function () {
+ // Create a mock auth provider for unauthenticated state
+ const mockAuthProvider = {
+ isConnected: sinon.stub().returns(false),
+ isConnectionValid: sinon.stub().returns(false),
+ activeConnection: undefined,
+ onDidChange: sinon.stub().returns({ dispose: sinon.stub() }),
+ } as any
+
+ const mockExtensionContext = {
+ subscriptions: [],
+ workspaceState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ globalState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ } as any
+
+ const unauthenticatedNode = new SageMakerUnifiedStudioRootNode(mockAuthProvider, mockExtensionContext)
+ const children = await unauthenticatedNode.getChildren()
+
+ assert.strictEqual(children.length, 2)
+ assert.strictEqual(children[0].id, 'smusLogin')
+ assert.strictEqual(children[1].id, 'smusLearnMore')
+
+ // Check login node
+ const loginTreeItem = await children[0].getTreeItem()
+ assert.strictEqual(loginTreeItem.label, 'Sign in to get started')
+ assert.strictEqual(loginTreeItem.contextValue, 'sageMakerUnifiedStudioLogin')
+ assert.deepStrictEqual(loginTreeItem.command, {
+ command: 'aws.smus.login',
+ title: 'Sign in to SageMaker Unified Studio',
+ })
+
+ // Check learn more node
+ const learnMoreTreeItem = await children[1].getTreeItem()
+ assert.strictEqual(learnMoreTreeItem.label, 'Learn more about SageMaker Unified Studio')
+ assert.strictEqual(learnMoreTreeItem.contextValue, 'sageMakerUnifiedStudioLearnMore')
+ assert.deepStrictEqual(learnMoreTreeItem.command, {
+ command: 'aws.smus.learnMore',
+ title: 'Learn more about SageMaker Unified Studio',
+ })
+ })
+
+ it('returns login node when DataZone client throws error', async function () {
+ // Create a mock auth provider that throws an error
+ const mockAuthProvider = {
+ isConnected: sinon.stub().throws(new Error('Auth provider error')),
+ isConnectionValid: sinon.stub().returns(false),
+ activeConnection: undefined,
+ onDidChange: sinon.stub().returns({ dispose: sinon.stub() }),
+ } as any
+
+ const mockExtensionContext = {
+ subscriptions: [],
+ workspaceState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ globalState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ } as any
+
+ const errorNode = new SageMakerUnifiedStudioRootNode(mockAuthProvider, mockExtensionContext)
+ const children = await errorNode.getChildren()
+
+ assert.strictEqual(children.length, 2)
+ assert.strictEqual(children[0].id, 'smusLogin')
+ assert.strictEqual(children[1].id, 'smusLearnMore')
+
+ // Check login node
+ const loginTreeItem = await children[0].getTreeItem()
+ assert.strictEqual(loginTreeItem.label, 'Sign in to get started')
+ assert.strictEqual(loginTreeItem.contextValue, 'sageMakerUnifiedStudioLogin')
+
+ // Check learn more node
+ const learnMoreTreeItem = await children[1].getTreeItem()
+ assert.strictEqual(learnMoreTreeItem.label, 'Learn more about SageMaker Unified Studio')
+ assert.strictEqual(learnMoreTreeItem.contextValue, 'sageMakerUnifiedStudioLearnMore')
+ })
+
+ it('returns root nodes when authenticated', async function () {
+ mockDataZoneClient.listProjects.resolves({ projects: [mockProject], nextToken: undefined })
+
+ const children = await rootNode.getChildren()
+
+ assert.strictEqual(children.length, 2)
+ assert.ok(children[0] instanceof SageMakerUnifiedStudioAuthInfoNode)
+ assert.ok(children[1] instanceof SageMakerUnifiedStudioProjectNode)
+ // The first child is the auth info node, the second is the project node
+ assert.strictEqual(children[0].id, 'smusAuthInfoNode')
+ assert.strictEqual(children[1].id, 'smusProjectNode')
+
+ assert.strictEqual(children.length, 2)
+ assert.strictEqual(children[1].id, 'smusProjectNode')
+
+ const treeItem = await children[1].getTreeItem()
+ assert.strictEqual(treeItem.label, 'Select a project')
+ assert.strictEqual(treeItem.contextValue, 'smusProjectSelectPicker')
+ assert.deepStrictEqual(treeItem.command, {
+ command: 'aws.smus.projectView',
+ title: 'Select Project',
+ arguments: [children[1]],
+ })
+ })
+
+ it('returns auth info node when connection is expired', async function () {
+ // Create a mock auth provider with expired connection
+ const mockAuthProvider = {
+ isConnected: sinon.stub().returns(true),
+ isConnectionValid: sinon.stub().returns(false),
+ activeConnection: { domainId: testDomainId, ssoRegion: 'us-west-2' },
+ onDidChange: sinon.stub().returns({ dispose: sinon.stub() }),
+ showReauthenticationPrompt: sinon.stub(),
+ } as any
+
+ const mockExtensionContext = {
+ subscriptions: [],
+ workspaceState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ globalState: {
+ get: sinon.stub(),
+ update: sinon.stub(),
+ },
+ } as any
+
+ const expiredNode = new SageMakerUnifiedStudioRootNode(mockAuthProvider, mockExtensionContext)
+ const children = await expiredNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.ok(children[0] instanceof SageMakerUnifiedStudioAuthInfoNode)
+ assert.ok(mockAuthProvider.showReauthenticationPrompt.calledOnce)
+ })
+ })
+
+ describe('refresh', function () {
+ it('fires change events', function () {
+ const onDidChangeTreeItemSpy = sinon.spy()
+ const onDidChangeChildrenSpy = sinon.spy()
+
+ rootNode.onDidChangeTreeItem(onDidChangeTreeItemSpy)
+ rootNode.onDidChangeChildren(onDidChangeChildrenSpy)
+
+ rootNode.refresh()
+
+ assert(onDidChangeTreeItemSpy.calledOnce)
+ assert(onDidChangeChildrenSpy.calledOnce)
+ })
+ })
+})
+
+describe('SelectSMUSProject', function () {
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+ let mockProjectNode: sinon.SinonStubbedInstance
+ let createQuickPickStub: sinon.SinonStub
+ let executeCommandStub: sinon.SinonStub
+
+ const testDomainId = 'test-domain-123'
+ const mockProject: DataZoneProject = {
+ id: 'project-123',
+ name: 'Test Project',
+ description: 'Test Description',
+ domainId: testDomainId,
+ updatedAt: new Date(),
+ }
+
+ const mockProject2: DataZoneProject = {
+ id: 'project-456',
+ name: 'Another Project',
+ description: 'Another Description',
+ domainId: testDomainId,
+ updatedAt: new Date(Date.now() - 86400000), // 1 day ago
+ }
+
+ beforeEach(function () {
+ // Create mock DataZone client
+ mockDataZoneClient = {
+ getDomainId: sinon.stub().returns(testDomainId),
+ listProjects: sinon.stub(),
+ fetchAllProjects: sinon.stub(),
+ } as any
+
+ // Create mock project node
+ mockProjectNode = {
+ setProject: sinon.stub(),
+ getProject: sinon.stub().returns(undefined),
+ project: undefined,
+ } as any
+
+ // Stub DataZoneClient static methods
+ sinon.stub(DataZoneClient, 'getInstance').returns(mockDataZoneClient as any)
+
+ // Stub SmusAuthenticationProvider
+ sinon.stub(SmusAuthenticationProvider, 'fromContext').returns({
+ isConnected: sinon.stub().returns(true),
+ isConnectionValid: sinon.stub().returns(true),
+ activeConnection: { domainId: testDomainId, ssoRegion: 'us-west-2' },
+ getDomainId: sinon.stub().returns(testDomainId),
+ getDomainRegion: sinon.stub().returns('us-west-2'),
+ } as any)
+
+ // Stub quickPick - return the project directly (not wrapped in an item)
+ const mockQuickPick = {
+ prompt: sinon.stub().resolves(mockProject),
+ }
+ createQuickPickStub = sinon.stub(pickerPrompter, 'createQuickPick').returns(mockQuickPick as any)
+
+ // Stub vscode.commands.executeCommand
+ executeCommandStub = sinon.stub(vscode.commands, 'executeCommand')
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ it('fetches all projects and sets the project for first time', async function () {
+ mockDataZoneClient.fetchAllProjects.resolves([mockProject, mockProject2])
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+
+ assert.strictEqual(result, mockProject)
+ assert.ok(mockDataZoneClient.fetchAllProjects.calledOnce)
+ assert.ok(mockDataZoneClient.fetchAllProjects.calledWith())
+ assert.ok(createQuickPickStub.calledOnce)
+ assert.ok(mockProjectNode.setProject.calledOnce)
+ assert.ok(executeCommandStub.calledWith('aws.smus.rootView.refresh'))
+ assertTelemetry('smus_accessProject', {
+ result: 'Succeeded',
+ smusProjectId: mockProject.id,
+ })
+ })
+
+ it('filters out GenerativeAIModelGovernanceProject', async function () {
+ const governanceProject: DataZoneProject = {
+ id: 'governance-123',
+ name: 'GenerativeAIModelGovernanceProject',
+ description: 'Governance project',
+ domainId: testDomainId,
+ updatedAt: new Date(),
+ }
+
+ mockDataZoneClient.fetchAllProjects.resolves([mockProject, governanceProject, mockProject2])
+
+ await selectSMUSProject(mockProjectNode as any)
+
+ // Verify that the governance project is filtered out
+ const quickPickCall = createQuickPickStub.getCall(0)
+ const items = quickPickCall.args[0]
+ assert.strictEqual(items.length, 2) // Should only have mockProject and mockProject2
+ assert.ok(!items.some((item: any) => item.data.name === 'GenerativeAIModelGovernanceProject'))
+ })
+
+ it('handles no active connection', async function () {
+ sinon.restore()
+ sinon.stub(SmusAuthenticationProvider, 'fromContext').returns({
+ activeConnection: undefined,
+ getDomainId: sinon.stub().returns(undefined),
+ } as any)
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+ assert.strictEqual(result, undefined)
+
+ assertTelemetry('smus_accessProject', {
+ result: 'Succeeded',
+ })
+ })
+
+ it('fetches all projects and switches the current project', async function () {
+ mockProjectNode = {
+ setProject: sinon.stub(),
+ getProject: sinon.stub().returns(mockProject),
+ project: mockProject,
+ } as any
+ mockDataZoneClient.fetchAllProjects.resolves([mockProject, mockProject2])
+
+ // Stub quickPick to return mockProject2 for the second test
+ const mockQuickPick = {
+ prompt: sinon.stub().resolves(mockProject2),
+ }
+ createQuickPickStub.restore() // Remove the previous stub
+ createQuickPickStub = sinon.stub(pickerPrompter, 'createQuickPick').returns(mockQuickPick as any)
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+
+ assert.strictEqual(result, mockProject2)
+ assert.ok(mockDataZoneClient.fetchAllProjects.calledOnce)
+ assert.ok(mockDataZoneClient.fetchAllProjects.calledWith())
+ assert.ok(createQuickPickStub.calledOnce)
+ assert.ok(mockProjectNode.setProject.calledOnce)
+ assert.ok(executeCommandStub.calledWith('aws.smus.rootView.refresh'))
+ assertTelemetry('smus_accessProject', {
+ result: 'Succeeded',
+ smusProjectId: mockProject2.id,
+ })
+ })
+
+ it('shows message when no projects found', async function () {
+ mockDataZoneClient.fetchAllProjects.resolves([])
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+
+ assert.strictEqual(result, undefined)
+ assert.ok(!mockProjectNode.setProject.called)
+ })
+
+ it('handles API errors gracefully', async function () {
+ const error = new Error('API error')
+ mockDataZoneClient.fetchAllProjects.rejects(error)
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+ assert.strictEqual(result, undefined)
+
+ assert.ok(!mockProjectNode.setProject.called)
+ assertTelemetry('smus_accessProject', {
+ result: 'Succeeded',
+ })
+ })
+
+ it('handles case when user cancels project selection', async function () {
+ mockDataZoneClient.fetchAllProjects.resolves([mockProject, mockProject2])
+
+ // Make quickPick return undefined (user cancelled)
+ const mockQuickPick = {
+ prompt: sinon.stub().resolves(undefined),
+ }
+ createQuickPickStub.returns(mockQuickPick as any)
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+
+ // Should return undefined
+ assert.strictEqual(result, undefined)
+
+ // Verify project was not set
+ assert.ok(!mockProjectNode.setProject.called)
+
+ // Verify refresh command was not called
+ assert.ok(!executeCommandStub.called)
+ })
+
+ it('handles empty projects list correctly', async function () {
+ mockDataZoneClient.fetchAllProjects.resolves([])
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+
+ assert.strictEqual(result, undefined)
+ assert.ok(mockDataZoneClient.fetchAllProjects.calledOnce)
+ assert.ok(!mockProjectNode.setProject.called)
+ assert.ok(!executeCommandStub.called)
+ })
+})
+
+describe('selectSMUSProject - Additional Tests', function () {
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+ let mockProjectNode: sinon.SinonStubbedInstance
+ let createQuickPickStub: sinon.SinonStub
+ let executeCommandStub: sinon.SinonStub
+
+ const testDomainId = 'test-domain-123'
+ const mockProject: DataZoneProject = {
+ id: 'project-123',
+ name: 'Test Project',
+ description: 'Test Description',
+ domainId: testDomainId,
+ updatedAt: new Date(),
+ }
+
+ beforeEach(function () {
+ mockDataZoneClient = {
+ getDomainId: sinon.stub().returns(testDomainId),
+ fetchAllProjects: sinon.stub(),
+ } as any
+
+ mockProjectNode = {
+ setProject: sinon.stub(),
+ } as any
+
+ sinon.stub(DataZoneClient, 'getInstance').returns(mockDataZoneClient as any)
+ sinon.stub(SmusAuthenticationProvider, 'fromContext').returns({
+ activeConnection: { domainId: testDomainId, ssoRegion: 'us-west-2' },
+ getDomainId: sinon.stub().returns(testDomainId),
+ getDomainRegion: sinon.stub().returns('us-west-2'),
+ } as any)
+
+ const mockQuickPick = {
+ prompt: sinon.stub().resolves(mockProject),
+ }
+ createQuickPickStub = sinon.stub(pickerPrompter, 'createQuickPick').returns(mockQuickPick as any)
+ executeCommandStub = sinon.stub(vscode.commands, 'executeCommand')
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ it('handles access denied error gracefully', async function () {
+ const accessDeniedError = new Error('Access denied')
+ accessDeniedError.name = 'AccessDeniedError'
+ mockDataZoneClient.fetchAllProjects.rejects(accessDeniedError)
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+
+ assert.strictEqual(result, undefined)
+ assert.ok(
+ createQuickPickStub.calledWith([
+ {
+ label: '$(error)',
+ description: "You don't have permissions to view projects. Please contact your administrator",
+ },
+ ])
+ )
+ })
+
+ it('shows "No projects found" message when projects list is empty', async function () {
+ mockDataZoneClient.fetchAllProjects.resolves([])
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+
+ assert.strictEqual(result, undefined)
+ const testWindow = getTestWindow()
+ assert.ok(testWindow.shownMessages.some((msg) => msg.message === 'No projects found in the domain'))
+ assert.ok(
+ createQuickPickStub.calledWith([
+ {
+ label: 'No projects found',
+ detail: '',
+ description: '',
+ data: {},
+ },
+ ])
+ )
+ })
+
+ it('handles invalid selected project object', async function () {
+ mockDataZoneClient.fetchAllProjects.resolves([mockProject])
+
+ // Mock quickPick to return an object with 'type' property (invalid selection)
+ const mockQuickPick = {
+ prompt: sinon.stub().resolves({ type: 'invalid', data: mockProject }),
+ }
+ createQuickPickStub.returns(mockQuickPick as any)
+
+ const result = await selectSMUSProject(mockProjectNode as any)
+
+ assert.deepStrictEqual(result, { type: 'invalid', data: mockProject })
+ assert.ok(!mockProjectNode.setProject.called)
+ assert.ok(!executeCommandStub.called)
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode.test.ts
new file mode 100644
index 00000000000..a44b2ec3e7d
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode.test.ts
@@ -0,0 +1,280 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import { SagemakerUnifiedStudioSpaceNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
+import { SageMakerUnifiedStudioSpacesParentNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode'
+import { SagemakerClient, SagemakerSpaceApp } from '../../../../shared/clients/sagemaker'
+import { SagemakerSpace } from '../../../../awsService/sagemaker/sagemakerSpace'
+
+describe('SagemakerUnifiedStudioSpaceNode', function () {
+ let spaceNode: SagemakerUnifiedStudioSpaceNode
+ let mockParent: SageMakerUnifiedStudioSpacesParentNode
+ let mockSagemakerClient: SagemakerClient
+ let mockSpaceApp: SagemakerSpaceApp
+ let mockSagemakerSpace: sinon.SinonStubbedInstance
+ let trackPendingNodeStub: sinon.SinonStub
+
+ beforeEach(function () {
+ trackPendingNodeStub = sinon.stub()
+ mockParent = {
+ trackPendingNode: trackPendingNodeStub,
+ } as any
+
+ mockSagemakerClient = {
+ describeApp: sinon.stub(),
+ describeSpace: sinon.stub(),
+ } as any
+
+ mockSpaceApp = {
+ SpaceName: 'test-space',
+ DomainId: 'domain-123',
+ Status: 'InService',
+ DomainSpaceKey: 'domain-123:test-space',
+ App: {
+ AppName: 'test-app',
+ Status: 'InService',
+ },
+ } as any
+
+ mockSagemakerSpace = {
+ label: 'test-space (Running)',
+ description: 'Private space',
+ tooltip: new vscode.MarkdownString('Space tooltip'),
+ iconPath: { light: 'light-icon', dark: 'dark-icon' },
+ contextValue: 'smusSpaceNode',
+ updateSpace: sinon.stub(),
+ setSpaceStatus: sinon.stub(),
+ isPending: sinon.stub().returns(false),
+ getStatus: sinon.stub().returns('Running'),
+ getAppStatus: sinon.stub().resolves('InService'),
+ name: 'test-space',
+ arn: 'arn:aws:sagemaker:us-west-2:123456789012:space/test-space',
+ getAppArn: sinon.stub().resolves('arn:aws:sagemaker:us-west-2:123456789012:app/test-app'),
+ getSpaceArn: sinon.stub().resolves('arn:aws:sagemaker:us-west-2:123456789012:space/test-space'),
+ updateSpaceAppStatus: sinon.stub().resolves(),
+ buildTooltip: sinon.stub().returns('Space tooltip'),
+ getAppIcon: sinon.stub().returns({ light: 'light-icon', dark: 'dark-icon' }),
+ DomainSpaceKey: 'domain-123:test-space',
+ } as any
+
+ sinon.stub(SagemakerSpace.prototype, 'constructor' as any).returns(mockSagemakerSpace)
+
+ spaceNode = new SagemakerUnifiedStudioSpaceNode(
+ mockParent,
+ mockSagemakerClient,
+ 'us-west-2',
+ mockSpaceApp,
+ true
+ )
+
+ // Replace the internal smSpace with our mock
+ ;(spaceNode as any).smSpace = mockSagemakerSpace
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('creates instance with correct properties', function () {
+ assert.strictEqual(spaceNode.id, 'smusSpaceNodetest-space')
+ assert.strictEqual(spaceNode.resource, spaceNode)
+ assert.strictEqual(spaceNode.regionCode, 'us-west-2')
+ assert.strictEqual(spaceNode.spaceApp, mockSpaceApp)
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('returns correct tree item', function () {
+ const treeItem = spaceNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'test-space (Running)')
+ assert.strictEqual(treeItem.description, 'Private space')
+ assert.strictEqual(treeItem.contextValue, 'smusSpaceNode')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ assert.ok(treeItem.iconPath)
+ assert.ok(treeItem.tooltip)
+ })
+ })
+
+ describe('getChildren', function () {
+ it('returns empty array', function () {
+ const children = spaceNode.getChildren()
+ assert.deepStrictEqual(children, [])
+ })
+ })
+
+ describe('getParent', function () {
+ it('returns parent node', function () {
+ const parent = spaceNode.getParent()
+ assert.strictEqual(parent, mockParent)
+ })
+ })
+
+ describe('refreshNode', function () {
+ it('fires change event', async function () {
+ const emitterSpy = sinon.spy(spaceNode['onDidChangeEmitter'], 'fire')
+ await spaceNode.refreshNode()
+ assert(emitterSpy.calledOnce)
+ })
+ })
+
+ describe('updateSpace', function () {
+ it('updates space and tracks pending node when pending', function () {
+ mockSagemakerSpace.isPending.returns(true)
+ const newSpaceApp = { ...mockSpaceApp, Status: 'Pending' }
+
+ spaceNode.updateSpace(newSpaceApp)
+
+ assert(mockSagemakerSpace.updateSpace.calledWith(newSpaceApp))
+ assert(trackPendingNodeStub.calledWith('domain-123:test-space'))
+ })
+
+ it('updates space without tracking when not pending', function () {
+ mockSagemakerSpace.isPending.returns(false)
+ const newSpaceApp = { ...mockSpaceApp, Status: 'InService' }
+
+ spaceNode.updateSpace(newSpaceApp)
+
+ assert(mockSagemakerSpace.updateSpace.calledWith(newSpaceApp))
+ assert(trackPendingNodeStub.notCalled)
+ })
+ })
+
+ describe('setSpaceStatus', function () {
+ it('delegates to SagemakerSpace', function () {
+ spaceNode.setSpaceStatus('InService', 'Running')
+ assert(mockSagemakerSpace.setSpaceStatus.calledWith('InService', 'Running'))
+ })
+ })
+
+ describe('isPending', function () {
+ it('delegates to SagemakerSpace', function () {
+ const result = spaceNode.isPending()
+ assert(mockSagemakerSpace.isPending.called)
+ assert.strictEqual(result, false)
+ })
+ })
+
+ describe('getStatus', function () {
+ it('delegates to SagemakerSpace', function () {
+ const result = spaceNode.getStatus()
+ assert(mockSagemakerSpace.getStatus.called)
+ assert.strictEqual(result, 'Running')
+ })
+ })
+
+ describe('getAppStatus', function () {
+ it('delegates to SagemakerSpace', async function () {
+ const result = await spaceNode.getAppStatus()
+ assert(mockSagemakerSpace.getAppStatus.called)
+ assert.strictEqual(result, 'InService')
+ })
+ })
+
+ describe('name property', function () {
+ it('returns space name', function () {
+ assert.strictEqual(spaceNode.name, 'test-space')
+ })
+ })
+
+ describe('arn property', function () {
+ it('returns space arn', function () {
+ assert.strictEqual(spaceNode.arn, 'arn:aws:sagemaker:us-west-2:123456789012:space/test-space')
+ })
+ })
+
+ describe('getAppArn', function () {
+ it('delegates to SagemakerSpace', async function () {
+ const result = await spaceNode.getAppArn()
+ assert(mockSagemakerSpace.getAppArn.called)
+ assert.strictEqual(result, 'arn:aws:sagemaker:us-west-2:123456789012:app/test-app')
+ })
+ })
+
+ describe('getSpaceArn', function () {
+ it('delegates to SagemakerSpace', async function () {
+ const result = await spaceNode.getSpaceArn()
+ assert(mockSagemakerSpace.getSpaceArn.called)
+ assert.strictEqual(result, 'arn:aws:sagemaker:us-west-2:123456789012:space/test-space')
+ })
+ })
+
+ describe('updateSpaceAppStatus', function () {
+ it('updates status and tracks pending node when pending', async function () {
+ mockSagemakerSpace.isPending.returns(true)
+
+ await spaceNode.updateSpaceAppStatus()
+
+ assert(mockSagemakerSpace.updateSpaceAppStatus.called)
+ assert(trackPendingNodeStub.calledWith('domain-123:test-space'))
+ })
+
+ it('updates status without tracking when not pending', async function () {
+ mockSagemakerSpace.isPending.returns(false)
+
+ await spaceNode.updateSpaceAppStatus()
+
+ assert(mockSagemakerSpace.updateSpaceAppStatus.called)
+ assert(trackPendingNodeStub.notCalled)
+ })
+ })
+
+ describe('buildTooltip', function () {
+ it('delegates to SagemakerSpace', function () {
+ const result = spaceNode.buildTooltip()
+ assert(mockSagemakerSpace.buildTooltip.called)
+ assert.strictEqual(result, 'Space tooltip')
+ })
+ })
+
+ describe('getAppIcon', function () {
+ it('delegates to SagemakerSpace', function () {
+ const result = spaceNode.getAppIcon()
+ assert(mockSagemakerSpace.getAppIcon.called)
+ assert.deepStrictEqual(result, { light: 'light-icon', dark: 'dark-icon' })
+ })
+ })
+
+ describe('DomainSpaceKey property', function () {
+ it('returns domain space key', function () {
+ assert.strictEqual(spaceNode.DomainSpaceKey, 'domain-123:test-space')
+ })
+ })
+
+ describe('SagemakerSpace getContext for SMUS', function () {
+ it('returns awsSagemakerSpaceRunningNode for running SMUS space with undefined RemoteAccess', function () {
+ // Create a space app without RemoteAccess setting (undefined)
+ const smusSpaceApp = {
+ SpaceName: 'test-space',
+ DomainId: 'domain-123',
+ Status: 'InService',
+ DomainSpaceKey: 'domain-123:test-space',
+ App: {
+ AppName: 'test-app',
+ Status: 'InService',
+ },
+ SpaceSettingsSummary: {
+ // RemoteAccess is undefined
+ },
+ } as any
+
+ // Create a real SagemakerSpace instance for SMUS to test the actual getContext logic
+ const realSagemakerSpace = new SagemakerSpace(
+ mockSagemakerClient,
+ 'us-west-2',
+ smusSpaceApp,
+ true // isSMUSSpace = true
+ )
+
+ const context = realSagemakerSpace.getContext()
+
+ assert.strictEqual(context, 'awsSagemakerSpaceRunningNode')
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode.test.ts
new file mode 100644
index 00000000000..31481e70953
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode.test.ts
@@ -0,0 +1,421 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import sinon from 'sinon'
+import * as vscode from 'vscode'
+import { SageMakerUnifiedStudioSpacesParentNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpacesParentNode'
+import { SageMakerUnifiedStudioComputeNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioComputeNode'
+import { SagemakerUnifiedStudioSpaceNode } from '../../../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
+import { DataZoneClient } from '../../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { SagemakerClient } from '../../../../shared/clients/sagemaker'
+import { SmusAuthenticationProvider } from '../../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+import { getLogger } from '../../../../shared/logger/logger'
+import { SmusUtils } from '../../../../sagemakerunifiedstudio/shared/smusUtils'
+
+describe('SageMakerUnifiedStudioSpacesParentNode', function () {
+ let spacesNode: SageMakerUnifiedStudioSpacesParentNode
+ let mockParent: SageMakerUnifiedStudioComputeNode
+ let mockExtensionContext: vscode.ExtensionContext
+ let mockAuthProvider: SmusAuthenticationProvider
+ let mockSagemakerClient: sinon.SinonStubbedInstance
+ let mockDataZoneClient: sinon.SinonStubbedInstance
+
+ beforeEach(function () {
+ mockParent = {} as any
+ mockExtensionContext = {
+ extensionUri: vscode.Uri.file('/test'),
+ } as any
+ mockAuthProvider = {
+ activeConnection: { domainId: 'test-domain', ssoRegion: 'us-west-2' },
+ } as any
+ mockSagemakerClient = sinon.createStubInstance(SagemakerClient)
+ mockSagemakerClient.fetchSpaceAppsAndDomains.resolves([new Map(), new Map()])
+
+ mockDataZoneClient = {
+ getInstance: sinon.stub(),
+ getUserId: sinon.stub(),
+ getDomainId: sinon.stub(),
+ getRegion: sinon.stub(),
+ getToolingEnvironmentId: sinon.stub(),
+ getEnvironmentDetails: sinon.stub(),
+ getToolingEnvironment: sinon.stub(),
+ } as any
+
+ sinon.stub(DataZoneClient, 'getInstance').resolves(mockDataZoneClient as any)
+ sinon.stub(getLogger(), 'debug')
+ sinon.stub(getLogger(), 'error')
+ sinon.stub(SmusUtils, 'extractSSOIdFromUserId').returns('user-12345')
+
+ spacesNode = new SageMakerUnifiedStudioSpacesParentNode(
+ mockParent,
+ 'project-123',
+ mockExtensionContext,
+ mockAuthProvider,
+ mockSagemakerClient as any
+ )
+ })
+
+ afterEach(function () {
+ spacesNode.pollingSet.clear()
+ sinon.restore()
+ })
+
+ describe('constructor', function () {
+ it('creates instance with correct properties', function () {
+ assert.strictEqual(spacesNode.id, 'smusSpacesParentNode')
+ assert.strictEqual(spacesNode.resource, spacesNode)
+ })
+ })
+
+ describe('getTreeItem', function () {
+ it('returns correct tree item', async function () {
+ const treeItem = await spacesNode.getTreeItem()
+
+ assert.strictEqual(treeItem.label, 'Spaces')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.Expanded)
+ assert.strictEqual(treeItem.contextValue, 'smusSpacesNode')
+ assert.ok(treeItem.iconPath)
+ })
+ })
+
+ describe('getParent', function () {
+ it('returns parent node', function () {
+ const parent = spacesNode.getParent()
+ assert.strictEqual(parent, mockParent)
+ })
+ })
+
+ describe('getProjectId', function () {
+ it('returns project ID', function () {
+ assert.strictEqual(spacesNode.getProjectId(), 'project-123')
+ })
+ })
+
+ describe('getAuthProvider', function () {
+ it('returns auth provider', function () {
+ assert.strictEqual(spacesNode.getAuthProvider(), mockAuthProvider)
+ })
+ })
+
+ describe('refreshNode', function () {
+ it('fires change event', async function () {
+ const emitterSpy = sinon.spy(spacesNode['onDidChangeEmitter'], 'fire')
+ await spacesNode.refreshNode()
+ assert(emitterSpy.calledOnce)
+ })
+ })
+
+ describe('trackPendingNode', function () {
+ it('adds node to polling set', function () {
+ const addSpy = sinon.spy(spacesNode.pollingSet, 'add')
+ spacesNode.trackPendingNode('test-key')
+ assert(addSpy.calledWith('test-key'))
+ })
+ })
+
+ describe('getSpaceNodes', function () {
+ it('returns space node when found', function () {
+ const mockSpaceNode = {} as SagemakerUnifiedStudioSpaceNode
+ spacesNode['sagemakerSpaceNodes'].set('test-key', mockSpaceNode)
+
+ const result = spacesNode.getSpaceNodes('test-key')
+ assert.strictEqual(result, mockSpaceNode)
+ })
+
+ it('throws error when node not found', function () {
+ assert.throws(
+ () => spacesNode.getSpaceNodes('non-existent'),
+ /Node with id non-existent from polling set not found/
+ )
+ })
+ })
+
+ describe('getSageMakerDomainId', function () {
+ it('throws error when no active connection', async function () {
+ const mockAuthProviderNoConnection = {
+ activeConnection: undefined,
+ } as any
+
+ const spacesNodeNoConnection = new SageMakerUnifiedStudioSpacesParentNode(
+ mockParent,
+ 'project-123',
+ mockExtensionContext,
+ mockAuthProviderNoConnection,
+ mockSagemakerClient as any
+ )
+
+ await assert.rejects(
+ async () => await spacesNodeNoConnection.getSageMakerDomainId(),
+ /No active connection found to get SageMaker domain ID/
+ )
+ })
+
+ it('throws error when DataZone client not initialized', async function () {
+ ;(DataZoneClient.getInstance as sinon.SinonStub).resolves(undefined)
+
+ await assert.rejects(
+ async () => await spacesNode.getSageMakerDomainId(),
+ /DataZone client is not initialized/
+ )
+ })
+
+ it('throws error when tooling environment ID not found', async function () {
+ mockDataZoneClient.getDomainId.returns('domain-123')
+ const error = new Error('Failed to get tooling environment ID: Environment not found')
+ mockDataZoneClient.getToolingEnvironment.rejects(error)
+
+ await assert.rejects(
+ async () => await spacesNode.getSageMakerDomainId(),
+ /Failed to get tooling environment ID: Environment not found/
+ )
+ })
+
+ it('throws error when no default environment found', async function () {
+ mockDataZoneClient.getDomainId.returns('domain-123')
+ const error = new Error('No default environment found for project')
+ mockDataZoneClient.getToolingEnvironment.rejects(error)
+
+ await assert.rejects(
+ async () => await spacesNode.getSageMakerDomainId(),
+ /No default environment found for project/
+ )
+ })
+
+ it('throws error when SageMaker domain ID not found in resources', async function () {
+ mockDataZoneClient.getDomainId.returns('domain-123')
+ mockDataZoneClient.getToolingEnvironment.resolves({
+ projectId: 'project-123',
+ domainId: 'domain-123',
+ createdBy: 'user',
+ name: 'test-env',
+ awsAccountRegion: 'us-west-2',
+ provisionedResources: [{ name: 'otherResource', value: 'value', type: 'OTHER' }],
+ } as any)
+
+ await assert.rejects(
+ async () => await spacesNode.getSageMakerDomainId(),
+ /No SageMaker domain found in the tooling environment/
+ )
+ })
+
+ it('returns SageMaker domain ID when found', async function () {
+ mockDataZoneClient.getDomainId.returns('domain-123')
+ mockDataZoneClient.getToolingEnvironment.resolves({
+ projectId: 'project-123',
+ domainId: 'domain-123',
+ createdBy: 'user',
+ name: 'test-env',
+ awsAccountRegion: 'us-west-2',
+ provisionedResources: [
+ {
+ name: 'sageMakerDomainId',
+ value: 'sagemaker-domain-123',
+ type: 'SAGEMAKER_DOMAIN',
+ },
+ ],
+ } as any)
+
+ const result = await spacesNode.getSageMakerDomainId()
+ assert.strictEqual(result, 'sagemaker-domain-123')
+ })
+ })
+
+ describe('getChildren', function () {
+ let updateChildrenStub: sinon.SinonStub
+ let mockSpaceNode1: SagemakerUnifiedStudioSpaceNode
+ let mockSpaceNode2: SagemakerUnifiedStudioSpaceNode
+
+ beforeEach(function () {
+ updateChildrenStub = sinon.stub(spacesNode as any, 'updateChildren').resolves()
+ mockSpaceNode1 = { id: 'space1' } as any
+ mockSpaceNode2 = { id: 'space2' } as any
+ })
+
+ it('returns space nodes when spaces exist', async function () {
+ spacesNode['sagemakerSpaceNodes'].set('space1', mockSpaceNode1)
+ spacesNode['sagemakerSpaceNodes'].set('space2', mockSpaceNode2)
+
+ const children = await spacesNode.getChildren()
+
+ assert.strictEqual(children.length, 2)
+ assert(children.includes(mockSpaceNode1))
+ assert(children.includes(mockSpaceNode2))
+ assert(updateChildrenStub.calledOnce)
+ })
+
+ it('returns no spaces found node when no spaces exist', async function () {
+ const children = await spacesNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ const noSpacesNode = children[0]
+ assert.strictEqual(noSpacesNode.id, 'smusNoSpaces')
+
+ const treeItem = await noSpacesNode.getTreeItem()
+ assert.strictEqual(treeItem.label, '[No Spaces found]')
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ })
+
+ it('returns no spaces found node when updateChildren throws error', async function () {
+ updateChildrenStub.rejects(new Error('Update failed'))
+
+ const children = await spacesNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ assert.strictEqual(children[0].id, 'smusNoSpaces')
+ })
+
+ it('returns access denied node when AccessDeniedException is thrown', async function () {
+ const accessDeniedError = new Error('Access denied')
+ accessDeniedError.name = 'AccessDeniedException'
+ updateChildrenStub.rejects(accessDeniedError)
+
+ const children = await spacesNode.getChildren()
+
+ assert.strictEqual(children.length, 1)
+ const accessDeniedNode = children[0]
+ assert.strictEqual(accessDeniedNode.id, 'smusAccessDenied')
+
+ const treeItem = await accessDeniedNode.getTreeItem()
+ assert.ok(treeItem)
+ assert.strictEqual(
+ treeItem.label,
+ "You don't have permission to view spaces. Please contact your administrator."
+ )
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ assert.ok(treeItem.iconPath)
+ assert.strictEqual((treeItem.iconPath as vscode.ThemeIcon).id, 'error')
+ })
+ })
+
+ describe('updatePendingNodes', function () {
+ it('updates pending space nodes and removes from polling set when not pending', async function () {
+ const mockSpaceNode = {
+ DomainSpaceKey: 'test-key',
+ updateSpaceAppStatus: sinon.stub().resolves(),
+ isPending: sinon.stub().returns(false),
+ refreshNode: sinon.stub().resolves(),
+ } as any
+
+ spacesNode['sagemakerSpaceNodes'].set('test-key', mockSpaceNode)
+ spacesNode.pollingSet.add('test-key')
+
+ await spacesNode['updatePendingNodes']()
+
+ assert(mockSpaceNode.updateSpaceAppStatus.calledOnce)
+ assert(mockSpaceNode.refreshNode.calledOnce)
+ assert(!spacesNode.pollingSet.has('test-key'))
+ })
+
+ it('keeps pending nodes in polling set', async function () {
+ const mockSpaceNode = {
+ DomainSpaceKey: 'test-key',
+ updateSpaceAppStatus: sinon.stub().resolves(),
+ isPending: sinon.stub().returns(true),
+ refreshNode: sinon.stub().resolves(),
+ } as any
+
+ spacesNode['sagemakerSpaceNodes'].set('test-key', mockSpaceNode)
+ spacesNode.pollingSet.add('test-key')
+
+ await spacesNode['updatePendingNodes']()
+
+ assert(mockSpaceNode.updateSpaceAppStatus.calledOnce)
+ assert(mockSpaceNode.refreshNode.notCalled)
+ assert(spacesNode.pollingSet.has('test-key'))
+ })
+ })
+
+ describe('getAccessDeniedChildren', function () {
+ it('returns access denied tree node with error icon', async function () {
+ const accessDeniedChildren = spacesNode['getAccessDeniedChildren']()
+
+ assert.strictEqual(accessDeniedChildren.length, 1)
+ const accessDeniedNode = accessDeniedChildren[0]
+ assert.strictEqual(accessDeniedNode.id, 'smusAccessDenied')
+
+ const treeItem = await accessDeniedNode.getTreeItem()
+ assert.ok(treeItem)
+ assert.strictEqual(
+ treeItem.label,
+ "You don't have permission to view spaces. Please contact your administrator."
+ )
+ assert.strictEqual(treeItem.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ assert.ok(treeItem.iconPath)
+ assert.strictEqual((treeItem.iconPath as vscode.ThemeIcon).id, 'error')
+ })
+ })
+
+ describe('updateChildren', function () {
+ beforeEach(function () {
+ mockDataZoneClient.getUserId.resolves('ABCA4NU3S7PEOLDQPLXYZ:user-12345678-d061-70a4-0bf2-eeee67a6ab12')
+ mockDataZoneClient.getDomainId.returns('domain-123')
+ mockDataZoneClient.getRegion.returns('us-west-2')
+ mockDataZoneClient.getToolingEnvironment.resolves({
+ awsAccountRegion: 'us-west-2',
+ provisionedResources: [{ name: 'sageMakerDomainId', value: 'sagemaker-domain-123' }],
+ } as any)
+ })
+
+ it('filters spaces by current user ownership', async function () {
+ const spaceApps = new Map([
+ [
+ 'space1',
+ {
+ DomainId: 'domain-123',
+ OwnershipSettingsSummary: { OwnerUserProfileName: 'user-12345' },
+ DomainSpaceKey: 'space1',
+ },
+ ],
+ [
+ 'space2',
+ {
+ DomainId: 'domain-123',
+ OwnershipSettingsSummary: { OwnerUserProfileName: 'other-user' },
+ DomainSpaceKey: 'space2',
+ },
+ ],
+ ])
+ const domains = new Map([['domain-123', { DomainId: 'domain-123' }]])
+
+ mockSagemakerClient.fetchSpaceAppsAndDomains.resolves([spaceApps, domains])
+
+ await spacesNode['updateChildren']()
+
+ assert.strictEqual(spacesNode['spaceApps'].size, 1)
+ assert(spacesNode['spaceApps'].has('space1'))
+ assert(!spacesNode['spaceApps'].has('space2'))
+ })
+
+ it('creates space nodes for filtered spaces', async function () {
+ const spaceApps = new Map([
+ [
+ 'space1',
+ {
+ DomainId: 'domain-123',
+ OwnershipSettingsSummary: { OwnerUserProfileName: 'user-12345' },
+ DomainSpaceKey: 'space1',
+ },
+ ],
+ ])
+ const domains = new Map([['domain-123', { DomainId: 'domain-123' }]])
+
+ mockSagemakerClient.fetchSpaceAppsAndDomains.resolves([spaceApps, domains])
+
+ await spacesNode['updateChildren']()
+
+ assert.strictEqual(spacesNode['sagemakerSpaceNodes'].size, 1)
+ assert(spacesNode['sagemakerSpaceNodes'].has('space1'))
+ })
+
+ it('throws AccessDeniedException when fetchSpaceAppsAndDomains fails with access denied', async function () {
+ const accessDeniedError = new Error('Access denied to spaces')
+ accessDeniedError.name = 'AccessDeniedException'
+ mockSagemakerClient.fetchSpaceAppsAndDomains.rejects(accessDeniedError)
+
+ await assert.rejects(async () => await spacesNode['updateChildren'](), /Access denied to spaces/)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/utils.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/utils.test.ts
new file mode 100644
index 00000000000..cd92aa42981
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/utils.test.ts
@@ -0,0 +1,252 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as vscode from 'vscode'
+import {
+ getLabel,
+ isLeafNode,
+ getIconForNodeType,
+ createTreeItem,
+ createColumnTreeItem,
+ createErrorTreeItem,
+ isRedLakeDatabase,
+ getTooltip,
+ getRedshiftTypeFromHost,
+} from '../../../../sagemakerunifiedstudio/explorer/nodes/utils'
+import { NodeType, ConnectionType, RedshiftType } from '../../../../sagemakerunifiedstudio/explorer/nodes/types'
+
+describe('utils', function () {
+ describe('getLabel', function () {
+ it('should return container labels for container nodes', function () {
+ assert.strictEqual(getLabel({ id: 'test', nodeType: NodeType.REDSHIFT_TABLE, isContainer: true }), 'Tables')
+ assert.strictEqual(getLabel({ id: 'test', nodeType: NodeType.REDSHIFT_VIEW, isContainer: true }), 'Views')
+ assert.strictEqual(
+ getLabel({ id: 'test', nodeType: NodeType.REDSHIFT_FUNCTION, isContainer: true }),
+ 'Functions'
+ )
+ assert.strictEqual(
+ getLabel({ id: 'test', nodeType: NodeType.REDSHIFT_STORED_PROCEDURE, isContainer: true }),
+ 'Stored Procedures'
+ )
+ })
+
+ it('should return path label when available', function () {
+ assert.strictEqual(
+ getLabel({ id: 'test', nodeType: NodeType.S3_FILE, path: { label: 'custom-label' } }),
+ 'custom-label'
+ )
+ })
+
+ it('should return S3 folder name with trailing slash', function () {
+ assert.strictEqual(
+ getLabel({ id: 'test', nodeType: NodeType.S3_FOLDER, path: { key: 'folder/subfolder/' } }),
+ 'subfolder/'
+ )
+ })
+
+ it('should return S3 file name', function () {
+ assert.strictEqual(
+ getLabel({ id: 'test', nodeType: NodeType.S3_FILE, path: { key: 'folder/file.txt' } }),
+ 'file.txt'
+ )
+ })
+
+ it('should return last part of ID as fallback', function () {
+ assert.strictEqual(getLabel({ id: 'parent/child/node', nodeType: NodeType.CONNECTION }), 'node')
+ })
+ })
+
+ describe('isLeafNode', function () {
+ it('should return false for container nodes', function () {
+ assert.strictEqual(isLeafNode({ nodeType: NodeType.REDSHIFT_TABLE, isContainer: true }), false)
+ })
+
+ it('should return true for leaf node types', function () {
+ assert.strictEqual(isLeafNode({ nodeType: NodeType.S3_FILE }), true)
+ assert.strictEqual(isLeafNode({ nodeType: NodeType.REDSHIFT_COLUMN }), true)
+ assert.strictEqual(isLeafNode({ nodeType: NodeType.ERROR }), true)
+ assert.strictEqual(isLeafNode({ nodeType: NodeType.LOADING }), true)
+ assert.strictEqual(isLeafNode({ nodeType: NodeType.EMPTY }), true)
+ })
+
+ it('should return false for non-leaf node types', function () {
+ assert.strictEqual(isLeafNode({ nodeType: NodeType.CONNECTION }), false)
+ assert.strictEqual(isLeafNode({ nodeType: NodeType.REDSHIFT_CLUSTER }), false)
+ })
+ })
+
+ describe('getIconForNodeType', function () {
+ it('should return correct icons for different node types', function () {
+ const errorIcon = getIconForNodeType(NodeType.ERROR)
+ const loadingIcon = getIconForNodeType(NodeType.LOADING)
+
+ assert.ok(errorIcon instanceof vscode.ThemeIcon)
+ assert.strictEqual((errorIcon as vscode.ThemeIcon).id, 'error')
+ assert.ok(loadingIcon instanceof vscode.ThemeIcon)
+ assert.strictEqual((loadingIcon as vscode.ThemeIcon).id, 'loading~spin')
+ })
+
+ it('should return different icons for container vs non-container nodes', function () {
+ const containerIcon = getIconForNodeType(NodeType.REDSHIFT_TABLE, true)
+ const nonContainerIcon = getIconForNodeType(NodeType.REDSHIFT_TABLE, false)
+
+ assert.ok(containerIcon instanceof vscode.ThemeIcon)
+ assert.ok(nonContainerIcon instanceof vscode.ThemeIcon)
+ assert.strictEqual((containerIcon as vscode.ThemeIcon).id, 'table')
+ assert.strictEqual((nonContainerIcon as vscode.ThemeIcon).id, 'aws-redshift-table')
+ })
+
+ it('should return custom icon for GLUE_CATALOG', function () {
+ const catalogIcon = getIconForNodeType(NodeType.GLUE_CATALOG)
+
+ // The catalog icon should be a custom icon, not a ThemeIcon
+ assert.ok(catalogIcon)
+ // We can't easily test the exact icon path in unit tests, but we can verify it's not a ThemeIcon
+ assert.ok(
+ !(catalogIcon instanceof vscode.ThemeIcon) ||
+ (catalogIcon as any).id === 'aws-sagemakerunifiedstudio-catalog'
+ )
+ })
+ })
+
+ describe('createTreeItem', function () {
+ it('should create tree item with correct properties', function () {
+ const item = createTreeItem('Test Label', NodeType.CONNECTION, false, false, 'Test Tooltip')
+
+ assert.strictEqual(item.label, 'Test Label')
+ assert.strictEqual(item.collapsibleState, vscode.TreeItemCollapsibleState.Collapsed)
+ assert.strictEqual(item.contextValue, NodeType.CONNECTION)
+ assert.strictEqual(item.tooltip, 'Test Tooltip')
+ })
+
+ it('should create leaf node with None collapsible state', function () {
+ const item = createTreeItem('Leaf Node', NodeType.S3_FILE, true)
+
+ assert.strictEqual(item.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ })
+ })
+
+ describe('createColumnTreeItem', function () {
+ it('should create column tree item with type description', function () {
+ const item = createColumnTreeItem('column_name', 'VARCHAR(255)', NodeType.REDSHIFT_COLUMN)
+
+ assert.strictEqual(item.label, 'column_name')
+ assert.strictEqual(item.description, 'VARCHAR(255)')
+ assert.strictEqual(item.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ assert.strictEqual(item.contextValue, NodeType.REDSHIFT_COLUMN)
+ assert.strictEqual(item.tooltip, 'column_name: VARCHAR(255)')
+ })
+ })
+
+ describe('createErrorTreeItem', function () {
+ it('should create error tree item', function () {
+ const item = createErrorTreeItem('Error message')
+
+ assert.strictEqual(item.label, 'Error message')
+ assert.strictEqual(item.collapsibleState, vscode.TreeItemCollapsibleState.None)
+ assert.ok(item.iconPath instanceof vscode.ThemeIcon)
+ assert.strictEqual((item.iconPath as vscode.ThemeIcon).id, 'error')
+ })
+ })
+
+ describe('isRedLakeDatabase', function () {
+ it('should return true for RedLake database names', function () {
+ assert.strictEqual(isRedLakeDatabase('database@catalog'), true)
+ assert.strictEqual(isRedLakeDatabase('my-db@my-catalog'), true)
+ assert.strictEqual(isRedLakeDatabase('test_db@test_catalog'), true)
+ })
+
+ it('should return false for regular database names', function () {
+ assert.strictEqual(isRedLakeDatabase('regular_database'), false)
+ assert.strictEqual(isRedLakeDatabase('dev'), false)
+ assert.strictEqual(isRedLakeDatabase(''), false)
+ assert.strictEqual(isRedLakeDatabase(undefined), false)
+ })
+ })
+
+ describe('getTooltip', function () {
+ it('should return correct tooltip for connection nodes', function () {
+ const redshiftData = {
+ id: 'conn1',
+ nodeType: NodeType.CONNECTION,
+ connectionType: ConnectionType.REDSHIFT,
+ }
+ const s3Data = {
+ id: 'conn2',
+ nodeType: NodeType.CONNECTION,
+ connectionType: ConnectionType.S3,
+ }
+
+ assert.strictEqual(getTooltip(redshiftData), 'Redshift Connection: conn1')
+ assert.strictEqual(getTooltip(s3Data), 'Connection: conn2\nType: S3')
+ })
+
+ it('should return correct tooltip for S3 nodes', function () {
+ const bucketData = {
+ id: 'bucket1',
+ nodeType: NodeType.S3_BUCKET,
+ path: { bucket: 'my-bucket' },
+ }
+ const fileData = {
+ id: 'file1',
+ nodeType: NodeType.S3_FILE,
+ path: { bucket: 'my-bucket', key: 'folder/file.txt' },
+ }
+
+ assert.strictEqual(getTooltip(bucketData), 'S3 Bucket: my-bucket')
+ assert.strictEqual(getTooltip(fileData), 'File: file.txt\nBucket: my-bucket')
+ })
+
+ it('should return correct tooltip for Redshift container nodes', function () {
+ const containerData = {
+ id: 'tables',
+ nodeType: NodeType.REDSHIFT_TABLE,
+ isContainer: true,
+ path: { schema: 'public' },
+ }
+
+ assert.strictEqual(getTooltip(containerData), 'Tables in public')
+ })
+
+ it('should return correct tooltip for Redshift object nodes', function () {
+ const tableData = {
+ id: 'table1',
+ nodeType: NodeType.REDSHIFT_TABLE,
+ path: { schema: 'public' },
+ }
+
+ assert.strictEqual(getTooltip(tableData), 'Table: public.table1')
+ })
+ })
+
+ describe('getRedshiftTypeFromHost', function () {
+ it('should return undefined for invalid hosts', function () {
+ assert.strictEqual(getRedshiftTypeFromHost(undefined), undefined)
+ assert.strictEqual(getRedshiftTypeFromHost(''), undefined)
+ assert.strictEqual(getRedshiftTypeFromHost('invalid-host'), undefined)
+ })
+
+ it('should identify serverless hosts', function () {
+ const serverlessHost = 'workgroup.123456789012.us-east-1.redshift-serverless.amazonaws.com'
+ assert.strictEqual(getRedshiftTypeFromHost(serverlessHost), RedshiftType.Serverless)
+ })
+
+ it('should identify cluster hosts', function () {
+ const clusterHost = 'cluster.123456789012.us-east-1.redshift.amazonaws.com'
+ assert.strictEqual(getRedshiftTypeFromHost(clusterHost), RedshiftType.Cluster)
+ })
+
+ it('should handle hosts with port numbers', function () {
+ const hostWithPort = 'cluster.123456789012.us-east-1.redshift.amazonaws.com:5439'
+ assert.strictEqual(getRedshiftTypeFromHost(hostWithPort), RedshiftType.Cluster)
+ })
+
+ it('should return undefined for unrecognized domains', function () {
+ const unknownHost = 'host.example.com'
+ assert.strictEqual(getRedshiftTypeFromHost(unknownHost), undefined)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/client/clientStore.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/client/clientStore.test.ts
new file mode 100644
index 00000000000..e2c14ace96a
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/client/clientStore.test.ts
@@ -0,0 +1,148 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import { ConnectionClientStore } from '../../../../sagemakerunifiedstudio/shared/client/connectionClientStore'
+import { S3Client } from '../../../../sagemakerunifiedstudio/shared/client/s3Client'
+import { SQLWorkbenchClient } from '../../../../sagemakerunifiedstudio/shared/client/sqlWorkbenchClient'
+import { GlueClient } from '../../../../sagemakerunifiedstudio/shared/client/glueClient'
+import { GlueCatalogClient } from '../../../../sagemakerunifiedstudio/shared/client/glueCatalogClient'
+import { ConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+
+describe('ClientStore', function () {
+ let sandbox: sinon.SinonSandbox
+ let clientStore: ConnectionClientStore
+
+ const mockCredentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+ clientStore = ConnectionClientStore.getInstance()
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ clientStore.clearAll()
+ })
+
+ describe('getInstance', function () {
+ it('should return singleton instance', function () {
+ const instance1 = ConnectionClientStore.getInstance()
+ const instance2 = ConnectionClientStore.getInstance()
+ assert.strictEqual(instance1, instance2)
+ })
+ })
+
+ describe('getClient', function () {
+ it('should create and cache client', function () {
+ const factory = sandbox.stub().returns({ test: 'client' })
+
+ const client1 = clientStore.getClient('conn-1', 'TestClient', factory)
+ const client2 = clientStore.getClient('conn-1', 'TestClient', factory)
+
+ assert.strictEqual(client1, client2)
+ assert.ok(factory.calledOnce)
+ })
+
+ it('should create separate clients for different connections', function () {
+ const factory = sandbox.stub()
+ factory.onFirstCall().returns({ test: 'client1' })
+ factory.onSecondCall().returns({ test: 'client2' })
+
+ const client1 = clientStore.getClient('conn-1', 'TestClient', factory)
+ const client2 = clientStore.getClient('conn-2', 'TestClient', factory)
+
+ assert.notStrictEqual(client1, client2)
+ assert.ok(factory.calledTwice)
+ })
+ })
+
+ describe('getS3Client', function () {
+ it('should create S3Client with credentials provider', function () {
+ sandbox.stub(S3Client.prototype, 'constructor' as any)
+
+ const client = clientStore.getS3Client(
+ 'conn-1',
+ 'us-east-1',
+ mockCredentialsProvider as ConnectionCredentialsProvider
+ )
+
+ assert.ok(client instanceof S3Client)
+ })
+ })
+
+ describe('getSQLWorkbenchClient', function () {
+ it('should create SQLWorkbenchClient with credentials provider', function () {
+ const stub = sandbox.stub(SQLWorkbenchClient, 'createWithCredentials').returns({} as any)
+
+ clientStore.getSQLWorkbenchClient(
+ 'conn-1',
+ 'us-east-1',
+ mockCredentialsProvider as ConnectionCredentialsProvider
+ )
+
+ assert.ok(stub.calledOnce)
+ })
+ })
+
+ describe('getGlueClient', function () {
+ it('should create GlueClient with credentials provider', function () {
+ sandbox.stub(GlueClient.prototype, 'constructor' as any)
+
+ const client = clientStore.getGlueClient(
+ 'conn-1',
+ 'us-east-1',
+ mockCredentialsProvider as ConnectionCredentialsProvider
+ )
+
+ assert.ok(client instanceof GlueClient)
+ })
+ })
+
+ describe('getGlueCatalogClient', function () {
+ it('should create GlueCatalogClient with credentials provider', function () {
+ const stub = sandbox.stub(GlueCatalogClient, 'createWithCredentials').returns({} as any)
+
+ clientStore.getGlueCatalogClient(
+ 'conn-1',
+ 'us-east-1',
+ mockCredentialsProvider as ConnectionCredentialsProvider
+ )
+
+ assert.ok(stub.calledOnce)
+ })
+ })
+
+ describe('clearConnection', function () {
+ it('should clear cached clients for specific connection', function () {
+ const factory = sandbox.stub().returns({ test: 'client' })
+
+ clientStore.getClient('conn-1', 'TestClient', factory)
+ clientStore.clearConnection('conn-1')
+ clientStore.getClient('conn-1', 'TestClient', factory)
+
+ assert.strictEqual(factory.callCount, 2)
+ })
+ })
+
+ describe('clearAll', function () {
+ it('should clear all cached clients', function () {
+ const factory = sandbox.stub().returns({ test: 'client' })
+
+ clientStore.getClient('conn-1', 'TestClient', factory)
+ clientStore.clearAll()
+ clientStore.getClient('conn-1', 'TestClient', factory)
+
+ assert.strictEqual(factory.callCount, 2)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/client/credentialsAdapter.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/client/credentialsAdapter.test.ts
new file mode 100644
index 00000000000..cfb2cfbce6e
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/client/credentialsAdapter.test.ts
@@ -0,0 +1,53 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import * as AWS from 'aws-sdk'
+import { adaptConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/shared/client/credentialsAdapter'
+
+describe('credentialsAdapter', function () {
+ let sandbox: sinon.SinonSandbox
+ let mockConnectionCredentialsProvider: any
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+ mockConnectionCredentialsProvider = {
+ getCredentials: sandbox.stub(),
+ }
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ describe('adaptConnectionCredentialsProvider', function () {
+ it('should create CredentialProviderChain', function () {
+ const chain = adaptConnectionCredentialsProvider(mockConnectionCredentialsProvider)
+ assert.ok(chain instanceof AWS.CredentialProviderChain)
+ })
+
+ it('should create credentials with provider function', function () {
+ const chain = adaptConnectionCredentialsProvider(mockConnectionCredentialsProvider)
+ assert.ok(chain.providers)
+ assert.strictEqual(chain.providers.length, 1)
+ assert.strictEqual(typeof chain.providers[0], 'function')
+ })
+
+ it('should create AWS Credentials object', function () {
+ const chain = adaptConnectionCredentialsProvider(mockConnectionCredentialsProvider)
+ const provider = chain.providers[0] as () => AWS.Credentials
+ const credentials = provider()
+ assert.ok(credentials instanceof AWS.Credentials)
+ })
+
+ it('should set needsRefresh to always return true', function () {
+ const chain = adaptConnectionCredentialsProvider(mockConnectionCredentialsProvider)
+ const provider = chain.providers[0] as () => AWS.Credentials
+ const credentials = provider()
+ assert.strictEqual(credentials.needsRefresh(), true)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/client/datazoneClient.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/client/datazoneClient.test.ts
new file mode 100644
index 00000000000..38dbd5e33f5
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/client/datazoneClient.test.ts
@@ -0,0 +1,483 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import { DataZoneClient } from '../../../../sagemakerunifiedstudio/shared/client/datazoneClient'
+import { SmusAuthenticationProvider } from '../../../../sagemakerunifiedstudio/auth/providers/smusAuthenticationProvider'
+import { GetEnvironmentCommandOutput } from '@aws-sdk/client-datazone/dist-types/commands/GetEnvironmentCommand'
+
+describe('DataZoneClient', () => {
+ let dataZoneClient: DataZoneClient
+ let mockAuthProvider: any
+ const testDomainId = 'dzd_domainId'
+ const testRegion = 'us-east-2'
+
+ beforeEach(async () => {
+ // Create mock connection object
+ const mockConnection = {
+ domainId: testDomainId,
+ ssoRegion: testRegion,
+ }
+
+ // Create mock auth provider
+ mockAuthProvider = {
+ isConnected: sinon.stub().returns(true),
+ getDomainId: sinon.stub().returns(testDomainId),
+ getDomainRegion: sinon.stub().returns(testRegion),
+ activeConnection: mockConnection,
+ onDidChangeActiveConnection: sinon.stub().returns({
+ dispose: sinon.stub(),
+ }),
+ } as any
+
+ // Set up the DataZoneClient using getInstance since constructor is private
+ DataZoneClient.dispose()
+ dataZoneClient = await DataZoneClient.getInstance(mockAuthProvider)
+ })
+
+ afterEach(() => {
+ sinon.restore()
+ })
+
+ describe('getInstance', () => {
+ it('should return singleton instance', async () => {
+ const instance1 = await DataZoneClient.getInstance(mockAuthProvider)
+ const instance2 = await DataZoneClient.getInstance(mockAuthProvider)
+
+ assert.strictEqual(instance1, instance2)
+ })
+
+ it('should create new instance after dispose', async () => {
+ const instance1 = await DataZoneClient.getInstance(mockAuthProvider)
+ DataZoneClient.dispose()
+ const instance2 = await DataZoneClient.getInstance(mockAuthProvider)
+
+ assert.notStrictEqual(instance1, instance2)
+ })
+ })
+
+ describe('dispose', () => {
+ it('should clear singleton instance', async () => {
+ const instance = await DataZoneClient.getInstance(mockAuthProvider)
+ DataZoneClient.dispose()
+
+ // Should create new instance after dispose
+ const newInstance = await DataZoneClient.getInstance(mockAuthProvider)
+ assert.notStrictEqual(instance, newInstance)
+ })
+ })
+
+ describe('getRegion', () => {
+ it('should return configured region', () => {
+ const result = dataZoneClient.getRegion()
+ assert.strictEqual(typeof result, 'string')
+ assert.ok(result.length > 0)
+ })
+ })
+
+ describe('listProjects', () => {
+ it('should list projects with pagination', async () => {
+ const mockDataZone = {
+ listProjects: sinon.stub().resolves({
+ items: [
+ {
+ id: 'project-1',
+ name: 'Project 1',
+ description: 'First project',
+ createdAt: new Date('2023-01-01'),
+ updatedAt: new Date('2023-01-02'),
+ },
+ ],
+ nextToken: 'next-token',
+ }),
+ }
+
+ // Mock the getDataZoneClient method
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ const result = await dataZoneClient.listProjects({
+ maxResults: 10,
+ })
+
+ assert.strictEqual(result.projects.length, 1)
+ assert.strictEqual(result.projects[0].id, 'project-1')
+ assert.strictEqual(result.projects[0].name, 'Project 1')
+ assert.strictEqual(result.projects[0].domainId, testDomainId)
+ assert.strictEqual(result.nextToken, 'next-token')
+ })
+
+ it('should handle empty results', async () => {
+ const mockDataZone = {
+ listProjects: sinon.stub().resolves({
+ items: [],
+ nextToken: undefined,
+ }),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ const result = await dataZoneClient.listProjects()
+
+ assert.strictEqual(result.projects.length, 0)
+ assert.strictEqual(result.nextToken, undefined)
+ })
+
+ it('should handle API errors', async () => {
+ const error = new Error('API Error')
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').rejects(error)
+
+ await assert.rejects(() => dataZoneClient.listProjects(), error)
+ })
+ })
+
+ describe('getProjectDefaultEnvironmentCreds', () => {
+ it('should get environment credentials for project', async () => {
+ const mockCredentials = {
+ accessKeyId: 'AKIATEST',
+ secretAccessKey: 'secret',
+ sessionToken: 'token',
+ }
+
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().resolves({
+ items: [{ id: 'blueprint-1', name: 'Tooling' }],
+ }),
+ listEnvironments: sinon.stub().resolves({
+ items: [{ id: 'env-1', name: 'Tooling' }],
+ }),
+ getEnvironmentCredentials: sinon.stub().resolves(mockCredentials),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ const result = await dataZoneClient.getProjectDefaultEnvironmentCreds('project-1')
+
+ assert.deepStrictEqual(result, mockCredentials)
+ assert.ok(
+ mockDataZone.listEnvironmentBlueprints.calledWith({
+ domainIdentifier: testDomainId,
+ managed: true,
+ name: 'Tooling',
+ })
+ )
+ assert.ok(
+ mockDataZone.listEnvironments.calledWith({
+ domainIdentifier: testDomainId,
+ projectIdentifier: 'project-1',
+ environmentBlueprintIdentifier: 'blueprint-1',
+ provider: 'Amazon SageMaker',
+ })
+ )
+ assert.ok(
+ mockDataZone.getEnvironmentCredentials.calledWith({
+ domainIdentifier: testDomainId,
+ environmentIdentifier: 'env-1',
+ })
+ )
+ })
+
+ it('should throw error when tooling blueprint not found', async () => {
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().resolves({
+ items: [],
+ }),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ await assert.rejects(
+ () => dataZoneClient.getProjectDefaultEnvironmentCreds('project-1'),
+ /Failed to get tooling blueprint/
+ )
+ })
+
+ it('should throw error when default environment not found', async () => {
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().resolves({
+ items: [{ id: 'blueprint-1', name: 'Tooling' }],
+ }),
+ listEnvironments: sinon.stub().resolves({
+ items: [],
+ }),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ await assert.rejects(
+ () => dataZoneClient.getProjectDefaultEnvironmentCreds('project-1'),
+ /Failed to find default Tooling environment/
+ )
+ })
+ })
+
+ describe('fetchAllProjects', function () {
+ it('fetches all projects by handling pagination', async function () {
+ const client = await DataZoneClient.getInstance(mockAuthProvider)
+
+ // Create a stub for listProjects that returns paginated results
+ const listProjectsStub = sinon.stub()
+
+ // First call returns first page with nextToken
+ listProjectsStub.onFirstCall().resolves({
+ projects: [
+ {
+ id: 'project-1',
+ name: 'Project 1',
+ description: 'First project',
+ domainId: testDomainId,
+ },
+ ],
+ nextToken: 'next-page-token',
+ })
+
+ // Second call returns second page with no nextToken
+ listProjectsStub.onSecondCall().resolves({
+ projects: [
+ {
+ id: 'project-2',
+ name: 'Project 2',
+ description: 'Second project',
+ domainId: testDomainId,
+ },
+ ],
+ nextToken: undefined,
+ })
+
+ // Replace the listProjects method with our stub
+ client.listProjects = listProjectsStub
+
+ // Call fetchAllProjects
+ const result = await client.fetchAllProjects()
+
+ // Verify results
+ assert.strictEqual(result.length, 2)
+ assert.strictEqual(result[0].id, 'project-1')
+ assert.strictEqual(result[1].id, 'project-2')
+
+ // Verify listProjects was called correctly
+ assert.strictEqual(listProjectsStub.callCount, 2)
+ assert.deepStrictEqual(listProjectsStub.firstCall.args[0], {
+ maxResults: 50,
+ nextToken: undefined,
+ })
+ assert.deepStrictEqual(listProjectsStub.secondCall.args[0], {
+ maxResults: 50,
+ nextToken: 'next-page-token',
+ })
+ })
+
+ it('returns empty array when no projects found', async function () {
+ const client = await DataZoneClient.getInstance(mockAuthProvider)
+
+ // Create a stub for listProjects that returns empty results
+ const listProjectsStub = sinon.stub().resolves({
+ projects: [],
+ nextToken: undefined,
+ })
+
+ // Replace the listProjects method with our stub
+ client.listProjects = listProjectsStub
+
+ // Call fetchAllProjects
+ const result = await client.fetchAllProjects()
+
+ // Verify results
+ assert.strictEqual(result.length, 0)
+ assert.strictEqual(listProjectsStub.callCount, 1)
+ })
+
+ it('handles errors gracefully', async function () {
+ const client = await DataZoneClient.getInstance(mockAuthProvider)
+
+ // Create a stub for listProjects that throws an error
+ const listProjectsStub = sinon.stub().rejects(new Error('API error'))
+
+ // Replace the listProjects method with our stub
+ client.listProjects = listProjectsStub
+
+ // Call fetchAllProjects and expect it to throw
+ await assert.rejects(() => client.fetchAllProjects(), /API error/)
+ })
+ })
+
+ describe('getToolingEnvironmentId', () => {
+ it('should get tooling environment ID successfully', async () => {
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().resolves({
+ items: [{ id: 'blueprint-1', name: 'Tooling' }],
+ }),
+ listEnvironments: sinon.stub().resolves({
+ items: [{ id: 'env-1', name: 'Tooling' }],
+ }),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ const result = await dataZoneClient.getToolingEnvironmentId('domain-1', 'project-1')
+
+ assert.strictEqual(result, 'env-1')
+ })
+
+ it('should handle listEnvironmentBlueprints error', async () => {
+ const error = new Error('Blueprint API Error')
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().rejects(error),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ await assert.rejects(() => dataZoneClient.getToolingEnvironmentId('domain-1', 'project-1'), error)
+ })
+
+ it('should handle listEnvironments error', async () => {
+ const error = new Error('Environment API Error')
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().resolves({
+ items: [{ id: 'blueprint-1', name: 'Tooling' }],
+ }),
+ listEnvironments: sinon.stub().rejects(error),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ await assert.rejects(() => dataZoneClient.getToolingEnvironmentId('domain-1', 'project-1'), error)
+ })
+ })
+
+ describe('getToolingEnvironment', () => {
+ beforeEach(() => {
+ mockAuthProvider = {} as SmusAuthenticationProvider
+ })
+
+ it('should return environment details when successful', async () => {
+ const mockEnvironment: GetEnvironmentCommandOutput = {
+ id: 'env-123',
+ awsAccountRegion: 'us-east-1',
+ projectId: undefined,
+ domainId: undefined,
+ createdBy: undefined,
+ name: undefined,
+ provider: undefined,
+ $metadata: {},
+ }
+
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().resolves({
+ items: [{ id: 'blueprint-1', name: 'Tooling' }],
+ }),
+ listEnvironments: sinon.stub().resolves({
+ items: [{ id: 'env-1', name: 'Tooling' }],
+ }),
+ getEnvironment: sinon.stub().resolves(mockEnvironment),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ const result = await dataZoneClient.getToolingEnvironment('project-123')
+
+ assert.strictEqual(result, mockEnvironment)
+ })
+
+ it('should throw error when no tooling environment ID found', async () => {
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().resolves({
+ items: [{ id: 'blueprint-1', name: 'Tooling' }],
+ }),
+ listEnvironments: sinon.stub().resolves({
+ items: [],
+ }),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ await assert.rejects(
+ () => dataZoneClient.getToolingEnvironment('project-123'),
+ /Failed to get tooling environment ID: No default Tooling environment found for project/
+ )
+ })
+
+ it('should throw error when getToolingEnvironmentId fails', async () => {
+ const mockDataZone = {
+ listEnvironmentBlueprints: sinon.stub().rejects(new Error('API error')),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ await assert.rejects(() => dataZoneClient.getToolingEnvironment('project-123'), /API error/)
+ })
+ })
+
+ describe('fetchAllProjectMemberships', () => {
+ it('should fetch all project memberships with pagination', async () => {
+ const mockDataZone = {
+ listProjectMemberships: sinon.stub(),
+ }
+
+ // First call returns first page with nextToken
+ mockDataZone.listProjectMemberships.onFirstCall().resolves({
+ members: [
+ {
+ memberDetails: {
+ user: {
+ userId: 'user-1',
+ },
+ },
+ },
+ ],
+ nextToken: 'next-token',
+ })
+
+ // Second call returns second page without nextToken
+ mockDataZone.listProjectMemberships.onSecondCall().resolves({
+ members: [
+ {
+ memberDetails: {
+ user: {
+ userId: 'user-2',
+ },
+ },
+ },
+ ],
+ nextToken: undefined,
+ })
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ const result = await dataZoneClient.fetchAllProjectMemberships('project-1')
+
+ assert.strictEqual(result.length, 2)
+ assert.strictEqual(result[0].memberDetails?.user?.userId, 'user-1')
+ assert.strictEqual(result[1].memberDetails?.user?.userId, 'user-2')
+ assert.strictEqual(mockDataZone.listProjectMemberships.callCount, 2)
+ })
+
+ it('should handle empty memberships', async () => {
+ const mockDataZone = {
+ listProjectMemberships: sinon.stub().resolves({
+ members: [],
+ nextToken: undefined,
+ }),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ const result = await dataZoneClient.fetchAllProjectMemberships('project-1')
+
+ assert.strictEqual(result.length, 0)
+ })
+
+ it('should handle API errors', async () => {
+ const error = new Error('Membership API Error')
+ const mockDataZone = {
+ listProjectMemberships: sinon.stub().rejects(error),
+ }
+
+ sinon.stub(dataZoneClient as any, 'getDataZoneClient').resolves(mockDataZone)
+
+ await assert.rejects(() => dataZoneClient.fetchAllProjectMemberships('project-1'), error)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/client/glueClient.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/client/glueClient.test.ts
new file mode 100644
index 00000000000..cd34fe7703e
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/client/glueClient.test.ts
@@ -0,0 +1,202 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import { GlueClient } from '../../../../sagemakerunifiedstudio/shared/client/glueClient'
+import { Glue, GetDatabasesCommand, GetTablesCommand, GetTableCommand } from '@aws-sdk/client-glue'
+import { ConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+
+describe('GlueClient', function () {
+ let sandbox: sinon.SinonSandbox
+ let glueClient: GlueClient
+ let mockGlue: sinon.SinonStubbedInstance
+
+ const mockCredentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+
+ mockGlue = {
+ send: sandbox.stub(),
+ } as any
+
+ sandbox.stub(Glue.prototype, 'send').callsFake(mockGlue.send)
+
+ glueClient = new GlueClient('us-east-1', mockCredentialsProvider as ConnectionCredentialsProvider)
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ describe('getDatabases', function () {
+ it('should get databases successfully', async function () {
+ const mockResponse = {
+ DatabaseList: [
+ { Name: 'database1', Description: 'Test database 1' },
+ { Name: 'database2', Description: 'Test database 2' },
+ ],
+ NextToken: 'next-token',
+ }
+
+ mockGlue.send.resolves(mockResponse)
+
+ const result = await glueClient.getDatabases('test-catalog', undefined, undefined, 'start-token')
+
+ assert.strictEqual(result.databases.length, 2)
+ assert.strictEqual(result.databases[0].Name, 'database1')
+ assert.strictEqual(result.nextToken, 'next-token')
+
+ const sendCall = mockGlue.send.getCall(0)
+ const command = sendCall.args[0] as GetDatabasesCommand
+ assert.ok(command instanceof GetDatabasesCommand)
+ })
+
+ it('should get databases without catalog ID', async function () {
+ const mockResponse = {
+ DatabaseList: [{ Name: 'default-db' }],
+ }
+
+ mockGlue.send.resolves(mockResponse)
+
+ const result = await glueClient.getDatabases()
+
+ assert.strictEqual(result.databases.length, 1)
+ assert.strictEqual(result.databases[0].Name, 'default-db')
+ assert.strictEqual(result.nextToken, undefined)
+ })
+
+ it('should handle errors when getting databases', async function () {
+ const error = new Error('Access denied')
+ mockGlue.send.rejects(error)
+
+ await assert.rejects(
+ async () => {
+ await glueClient.getDatabases('test-catalog')
+ },
+ {
+ message: 'Access denied',
+ }
+ )
+ })
+ })
+
+ describe('getTables', function () {
+ it('should get tables successfully', async function () {
+ const mockResponse = {
+ TableList: [
+ { Name: 'table1', DatabaseName: 'test-db' },
+ { Name: 'table2', DatabaseName: 'test-db' },
+ ],
+ NextToken: 'next-token',
+ }
+
+ mockGlue.send.resolves(mockResponse)
+
+ const result = await glueClient.getTables('test-db', 'test-catalog', undefined, 'start-token')
+
+ assert.strictEqual(result.tables.length, 2)
+ assert.strictEqual(result.tables[0].Name, 'table1')
+ assert.strictEqual(result.nextToken, 'next-token')
+
+ const sendCall = mockGlue.send.getCall(0)
+ const command = sendCall.args[0] as GetTablesCommand
+ assert.ok(command instanceof GetTablesCommand)
+ })
+
+ it('should get tables without catalog ID', async function () {
+ const mockResponse = {
+ TableList: [{ Name: 'default-table' }],
+ }
+
+ mockGlue.send.resolves(mockResponse)
+
+ const result = await glueClient.getTables('test-db')
+
+ assert.strictEqual(result.tables.length, 1)
+ assert.strictEqual(result.tables[0].Name, 'default-table')
+ })
+
+ it('should handle errors when getting tables', async function () {
+ const error = new Error('Database not found')
+ mockGlue.send.rejects(error)
+
+ await assert.rejects(
+ async () => {
+ await glueClient.getTables('nonexistent-db')
+ },
+ {
+ message: 'Database not found',
+ }
+ )
+ })
+ })
+
+ describe('getTable', function () {
+ it('should get table details successfully', async function () {
+ const mockResponse = {
+ Table: {
+ Name: 'test-table',
+ DatabaseName: 'test-db',
+ StorageDescriptor: {
+ Columns: [
+ { Name: 'col1', Type: 'string' },
+ { Name: 'col2', Type: 'int' },
+ ],
+ },
+ PartitionKeys: [{ Name: 'partition_col', Type: 'date' }],
+ },
+ }
+
+ mockGlue.send.resolves(mockResponse)
+
+ const result = await glueClient.getTable('test-db', 'test-table', 'test-catalog')
+
+ assert.strictEqual(result?.Name, 'test-table')
+ assert.strictEqual(result?.StorageDescriptor?.Columns?.length, 2)
+ assert.strictEqual(result?.PartitionKeys?.length, 1)
+
+ const sendCall = mockGlue.send.getCall(0)
+ const command = sendCall.args[0] as GetTableCommand
+ assert.ok(command instanceof GetTableCommand)
+ })
+
+ it('should get table without catalog ID', async function () {
+ const mockResponse = {
+ Table: {
+ Name: 'default-table',
+ DatabaseName: 'default-db',
+ },
+ }
+
+ mockGlue.send.resolves(mockResponse)
+
+ const result = await glueClient.getTable('default-db', 'default-table')
+
+ assert.strictEqual(result?.Name, 'default-table')
+ })
+
+ it('should handle errors when getting table', async function () {
+ const error = new Error('Table not found')
+ mockGlue.send.rejects(error)
+
+ await assert.rejects(
+ async () => {
+ await glueClient.getTable('test-db', 'nonexistent-table')
+ },
+ {
+ message: 'Table not found',
+ }
+ )
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/client/gluePrivateClient.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/client/gluePrivateClient.test.ts
new file mode 100644
index 00000000000..22a97d82caf
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/client/gluePrivateClient.test.ts
@@ -0,0 +1,143 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import globals from '../../../../shared/extensionGlobals'
+import { GlueCatalogClient } from '../../../../sagemakerunifiedstudio/shared/client/glueCatalogClient'
+import { ConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+
+describe('GlueCatalogClient', function () {
+ let sandbox: sinon.SinonSandbox
+ let mockGlueClient: any
+ let mockSdkClientBuilder: any
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+
+ mockGlueClient = {
+ getCatalogs: sandbox.stub().returns({
+ promise: sandbox.stub().resolves({
+ CatalogList: [
+ {
+ Name: 'test-catalog',
+ CatalogType: 'HIVE',
+ Parameters: { key1: 'value1' },
+ },
+ ],
+ }),
+ }),
+ }
+
+ mockSdkClientBuilder = {
+ createAwsService: sandbox.stub().resolves(mockGlueClient),
+ }
+
+ sandbox.stub(globals, 'sdkClientBuilder').value(mockSdkClientBuilder)
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ // Reset singleton instance
+ ;(GlueCatalogClient as any).instance = undefined
+ })
+
+ describe('getInstance', function () {
+ it('should create singleton instance', function () {
+ const client1 = GlueCatalogClient.getInstance('us-east-1')
+ const client2 = GlueCatalogClient.getInstance('us-east-1')
+
+ assert.strictEqual(client1, client2)
+ })
+
+ it('should return region correctly', function () {
+ const client = GlueCatalogClient.getInstance('us-west-2')
+ assert.strictEqual(client.getRegion(), 'us-west-2')
+ })
+ })
+
+ describe('createWithCredentials', function () {
+ it('should create client with credentials', function () {
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+
+ const client = GlueCatalogClient.createWithCredentials(
+ 'us-east-1',
+ credentialsProvider as ConnectionCredentialsProvider
+ )
+ assert.strictEqual(client.getRegion(), 'us-east-1')
+ })
+ })
+
+ describe('getCatalogs', function () {
+ it('should return catalogs successfully', async function () {
+ const client = GlueCatalogClient.getInstance('us-east-1')
+ const catalogs = await client.getCatalogs()
+
+ assert.strictEqual(catalogs.catalogs.length, 1)
+ assert.strictEqual(catalogs.catalogs[0].Name, 'test-catalog')
+ assert.strictEqual(catalogs.catalogs[0].CatalogType, 'HIVE')
+ assert.deepStrictEqual(catalogs.catalogs[0].Parameters, { key1: 'value1' })
+ })
+
+ it('should return empty array when no catalogs found', async function () {
+ mockGlueClient.getCatalogs.returns({
+ promise: sandbox.stub().resolves({ CatalogList: [] }),
+ })
+
+ const client = GlueCatalogClient.getInstance('us-east-1')
+ const catalogs = await client.getCatalogs()
+
+ assert.strictEqual(catalogs.catalogs.length, 0)
+ })
+
+ it('should handle API errors', async function () {
+ const error = new Error('API Error')
+ mockGlueClient.getCatalogs.returns({
+ promise: sandbox.stub().rejects(error),
+ })
+
+ const client = GlueCatalogClient.getInstance('us-east-1')
+
+ await assert.rejects(async () => await client.getCatalogs(), error)
+ })
+
+ it('should create client with credentials when provided', async function () {
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+
+ const client = GlueCatalogClient.createWithCredentials(
+ 'us-east-1',
+ credentialsProvider as ConnectionCredentialsProvider
+ )
+ await client.getCatalogs()
+
+ assert.ok(mockSdkClientBuilder.createAwsService.calledOnce)
+ const callArgs = mockSdkClientBuilder.createAwsService.getCall(0).args[1]
+ assert.ok(callArgs.credentialProvider)
+ assert.strictEqual(callArgs.region, 'us-east-1')
+ })
+
+ it('should create client without credentials when not provided', async function () {
+ const client = GlueCatalogClient.getInstance('us-east-1')
+ await client.getCatalogs()
+
+ assert.ok(mockSdkClientBuilder.createAwsService.calledOnce)
+ const callArgs = mockSdkClientBuilder.createAwsService.getCall(0).args[1]
+ assert.strictEqual(callArgs.region, 'us-east-1')
+ assert.ok(!callArgs.credentials)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/client/s3Client.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/client/s3Client.test.ts
new file mode 100644
index 00000000000..714ced3d446
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/client/s3Client.test.ts
@@ -0,0 +1,306 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import { S3Client } from '../../../../sagemakerunifiedstudio/shared/client/s3Client'
+import { S3 } from '@aws-sdk/client-s3'
+import { ConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+
+describe('S3Client', function () {
+ let sandbox: sinon.SinonSandbox
+ let mockS3: sinon.SinonStubbedInstance
+ let s3Client: S3Client
+
+ const mockCredentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+
+ mockS3 = {
+ listObjectsV2: sandbox.stub(),
+ } as any
+
+ sandbox.stub(S3.prototype, 'constructor' as any)
+ sandbox.stub(S3.prototype, 'listObjectsV2').callsFake(mockS3.listObjectsV2)
+
+ s3Client = new S3Client('us-east-1', mockCredentialsProvider as ConnectionCredentialsProvider)
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ describe('constructor', function () {
+ it('should create client with correct properties', function () {
+ const client = new S3Client('us-west-2', mockCredentialsProvider as ConnectionCredentialsProvider)
+ assert.ok(client)
+ })
+ })
+
+ describe('listPaths', function () {
+ it('should list folders and files successfully', async function () {
+ const mockResponse = {
+ CommonPrefixes: [{ Prefix: 'folder1/' }, { Prefix: 'folder2/' }],
+ Contents: [
+ {
+ Key: 'file1.txt',
+ Size: 1024,
+ LastModified: new Date('2023-01-01'),
+ },
+ {
+ Key: 'file2.txt',
+ Size: 2048,
+ LastModified: new Date('2023-01-02'),
+ },
+ ],
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('test-bucket')
+
+ assert.strictEqual(result.paths.length, 4)
+ const paths = result.paths
+
+ // Check folders
+ assert.strictEqual(paths[0].displayName, 'folder1')
+ assert.strictEqual(paths[0].isFolder, true)
+ assert.strictEqual(paths[0].bucket, 'test-bucket')
+ assert.strictEqual(paths[0].prefix, 'folder1/')
+
+ assert.strictEqual(paths[1].displayName, 'folder2')
+ assert.strictEqual(paths[1].isFolder, true)
+
+ // Check files
+ assert.strictEqual(paths[2].displayName, 'file1.txt')
+ assert.strictEqual(paths[2].isFolder, false)
+ assert.strictEqual(paths[2].size, 1024)
+ assert.deepStrictEqual(paths[2].lastModified, new Date('2023-01-01'))
+
+ assert.strictEqual(paths[3].displayName, 'file2.txt')
+ assert.strictEqual(paths[3].isFolder, false)
+ assert.strictEqual(paths[3].size, 2048)
+ })
+
+ it('should list paths with prefix', async function () {
+ const mockResponse = {
+ CommonPrefixes: [{ Prefix: 'prefix/subfolder/' }],
+ Contents: [
+ {
+ Key: 'prefix/file.txt',
+ Size: 512,
+ LastModified: new Date('2023-01-01'),
+ },
+ ],
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('test-bucket', 'prefix/')
+
+ assert.strictEqual(result.paths.length, 2)
+ const paths = result.paths
+ assert.strictEqual(paths[0].displayName, 'subfolder')
+ assert.strictEqual(paths[0].isFolder, true)
+ assert.strictEqual(paths[1].displayName, 'file.txt')
+ assert.strictEqual(paths[1].isFolder, false)
+
+ // Verify API call
+ assert.ok(mockS3.listObjectsV2.calledOnce)
+ const callArgs = mockS3.listObjectsV2.getCall(0).args[0]
+ assert.strictEqual(callArgs.Bucket, 'test-bucket')
+ assert.strictEqual(callArgs.Prefix, 'prefix/')
+ assert.strictEqual(callArgs.Delimiter, '/')
+ assert.strictEqual(callArgs.ContinuationToken, undefined)
+ })
+
+ it('should return empty array when no objects found', async function () {
+ const mockResponse = {
+ CommonPrefixes: [],
+ Contents: [],
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('empty-bucket')
+
+ assert.strictEqual(result.paths.length, 0)
+ })
+
+ it('should handle response with only folders', async function () {
+ const mockResponse = {
+ CommonPrefixes: [{ Prefix: 'folder1/' }, { Prefix: 'folder2/' }],
+ Contents: undefined,
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('test-bucket')
+
+ assert.strictEqual(result.paths.length, 2)
+ const paths = result.paths
+ assert.strictEqual(paths[0].isFolder, true)
+ assert.strictEqual(paths[1].isFolder, true)
+ })
+
+ it('should handle response with only files', async function () {
+ const mockResponse = {
+ CommonPrefixes: undefined,
+ Contents: [
+ {
+ Key: 'file1.txt',
+ Size: 1024,
+ LastModified: new Date('2023-01-01'),
+ },
+ ],
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('test-bucket')
+
+ assert.strictEqual(result.paths.length, 1)
+ const paths = result.paths
+ assert.strictEqual(paths[0].isFolder, false)
+ assert.strictEqual(paths[0].displayName, 'file1.txt')
+ })
+
+ it('should filter out folder markers and prefix matches', async function () {
+ const mockResponse = {
+ CommonPrefixes: [{ Prefix: 'folder/' }],
+ Contents: [
+ {
+ Key: 'prefix/',
+ Size: 0,
+ LastModified: new Date('2023-01-01'),
+ },
+ {
+ Key: 'prefix/file.txt',
+ Size: 1024,
+ LastModified: new Date('2023-01-01'),
+ },
+ {
+ Key: 'prefix/folder/',
+ Size: 0,
+ LastModified: new Date('2023-01-01'),
+ },
+ ],
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('test-bucket', 'prefix/')
+
+ // Should only include the folder from CommonPrefixes and the file (not folder markers)
+ assert.strictEqual(result.paths.length, 2)
+ const paths = result.paths
+ assert.strictEqual(paths[0].displayName, 'folder')
+ assert.strictEqual(paths[0].isFolder, true)
+ assert.strictEqual(paths[1].displayName, 'file.txt')
+ assert.strictEqual(paths[1].isFolder, false)
+ })
+
+ it('should handle API errors', async function () {
+ const error = new Error('S3 API Error')
+ mockS3.listObjectsV2.rejects(error)
+
+ await assert.rejects(async () => await s3Client.listPaths('test-bucket'), error)
+ })
+
+ it('should handle missing object properties gracefully', async function () {
+ const mockResponse = {
+ CommonPrefixes: [{ Prefix: undefined }, { Prefix: 'valid-folder/' }],
+ Contents: [
+ {
+ Key: undefined,
+ Size: 1024,
+ },
+ {
+ Key: 'valid-file.txt',
+ Size: undefined,
+ LastModified: undefined,
+ },
+ ],
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('test-bucket')
+
+ // Should only include valid entries
+ assert.strictEqual(result.paths.length, 2)
+ const paths = result.paths
+ assert.strictEqual(paths[0].displayName, 'valid-folder')
+ assert.strictEqual(paths[0].isFolder, true)
+ assert.strictEqual(paths[1].displayName, 'valid-file.txt')
+ assert.strictEqual(paths[1].isFolder, false)
+ assert.strictEqual(paths[1].size, undefined)
+ assert.strictEqual(paths[1].lastModified, undefined)
+ })
+
+ it('should create S3 client on first use', async function () {
+ const mockResponse = { CommonPrefixes: [], Contents: [] }
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ await s3Client.listPaths('test-bucket')
+
+ // Verify S3 client was created with correct parameters
+ assert.ok(S3.prototype.constructor)
+ })
+
+ it('should reuse existing S3 client on subsequent calls', async function () {
+ const mockResponse = { CommonPrefixes: [], Contents: [] }
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ // Make multiple calls
+ await s3Client.listPaths('test-bucket')
+ await s3Client.listPaths('test-bucket')
+
+ // S3 constructor should only be called once (during first call)
+ assert.ok(mockS3.listObjectsV2.calledTwice)
+ })
+
+ it('should handle ContinuationToken for pagination', async function () {
+ const mockResponse = {
+ CommonPrefixes: [{ Prefix: 'folder1/' }],
+ Contents: [{ Key: 'file1.txt', Size: 1024 }],
+ NextContinuationToken: 'next-token-123',
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('test-bucket', 'prefix/', 'continuation-token')
+
+ assert.strictEqual(result.paths.length, 2)
+ assert.strictEqual(result.nextToken, 'next-token-123')
+
+ // Verify ContinuationToken was passed
+ const callArgs = mockS3.listObjectsV2.getCall(0).args[0]
+ assert.strictEqual(callArgs.ContinuationToken, 'continuation-token')
+ })
+
+ it('should return undefined nextToken when no more pages', async function () {
+ const mockResponse = {
+ CommonPrefixes: [{ Prefix: 'folder1/' }],
+ Contents: [{ Key: 'file1.txt', Size: 1024 }],
+ NextContinuationToken: undefined,
+ }
+
+ mockS3.listObjectsV2.resolves(mockResponse)
+
+ const result = await s3Client.listPaths('test-bucket')
+
+ assert.strictEqual(result.paths.length, 2)
+ assert.strictEqual(result.nextToken, undefined)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/client/sqlWorkbenchClient.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/client/sqlWorkbenchClient.test.ts
new file mode 100644
index 00000000000..e4b1dc50a85
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/client/sqlWorkbenchClient.test.ts
@@ -0,0 +1,249 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import {
+ SQLWorkbenchClient,
+ generateSqlWorkbenchArn,
+ createRedshiftConnectionConfig,
+} from '../../../../sagemakerunifiedstudio/shared/client/sqlWorkbenchClient'
+import { STSClient } from '@aws-sdk/client-sts'
+import globals from '../../../../shared/extensionGlobals'
+import { ConnectionCredentialsProvider } from '../../../../sagemakerunifiedstudio/auth/providers/connectionCredentialsProvider'
+
+describe('SQLWorkbenchClient', function () {
+ let sandbox: sinon.SinonSandbox
+ let mockSqlClient: any
+ let mockSdkClientBuilder: any
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+
+ mockSqlClient = {
+ getResources: sandbox.stub().returns({
+ promise: sandbox.stub().resolves({
+ resources: [{ name: 'test-resource' }],
+ nextToken: 'next-token',
+ }),
+ }),
+ executeQuery: sandbox.stub().returns({
+ promise: sandbox.stub().resolves({
+ queryExecutions: [{ queryExecutionId: 'test-execution-id' }],
+ }),
+ }),
+ }
+
+ mockSdkClientBuilder = {
+ createAwsService: sandbox.stub().resolves(mockSqlClient),
+ }
+
+ sandbox.stub(globals, 'sdkClientBuilder').value(mockSdkClientBuilder)
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ // Reset singleton instance
+ ;(SQLWorkbenchClient as any).instance = undefined
+ })
+
+ describe('getInstance', function () {
+ it('should create singleton instance', function () {
+ const client1 = SQLWorkbenchClient.getInstance('us-east-1')
+ const client2 = SQLWorkbenchClient.getInstance('us-east-1')
+
+ assert.strictEqual(client1, client2)
+ })
+
+ it('should return region correctly', function () {
+ const client = SQLWorkbenchClient.getInstance('us-west-2')
+ assert.strictEqual(client.getRegion(), 'us-west-2')
+ })
+ })
+
+ describe('createWithCredentials', function () {
+ it('should create client with credentials', function () {
+ const credentialsProvider = {
+ getCredentials: async () => ({
+ accessKeyId: 'test-key',
+ secretAccessKey: 'test-secret',
+ sessionToken: 'test-token',
+ }),
+ }
+
+ const client = SQLWorkbenchClient.createWithCredentials(
+ 'us-east-1',
+ credentialsProvider as ConnectionCredentialsProvider
+ )
+ assert.strictEqual(client.getRegion(), 'us-east-1')
+ })
+ })
+
+ describe('getResources', function () {
+ it('should get resources with connection', async function () {
+ const client = SQLWorkbenchClient.getInstance('us-east-1')
+ const connectionConfig = {
+ id: 'test-id',
+ type: 'test-type',
+ databaseType: 'REDSHIFT',
+ connectableResourceIdentifier: 'test-identifier',
+ connectableResourceType: 'CLUSTER',
+ database: 'test-db',
+ }
+
+ const result = await client.getResources({
+ connection: connectionConfig,
+ resourceType: 'TABLE',
+ maxItems: 50,
+ })
+
+ assert.deepStrictEqual(result.resources, [{ name: 'test-resource' }])
+ assert.strictEqual(result.nextToken, 'next-token')
+ })
+
+ it('should handle API errors', async function () {
+ const error = new Error('API Error')
+ mockSqlClient.getResources.returns({
+ promise: sandbox.stub().rejects(error),
+ })
+
+ const client = SQLWorkbenchClient.getInstance('us-east-1')
+
+ await assert.rejects(
+ async () =>
+ await client.getResources({
+ connection: {
+ id: '',
+ type: '',
+ databaseType: '',
+ connectableResourceIdentifier: '',
+ connectableResourceType: '',
+ database: '',
+ },
+ resourceType: '',
+ }),
+ error
+ )
+ })
+ })
+
+ describe('executeQuery', function () {
+ it('should execute query successfully', async function () {
+ const client = SQLWorkbenchClient.getInstance('us-east-1')
+ const connectionConfig = {
+ id: 'test-id',
+ type: 'test-type',
+ databaseType: 'REDSHIFT',
+ connectableResourceIdentifier: 'test-identifier',
+ connectableResourceType: 'CLUSTER',
+ database: 'test-db',
+ }
+
+ const result = await client.executeQuery(connectionConfig, 'SELECT 1')
+
+ assert.strictEqual(result, 'test-execution-id')
+ })
+
+ it('should handle query execution errors', async function () {
+ const error = new Error('Query Error')
+ mockSqlClient.executeQuery.returns({
+ promise: sandbox.stub().rejects(error),
+ })
+
+ const client = SQLWorkbenchClient.getInstance('us-east-1')
+ const connectionConfig = {
+ id: 'test-id',
+ type: 'test-type',
+ databaseType: 'REDSHIFT',
+ connectableResourceIdentifier: 'test-identifier',
+ connectableResourceType: 'CLUSTER',
+ database: 'test-db',
+ }
+
+ await assert.rejects(async () => await client.executeQuery(connectionConfig, 'SELECT 1'), error)
+ })
+ })
+})
+
+describe('generateSqlWorkbenchArn', function () {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ it('should generate ARN with provided account ID', async function () {
+ const arn = await generateSqlWorkbenchArn('us-east-1', '123456789012')
+
+ assert.ok(arn.startsWith('arn:aws:sqlworkbench:us-east-1:123456789012:connection/'))
+ assert.ok(arn.includes('-'))
+ })
+})
+
+describe('createRedshiftConnectionConfig', function () {
+ let sandbox: sinon.SinonSandbox
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+ sandbox.stub(STSClient.prototype, 'send').resolves({ Account: '123456789012' })
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ it('should create serverless connection config', async function () {
+ const config = await createRedshiftConnectionConfig(
+ 'test-workgroup.123456789012.us-east-1.redshift-serverless.amazonaws.com',
+ 'test-db',
+ '123456789012',
+ 'us-east-1',
+ '',
+ false
+ )
+
+ assert.strictEqual(config.databaseType, 'REDSHIFT')
+ assert.strictEqual(config.connectableResourceType, 'WORKGROUP')
+ assert.strictEqual(config.connectableResourceIdentifier, 'test-workgroup')
+ assert.strictEqual(config.database, 'test-db')
+ assert.strictEqual(config.type, '4') // FEDERATED
+ })
+
+ it('should create cluster connection config', async function () {
+ const config = await createRedshiftConnectionConfig(
+ 'test-cluster.123456789012.us-east-1.redshift.amazonaws.com',
+ 'test-db',
+ '123456789012',
+ 'us-east-1',
+ '',
+ false
+ )
+
+ assert.strictEqual(config.databaseType, 'REDSHIFT')
+ assert.strictEqual(config.connectableResourceType, 'CLUSTER')
+ assert.strictEqual(config.connectableResourceIdentifier, 'test-cluster')
+ assert.strictEqual(config.database, 'test-db')
+ assert.strictEqual(config.type, '5') // TEMPORARY_CREDENTIALS_WITH_IAM
+ })
+
+ it('should create config with secret authentication', async function () {
+ const config = await createRedshiftConnectionConfig(
+ 'test-cluster.123456789012.us-east-1.redshift.amazonaws.com',
+ 'test-db',
+ '123456789012',
+ 'us-east-1',
+ 'arn:aws:secretsmanager:us-east-1:123456789012:secret:test-secret',
+ false
+ )
+
+ assert.strictEqual(config.type, '6') // SECRET
+ assert.ok(config.auth)
+ assert.strictEqual(config.auth.secretArn, 'arn:aws:secretsmanager:us-east-1:123456789012:secret:test-secret')
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/smusUtils.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/smusUtils.test.ts
new file mode 100644
index 00000000000..e8f17d486a3
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/smusUtils.test.ts
@@ -0,0 +1,463 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import {
+ SmusUtils,
+ SmusErrorCodes,
+ SmusTimeouts,
+ SmusCredentialExpiry,
+ validateCredentialFields,
+} from '../../../sagemakerunifiedstudio/shared/smusUtils'
+import { ToolkitError } from '../../../shared/errors'
+import * as extensionUtilities from '../../../shared/extensionUtilities'
+import * as resourceMetadataUtils from '../../../sagemakerunifiedstudio/shared/utils/resourceMetadataUtils'
+import fetch from 'node-fetch'
+
+describe('SmusUtils', () => {
+ const testDomainUrl = 'https://dzd_domainId.sagemaker.us-east-2.on.aws'
+ const testDomainIdLowercase = 'dzd_domainid' // Domain IDs get lowercased by URL parsing
+ const testRegion = 'us-east-2'
+
+ afterEach(() => {
+ sinon.restore()
+ })
+
+ describe('extractDomainIdFromUrl', () => {
+ it('should extract domain ID from valid URL', () => {
+ const result = SmusUtils.extractDomainIdFromUrl(testDomainUrl)
+ assert.strictEqual(result, testDomainIdLowercase)
+ })
+
+ it('should return undefined for invalid URL', () => {
+ const result = SmusUtils.extractDomainIdFromUrl('invalid-url')
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should handle URLs with dzd- prefix', () => {
+ const urlWithDash = 'https://dzd-domainId.sagemaker.us-east-2.on.aws'
+ const result = SmusUtils.extractDomainIdFromUrl(urlWithDash)
+ assert.strictEqual(result, 'dzd-domainid')
+ })
+
+ it('should handle URLs with dzd_ prefix', () => {
+ const urlWithUnderscore = 'https://dzd_domainId.sagemaker.us-east-2.on.aws'
+ const result = SmusUtils.extractDomainIdFromUrl(urlWithUnderscore)
+ assert.strictEqual(result, testDomainIdLowercase)
+ })
+ })
+
+ describe('extractRegionFromUrl', () => {
+ it('should extract region from valid URL', () => {
+ const result = SmusUtils.extractRegionFromUrl(testDomainUrl)
+ assert.strictEqual(result, testRegion)
+ })
+
+ it('should return fallback region for invalid URL', () => {
+ const result = SmusUtils.extractRegionFromUrl('invalid-url', 'us-west-2')
+ assert.strictEqual(result, 'us-west-2')
+ })
+
+ it('should return default fallback region when not specified', () => {
+ const result = SmusUtils.extractRegionFromUrl('invalid-url')
+ assert.strictEqual(result, 'us-east-1')
+ })
+
+ it('should handle different regions', () => {
+ const urlWithDifferentRegion = 'https://dzd_test.sagemaker.eu-west-1.on.aws'
+ const result = SmusUtils.extractRegionFromUrl(urlWithDifferentRegion)
+ assert.strictEqual(result, 'eu-west-1')
+ })
+
+ it('should handle non-prod stages', () => {
+ const urlWithStage = 'https://dzd_test.sagemaker-gamma.us-west-2.on.aws'
+ const result = SmusUtils.extractRegionFromUrl(urlWithStage)
+ assert.strictEqual(result, 'us-west-2')
+ })
+ })
+
+ describe('extractDomainInfoFromUrl', () => {
+ it('should extract both domain ID and region', () => {
+ const result = SmusUtils.extractDomainInfoFromUrl(testDomainUrl)
+ assert.strictEqual(result.domainId, testDomainIdLowercase)
+ assert.strictEqual(result.region, testRegion)
+ })
+
+ it('should use fallback region when extraction fails', () => {
+ const result = SmusUtils.extractDomainInfoFromUrl('invalid-url', 'us-west-2')
+ assert.strictEqual(result.domainId, undefined)
+ assert.strictEqual(result.region, 'us-west-2')
+ })
+ })
+
+ describe('validateDomainUrl', () => {
+ it('should return undefined for valid URL', () => {
+ const result = SmusUtils.validateDomainUrl(testDomainUrl)
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should return error for empty URL', () => {
+ const result = SmusUtils.validateDomainUrl('')
+ assert.strictEqual(result, 'Domain URL is required')
+ })
+
+ it('should return error for whitespace-only URL', () => {
+ const result = SmusUtils.validateDomainUrl(' ')
+ assert.strictEqual(result, 'Domain URL is required')
+ })
+
+ it('should return error for non-HTTPS URL', () => {
+ const result = SmusUtils.validateDomainUrl('http://dzd_test.sagemaker.us-east-1.on.aws')
+ assert.strictEqual(result, 'Domain URL must use HTTPS (https://)')
+ })
+
+ it('should return error for non-SageMaker domain', () => {
+ const result = SmusUtils.validateDomainUrl('https://example.com')
+ assert.strictEqual(
+ result,
+ 'URL must be a valid SageMaker Unified Studio domain (e.g., https://dzd_xxxxxxxxx.sagemaker.us-east-1.on.aws)'
+ )
+ })
+
+ it('should return error for URL without domain ID', () => {
+ const result = SmusUtils.validateDomainUrl('https://invalid.sagemaker.us-east-1.on.aws')
+ assert.strictEqual(result, 'URL must contain a valid domain ID (starting with dzd- or dzd_)')
+ })
+
+ it('should return error for invalid URL format', () => {
+ const result = SmusUtils.validateDomainUrl('not-a-url')
+ assert.strictEqual(result, 'Domain URL must use HTTPS (https://)')
+ })
+
+ it('should handle URLs with dzd- prefix', () => {
+ const urlWithDash = 'https://dzd-domainId.sagemaker.us-east-2.on.aws'
+ const result = SmusUtils.validateDomainUrl(urlWithDash)
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should handle URLs with dzd_ prefix', () => {
+ const urlWithUnderscore = 'https://dzd_domainId.sagemaker.us-east-2.on.aws'
+ const result = SmusUtils.validateDomainUrl(urlWithUnderscore)
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should trim whitespace from URL', () => {
+ const urlWithWhitespace = ' https://dzd_domainId.sagemaker.us-east-2.on.aws '
+ const result = SmusUtils.validateDomainUrl(urlWithWhitespace)
+ assert.strictEqual(result, undefined)
+ })
+ })
+
+ describe('constants', () => {
+ it('should export SmusErrorCodes with correct values', () => {
+ assert.strictEqual(SmusErrorCodes.NoActiveConnection, 'NoActiveConnection')
+ assert.strictEqual(SmusErrorCodes.ApiTimeout, 'ApiTimeout')
+ assert.strictEqual(SmusErrorCodes.SmusLoginFailed, 'SmusLoginFailed')
+ assert.strictEqual(SmusErrorCodes.RedeemAccessTokenFailed, 'RedeemAccessTokenFailed')
+ })
+
+ it('should export SmusTimeouts with correct values', () => {
+ assert.strictEqual(SmusTimeouts.apiCallTimeoutMs, 10 * 1000)
+ })
+
+ it('should export SmusCredentialExpiry with correct values', () => {
+ assert.strictEqual(SmusCredentialExpiry.derExpiryMs, 10 * 60 * 1000)
+ assert.strictEqual(SmusCredentialExpiry.projectExpiryMs, 10 * 60 * 1000)
+ assert.strictEqual(SmusCredentialExpiry.connectionExpiryMs, 10 * 60 * 1000)
+ })
+ })
+
+ describe('getSsoInstanceInfo', () => {
+ let fetchStub: sinon.SinonStub
+
+ beforeEach(() => {
+ fetchStub = sinon.stub(fetch, 'default' as any)
+ })
+
+ afterEach(() => {
+ fetchStub.restore()
+ })
+
+ it('should throw error for invalid domain URL', async () => {
+ await assert.rejects(
+ () => SmusUtils.getSsoInstanceInfo('invalid-url'),
+ (error: any) => {
+ assert.strictEqual(error.code, 'InvalidDomainUrl')
+ return true
+ }
+ )
+ })
+
+ it('should throw error for URL without domain ID', async () => {
+ await assert.rejects(
+ () => SmusUtils.getSsoInstanceInfo('https://invalid.sagemaker.us-east-1.on.aws'),
+ (error: any) => {
+ assert.strictEqual(error.code, 'InvalidDomainUrl')
+ return true
+ }
+ )
+ })
+
+ it('should handle timeout errors', async () => {
+ const timeoutError = new Error('Request timeout')
+ timeoutError.name = 'AbortError'
+ fetchStub.rejects(timeoutError)
+
+ await assert.rejects(
+ () => SmusUtils.getSsoInstanceInfo(testDomainUrl),
+ (error: any) => {
+ assert.strictEqual(error.code, SmusErrorCodes.ApiTimeout)
+ assert.ok(error.message.includes('timed out after 10 seconds'))
+ return true
+ }
+ )
+ })
+
+ it('should handle login failure errors', async () => {
+ fetchStub.resolves({
+ ok: false,
+ status: 401,
+ statusText: 'Unauthorized',
+ })
+
+ await assert.rejects(
+ () => SmusUtils.getSsoInstanceInfo(testDomainUrl),
+ (error: any) => {
+ assert.strictEqual(error.code, SmusErrorCodes.SmusLoginFailed)
+ assert.ok(error.message.includes('401'))
+ return true
+ }
+ )
+ })
+
+ it('should successfully extract SSO instance info', async () => {
+ const mockResponse = {
+ ok: true,
+ json: sinon.stub().resolves({
+ redirectUrl:
+ 'https://example.com/oauth/authorize?client_id=arn%3Aaws%3Asso%3A%3A123456789%3Aapplication%2Fssoins-123%2Fapl-456',
+ }),
+ }
+ fetchStub.resolves(mockResponse)
+
+ const result = await SmusUtils.getSsoInstanceInfo(testDomainUrl)
+
+ assert.strictEqual(result.ssoInstanceId, 'ssoins-123')
+ assert.strictEqual(result.issuerUrl, 'https://identitycenter.amazonaws.com/ssoins-123')
+ assert.strictEqual(result.clientId, 'arn:aws:sso::123456789:application/ssoins-123/apl-456')
+ assert.strictEqual(result.region, testRegion)
+ })
+
+ it('should throw error for missing redirect URL', async () => {
+ const mockResponse = {
+ ok: true,
+ json: sinon.stub().resolves({}),
+ }
+ fetchStub.resolves(mockResponse)
+
+ await assert.rejects(
+ () => SmusUtils.getSsoInstanceInfo(testDomainUrl),
+ (error: any) => {
+ assert.strictEqual(error.code, 'InvalidLoginResponse')
+ return true
+ }
+ )
+ })
+
+ it('should throw error for missing client_id in redirect URL', async () => {
+ const mockResponse = {
+ ok: true,
+ json: sinon.stub().resolves({
+ redirectUrl: 'https://example.com/oauth/authorize',
+ }),
+ }
+ fetchStub.resolves(mockResponse)
+
+ await assert.rejects(
+ () => SmusUtils.getSsoInstanceInfo(testDomainUrl),
+ (error: any) => {
+ assert.strictEqual(error.code, 'InvalidRedirectUrl')
+ return true
+ }
+ )
+ })
+
+ it('should throw error for invalid ARN format', async () => {
+ const mockResponse = {
+ ok: true,
+ json: sinon.stub().resolves({
+ redirectUrl: 'https://example.com/oauth/authorize?client_id=invalid-arn',
+ }),
+ }
+ fetchStub.resolves(mockResponse)
+
+ await assert.rejects(
+ () => SmusUtils.getSsoInstanceInfo(testDomainUrl),
+ (error: any) => {
+ assert.strictEqual(error.code, 'InvalidArnFormat')
+ return true
+ }
+ )
+ })
+ })
+
+ describe('extractSSOIdFromUserId', () => {
+ it('should extract SSO ID from valid user ID', () => {
+ const result = SmusUtils.extractSSOIdFromUserId('user-12345678-abcd-efgh-ijkl-123456789012')
+ assert.strictEqual(result, '12345678-abcd-efgh-ijkl-123456789012')
+ })
+
+ it('should throw error for invalid user ID format', () => {
+ assert.throws(
+ () => SmusUtils.extractSSOIdFromUserId('invalid-format'),
+ /Invalid UserId format: invalid-format/
+ )
+ })
+
+ it('should throw error for empty user ID', () => {
+ assert.throws(() => SmusUtils.extractSSOIdFromUserId(''), /Invalid UserId format: /)
+ })
+
+ it('should throw error for user ID without prefix', () => {
+ assert.throws(
+ () => SmusUtils.extractSSOIdFromUserId('12345678-abcd-efgh-ijkl-123456789012'),
+ /Invalid UserId format: 12345678-abcd-efgh-ijkl-123456789012/
+ )
+ })
+ })
+
+ describe('validateCredentialFields', () => {
+ it('should not throw for valid credentials', () => {
+ const validCredentials = {
+ accessKeyId: 'AKIAIOSFODNN7EXAMPLE',
+ secretAccessKey: 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
+ sessionToken:
+ 'AQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKwRcOIfrRh3c/LTo6UDdyJwOOvEVPvLXCrrrUtdnniCEXAMPLE',
+ }
+
+ assert.doesNotThrow(() => {
+ validateCredentialFields(validCredentials, 'TestError', 'test context')
+ })
+ })
+
+ it('should throw for missing accessKeyId', () => {
+ const invalidCredentials = {
+ secretAccessKey: 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
+ sessionToken: 'token',
+ }
+
+ assert.throws(
+ () => validateCredentialFields(invalidCredentials, 'TestError', 'test context'),
+ (error: any) => {
+ assert.ok(error instanceof ToolkitError)
+ assert.strictEqual(error.code, 'TestError')
+ assert.ok(error.message.includes('Invalid accessKeyId in test context'))
+ return true
+ }
+ )
+ })
+
+ it('should throw for invalid accessKeyId type', () => {
+ const invalidCredentials = {
+ accessKeyId: 123,
+ secretAccessKey: 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
+ sessionToken: 'token',
+ }
+
+ assert.throws(
+ () => validateCredentialFields(invalidCredentials, 'TestError', 'test context'),
+ (error: any) => {
+ assert.ok(error instanceof ToolkitError)
+ assert.strictEqual(error.code, 'TestError')
+ assert.ok(error.message.includes('Invalid accessKeyId in test context: number'))
+ return true
+ }
+ )
+ })
+
+ it('should throw for missing secretAccessKey', () => {
+ const invalidCredentials = {
+ accessKeyId: 'AKIAIOSFODNN7EXAMPLE',
+ sessionToken: 'token',
+ }
+
+ assert.throws(
+ () => validateCredentialFields(invalidCredentials, 'TestError', 'test context'),
+ (error: any) => {
+ assert.ok(error instanceof ToolkitError)
+ assert.strictEqual(error.code, 'TestError')
+ assert.ok(error.message.includes('Invalid secretAccessKey in test context'))
+ return true
+ }
+ )
+ })
+
+ it('should throw for missing sessionToken', () => {
+ const invalidCredentials = {
+ accessKeyId: 'AKIAIOSFODNN7EXAMPLE',
+ secretAccessKey: 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
+ }
+
+ assert.throws(
+ () => validateCredentialFields(invalidCredentials, 'TestError', 'test context'),
+ (error: any) => {
+ assert.ok(error instanceof ToolkitError)
+ assert.strictEqual(error.code, 'TestError')
+ assert.ok(error.message.includes('Invalid sessionToken in test context'))
+ return true
+ }
+ )
+ })
+ })
+
+ describe('isInSmusSpaceEnvironment', () => {
+ let isSageMakerStub: sinon.SinonStub
+ let getResourceMetadataStub: sinon.SinonStub
+
+ beforeEach(() => {
+ isSageMakerStub = sinon.stub(extensionUtilities, 'isSageMaker')
+ getResourceMetadataStub = sinon.stub(resourceMetadataUtils, 'getResourceMetadata')
+ })
+
+ it('should return true when in SMUS space with DataZone domain ID', () => {
+ isSageMakerStub.withArgs('SMUS').returns(true)
+ getResourceMetadataStub.returns({
+ AdditionalMetadata: {
+ DataZoneDomainId: 'dz-domain-123',
+ },
+ })
+
+ const result = SmusUtils.isInSmusSpaceEnvironment()
+ assert.strictEqual(result, true)
+ })
+
+ it('should return false when not in SMUS space', () => {
+ isSageMakerStub.withArgs('SMUS').returns(false)
+ isSageMakerStub.withArgs('SMUS-SPACE-REMOTE-ACCESS').returns(false)
+
+ const result = SmusUtils.isInSmusSpaceEnvironment()
+ assert.strictEqual(result, false)
+ })
+
+ it('should return false when in SMUS space but no resource metadata', () => {
+ isSageMakerStub.withArgs('SMUS').returns(true)
+ getResourceMetadataStub.returns(undefined)
+
+ const result = SmusUtils.isInSmusSpaceEnvironment()
+ assert.strictEqual(result, false)
+ })
+
+ it('should return false when in SMUS space but no DataZone domain ID', () => {
+ isSageMakerStub.withArgs('SMUS').returns(true)
+ getResourceMetadataStub.returns({
+ AdditionalMetadata: {},
+ })
+
+ const result = SmusUtils.isInSmusSpaceEnvironment()
+ assert.strictEqual(result, false)
+ })
+ })
+})
diff --git a/packages/core/src/test/sagemakerunifiedstudio/shared/utils/resourceMetadataUtils.test.ts b/packages/core/src/test/sagemakerunifiedstudio/shared/utils/resourceMetadataUtils.test.ts
new file mode 100644
index 00000000000..3580a730fbc
--- /dev/null
+++ b/packages/core/src/test/sagemakerunifiedstudio/shared/utils/resourceMetadataUtils.test.ts
@@ -0,0 +1,292 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import * as sinon from 'sinon'
+import { fs } from '../../../../shared/fs/fs'
+import * as extensionUtilities from '../../../../shared/extensionUtilities'
+import {
+ initializeResourceMetadata,
+ getResourceMetadata,
+ resourceMetadataFileExists,
+ resetResourceMetadata,
+ ResourceMetadata,
+} from '../../../../sagemakerunifiedstudio/shared/utils/resourceMetadataUtils'
+
+describe('resourceMetadataUtils', function () {
+ let sandbox: sinon.SinonSandbox
+
+ const mockMetadata: ResourceMetadata = {
+ AppType: 'JupyterServer',
+ DomainId: 'domain-12345',
+ SpaceName: 'test-space',
+ UserProfileName: 'test-user',
+ ExecutionRoleArn: 'arn:aws:iam::123456789012:role/test-role',
+ ResourceArn: 'arn:aws:sagemaker:us-west-2:123456789012:app/domain-12345/test-user/jupyterserver/test-app',
+ ResourceName: 'test-app',
+ AppImageVersion: '1.0.0',
+ AdditionalMetadata: {
+ DataZoneDomainId: 'dz-domain-123',
+ DataZoneDomainRegion: 'us-west-2',
+ DataZoneEndpoint: 'https://datazone.us-west-2.amazonaws.com',
+ DataZoneEnvironmentId: 'env-123',
+ DataZoneProjectId: 'project-456',
+ DataZoneScopeName: 'test-scope',
+ DataZoneStage: 'prod',
+ DataZoneUserId: 'user-789',
+ PrivateSubnets: 'subnet-123,subnet-456',
+ ProjectS3Path: 's3://test-bucket/project/',
+ SecurityGroup: 'sg-123456789',
+ },
+ ResourceArnCaseSensitive:
+ 'arn:aws:sagemaker:us-west-2:123456789012:app/domain-12345/test-user/JupyterServer/test-app',
+ IpAddressType: 'IPv4',
+ }
+
+ beforeEach(function () {
+ sandbox = sinon.createSandbox()
+ resetResourceMetadata()
+ })
+
+ afterEach(function () {
+ sandbox.restore()
+ })
+
+ describe('initializeResourceMetadata()', function () {
+ it('should initialize metadata when file exists and is valid JSON', async function () {
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(mockMetadata))
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.deepStrictEqual(result, mockMetadata)
+ })
+
+ it('should not initialize when not in SMUS environment', async function () {
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(false)
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should not throw when file does not exist', async function () {
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(false)
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should handle invalid JSON gracefully', async function () {
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves('{ invalid json }')
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should handle file read errors gracefully', async function () {
+ const error = new Error('File read error')
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').rejects(error)
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should handle metadata with missing optional fields', async function () {
+ const minimalMetadata: ResourceMetadata = {
+ DomainId: 'domain-123',
+ }
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(minimalMetadata))
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.deepStrictEqual(result, minimalMetadata)
+ })
+
+ it('should handle metadata with empty AdditionalMetadata', async function () {
+ const metadataWithEmptyAdditional: ResourceMetadata = {
+ DomainId: 'domain-123',
+ AdditionalMetadata: {},
+ }
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(metadataWithEmptyAdditional))
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.deepStrictEqual(result, metadataWithEmptyAdditional)
+ })
+
+ it('should handle empty JSON file', async function () {
+ const emptyMetadata: ResourceMetadata = {}
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(emptyMetadata))
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.deepStrictEqual(result, emptyMetadata)
+ })
+
+ it('should handle very large JSON files', async function () {
+ const largeMetadata = {
+ ...mockMetadata,
+ LargeField: 'x'.repeat(10000),
+ }
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(largeMetadata))
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.strictEqual((result as any).LargeField?.length, 10000)
+ })
+
+ it('should handle JSON with unexpected additional fields', async function () {
+ const metadataWithExtraFields = {
+ ...mockMetadata,
+ UnexpectedField: 'unexpected-value',
+ AdditionalMetadata: {
+ ...mockMetadata.AdditionalMetadata,
+ UnexpectedNestedField: 'unexpected-nested-value',
+ },
+ }
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(metadataWithExtraFields))
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.strictEqual((result as any).UnexpectedField, 'unexpected-value')
+ assert.strictEqual((result as any).AdditionalMetadata?.UnexpectedNestedField, 'unexpected-nested-value')
+ })
+
+ it('should handle JSON with undefined values', async function () {
+ const metadataWithUndefined = {
+ DomainId: undefined,
+ AdditionalMetadata: {
+ DataZoneDomainId: undefined,
+ },
+ }
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(metadataWithUndefined))
+
+ await initializeResourceMetadata()
+ const result = getResourceMetadata()
+
+ assert.strictEqual(result?.DomainId, undefined)
+ assert.strictEqual(result?.AdditionalMetadata?.DataZoneDomainId, undefined)
+ })
+ })
+
+ describe('getResourceMetadata()', function () {
+ it('should return undefined when not initialized', function () {
+ const result = getResourceMetadata()
+ assert.strictEqual(result, undefined)
+ })
+
+ it('should return cached metadata after initialization', async function () {
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(mockMetadata))
+
+ await initializeResourceMetadata()
+
+ const result = getResourceMetadata()
+ assert.deepStrictEqual(result, mockMetadata)
+ })
+
+ it('should return the same instance on multiple calls', async function () {
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ sandbox.stub(fs, 'existsFile').resolves(true)
+ sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(mockMetadata))
+
+ await initializeResourceMetadata()
+
+ const result1 = getResourceMetadata()
+ const result2 = getResourceMetadata()
+
+ assert.strictEqual(result1, result2)
+ assert.deepStrictEqual(result1, mockMetadata)
+ })
+ })
+
+ describe('resetResourceMetadata()', function () {
+ it('should reset cached metadata and allow re-initialization', async function () {
+ sandbox.stub(extensionUtilities, 'isSageMaker').withArgs('SMUS').returns(true)
+ const existsFileStub = sandbox.stub(fs, 'existsFile').resolves(true)
+ const readFileTextStub = sandbox.stub(fs, 'readFileText').resolves(JSON.stringify(mockMetadata))
+
+ await initializeResourceMetadata()
+ const cached1 = getResourceMetadata()
+ assert.deepStrictEqual(cached1, mockMetadata)
+
+ sinon.assert.calledOnce(existsFileStub)
+ sinon.assert.calledOnce(readFileTextStub)
+
+ resetResourceMetadata()
+
+ const cached2 = getResourceMetadata()
+ assert.strictEqual(cached2, undefined)
+
+ await initializeResourceMetadata()
+ const cached3 = getResourceMetadata()
+ assert.deepStrictEqual(cached3, mockMetadata)
+
+ sinon.assert.calledTwice(existsFileStub)
+ sinon.assert.calledTwice(readFileTextStub)
+ })
+ })
+
+ describe('resourceMetadataFileExists()', function () {
+ it('should return true when file exists', async function () {
+ const existsStub = sandbox.stub(fs, 'existsFile').resolves(true)
+
+ const result = await resourceMetadataFileExists()
+
+ assert.strictEqual(result, true)
+ sinon.assert.calledOnceWithExactly(existsStub, '/opt/ml/metadata/resource-metadata.json')
+ })
+
+ it('should return false when file does not exist', async function () {
+ sandbox.stub(fs, 'existsFile').resolves(false)
+
+ const result = await resourceMetadataFileExists()
+
+ assert.strictEqual(result, false)
+ })
+
+ it('should return false and log error when fs.existsFile throws', async function () {
+ const error = new Error('Permission denied')
+ sandbox.stub(fs, 'existsFile').rejects(error)
+
+ const result = await resourceMetadataFileExists()
+
+ assert.strictEqual(result, false)
+ })
+ })
+})
diff --git a/packages/core/src/test/shared/clients/sagemakerClient.test.ts b/packages/core/src/test/shared/clients/sagemakerClient.test.ts
index 888d2222692..ecd60af5ad1 100644
--- a/packages/core/src/test/shared/clients/sagemakerClient.test.ts
+++ b/packages/core/src/test/shared/clients/sagemakerClient.test.ts
@@ -6,9 +6,11 @@
import * as sinon from 'sinon'
import * as assert from 'assert'
import { SagemakerClient } from '../../../shared/clients/sagemaker'
-import { AppDetails, SpaceDetails, DescribeDomainCommandOutput } from '@aws-sdk/client-sagemaker'
+import { AppDetails, SpaceDetails, DescribeDomainCommandOutput, AppType } from '@aws-sdk/client-sagemaker'
import { DescribeDomainResponse } from '@amzn/sagemaker-client'
import { intoCollection } from '../../../shared/utilities/collectionUtils'
+import { ToolkitError } from '../../../shared/errors'
+import { getTestWindow } from '../vscode/window'
describe('SagemakerClient.fetchSpaceAppsAndDomains', function () {
const region = 'test-region'
@@ -91,10 +93,6 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () {
listAppsStub.returns(intoCollection([{ AppName: 'app1', DomainId: 'domain1', SpaceName: 'space1' }]))
const [spaceApps] = await client.fetchSpaceAppsAndDomains()
- for (const space of spaceApps) {
- console.log(space[0])
- console.log(space[1])
- }
const spaceAppKey2 = 'domain2__space2'
const spaceAppKey3 = 'domain2__space3'
@@ -104,121 +102,287 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () {
assert.strictEqual(spaceApps.get(spaceAppKey3)?.App, undefined)
})
- describe('SagemakerClient.startSpace', function () {
- const region = 'test-region'
- let client: SagemakerClient
- let describeSpaceStub: sinon.SinonStub
- let updateSpaceStub: sinon.SinonStub
- let waitForSpaceStub: sinon.SinonStub
- let createAppStub: sinon.SinonStub
-
- beforeEach(function () {
- client = new SagemakerClient(region)
- describeSpaceStub = sinon.stub(client, 'describeSpace')
- updateSpaceStub = sinon.stub(client, 'updateSpace')
- waitForSpaceStub = sinon.stub(client as any, 'waitForSpaceInService')
- createAppStub = sinon.stub(client, 'createApp')
- })
+ it('filters out unified studio domains when filterSmusDomains is true', async function () {
+ const [spaceApps] = await client.fetchSpaceAppsAndDomains(undefined, true)
- afterEach(function () {
- sinon.restore()
- })
+ assert.strictEqual(spaceApps.size, 3)
+ assert.ok(!spaceApps.has('domain3__space4'))
+ })
+
+ it('includes unified studio domains when filterSmusDomains is false', async function () {
+ const [spaceApps] = await client.fetchSpaceAppsAndDomains(undefined, false)
+
+ assert.strictEqual(spaceApps.size, 4)
+ assert.ok(spaceApps.has('domain3__space4'))
+ })
+
+ it('handles AccessDeniedException and shows error message', async function () {
+ sinon.stub(client, 'listSpaceApps').rejects({ name: 'AccessDeniedException' })
+
+ await assert.rejects(client.fetchSpaceAppsAndDomains())
+
+ const messages = getTestWindow().shownMessages
+ assert.ok(messages.some((m) => m.message.includes('AccessDeniedException')))
+ })
+})
+
+describe('SagemakerClient.listSpaceApps', function () {
+ const region = 'test-region'
+ let client: SagemakerClient
+
+ const appDetails: AppDetails[] = [
+ { AppName: 'app1', DomainId: 'domain1', SpaceName: 'space1', AppType: AppType.CodeEditor },
+ { AppName: 'app2', DomainId: 'domain2', SpaceName: 'space2', AppType: AppType.JupyterLab },
+ { AppName: 'app3', DomainId: 'domain2', SpaceName: 'space3', AppType: 'Studio' as any },
+ ]
+
+ const spaceDetails: SpaceDetails[] = [
+ { SpaceName: 'space1', DomainId: 'domain1' },
+ { SpaceName: 'space2', DomainId: 'domain2' },
+ { SpaceName: 'space3', DomainId: 'domain2' },
+ ]
+
+ beforeEach(function () {
+ client = new SagemakerClient(region)
+ sinon.stub(client, 'listApps').returns(intoCollection([appDetails]))
+ sinon.stub(client, 'listSpaces').returns(intoCollection([spaceDetails]))
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ it('returns space apps with correct mapping', async function () {
+ const spaceApps = await client.listSpaceApps()
+
+ assert.strictEqual(spaceApps.size, 3)
+ assert.strictEqual(spaceApps.get('domain1__space1')?.App?.AppName, 'app1')
+ assert.strictEqual(spaceApps.get('domain2__space2')?.App?.AppName, 'app2')
+ assert.strictEqual(spaceApps.get('domain2__space3')?.App, undefined) // Studio app filtered out
+ })
+
+ it('filters by domain when domainId provided', async function () {
+ const newClient = new SagemakerClient(region)
+ const listAppsStub = sinon.stub(newClient, 'listApps').returns(intoCollection([]))
+ const listSpacesStub = sinon.stub(newClient, 'listSpaces').returns(intoCollection([]))
+
+ await newClient.listSpaceApps('domain1')
+
+ sinon.assert.calledWith(listAppsStub, { DomainIdEquals: 'domain1' })
+ sinon.assert.calledWith(listSpacesStub, { DomainIdEquals: 'domain1' })
+ })
+})
+
+describe('SagemakerClient.waitForAppInService', function () {
+ const region = 'test-region'
+ let client: SagemakerClient
+ let describeAppStub: sinon.SinonStub
+
+ beforeEach(function () {
+ client = new SagemakerClient(region)
+ describeAppStub = sinon.stub(client, 'describeApp')
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
- it('enables remote access and starts the app', async function () {
- describeSpaceStub.resolves({
- SpaceSettings: {
- RemoteAccess: 'DISABLED',
- AppType: 'CodeEditor',
- CodeEditorAppSettings: {
- DefaultResourceSpec: {
- InstanceType: 'ml.t3.large',
- SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:img',
- SageMakerImageVersionAlias: '1.0.0',
- },
+ it('resolves when app reaches InService status', async function () {
+ describeAppStub.resolves({ Status: 'InService' })
+
+ await client.waitForAppInService('domain1', 'space1', 'CodeEditor')
+
+ sinon.assert.calledOnce(describeAppStub)
+ })
+
+ it('throws error when app status is Failed', async function () {
+ describeAppStub.resolves({ Status: 'Failed' })
+
+ await assert.rejects(
+ client.waitForAppInService('domain1', 'space1', 'CodeEditor'),
+ /App failed to start. Status: Failed/
+ )
+ })
+
+ it('throws error when app status is DeleteFailed', async function () {
+ describeAppStub.resolves({ Status: 'DeleteFailed' })
+
+ await assert.rejects(
+ client.waitForAppInService('domain1', 'space1', 'CodeEditor'),
+ /App failed to start. Status: DeleteFailed/
+ )
+ })
+
+ it('times out after max retries', async function () {
+ describeAppStub.resolves({ Status: 'Pending' })
+
+ await assert.rejects(
+ client.waitForAppInService('domain1', 'space1', 'CodeEditor', 2, 10),
+ /Timed out waiting for app/
+ )
+ })
+})
+
+describe('SagemakerClient.startSpace', function () {
+ const region = 'test-region'
+ let client: SagemakerClient
+ let describeSpaceStub: sinon.SinonStub
+ let updateSpaceStub: sinon.SinonStub
+ let waitForSpaceStub: sinon.SinonStub
+ let createAppStub: sinon.SinonStub
+
+ beforeEach(function () {
+ client = new SagemakerClient(region)
+ describeSpaceStub = sinon.stub(client, 'describeSpace')
+ updateSpaceStub = sinon.stub(client, 'updateSpace')
+ waitForSpaceStub = sinon.stub(client as any, 'waitForSpaceInService')
+ createAppStub = sinon.stub(client, 'createApp')
+ })
+
+ afterEach(function () {
+ sinon.restore()
+ })
+
+ it('enables remote access and starts the app', async function () {
+ describeSpaceStub.resolves({
+ SpaceSettings: {
+ RemoteAccess: 'DISABLED',
+ AppType: 'CodeEditor',
+ CodeEditorAppSettings: {
+ DefaultResourceSpec: {
+ InstanceType: 'ml.t3.large',
+ SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:img',
+ SageMakerImageVersionAlias: '1.0.0',
},
},
- })
+ },
+ })
- updateSpaceStub.resolves({})
- waitForSpaceStub.resolves()
- createAppStub.resolves({})
+ updateSpaceStub.resolves({})
+ waitForSpaceStub.resolves()
+ createAppStub.resolves({})
- await client.startSpace('my-space', 'my-domain')
+ await client.startSpace('my-space', 'my-domain')
- sinon.assert.calledOnce(updateSpaceStub)
- sinon.assert.calledOnce(waitForSpaceStub)
- sinon.assert.calledOnce(createAppStub)
- })
+ sinon.assert.calledOnce(updateSpaceStub)
+ sinon.assert.calledOnce(waitForSpaceStub)
+ sinon.assert.calledOnce(createAppStub)
+ })
- it('skips enabling remote access if already enabled', async function () {
- describeSpaceStub.resolves({
- SpaceSettings: {
- RemoteAccess: 'ENABLED',
- AppType: 'CodeEditor',
- CodeEditorAppSettings: {
- DefaultResourceSpec: {
- InstanceType: 'ml.t3.large',
- SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:img',
- SageMakerImageVersionAlias: '1.0.0',
- },
+ it('skips enabling remote access if already enabled', async function () {
+ describeSpaceStub.resolves({
+ SpaceSettings: {
+ RemoteAccess: 'ENABLED',
+ AppType: 'CodeEditor',
+ CodeEditorAppSettings: {
+ DefaultResourceSpec: {
+ InstanceType: 'ml.t3.large',
+ SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:img',
+ SageMakerImageVersionAlias: '1.0.0',
},
},
- })
+ },
+ })
- createAppStub.resolves({})
+ createAppStub.resolves({})
- await client.startSpace('my-space', 'my-domain')
+ await client.startSpace('my-space', 'my-domain')
- sinon.assert.notCalled(updateSpaceStub)
- sinon.assert.notCalled(waitForSpaceStub)
- sinon.assert.calledOnce(createAppStub)
+ sinon.assert.notCalled(updateSpaceStub)
+ sinon.assert.notCalled(waitForSpaceStub)
+ sinon.assert.calledOnce(createAppStub)
+ })
+
+ it('throws error on unsupported app type', async function () {
+ describeSpaceStub.resolves({
+ SpaceSettings: {
+ RemoteAccess: 'ENABLED',
+ AppType: 'Studio',
+ },
})
- it('throws error on unsupported app type', async function () {
- describeSpaceStub.resolves({
- SpaceSettings: {
- RemoteAccess: 'ENABLED',
- AppType: 'Studio',
+ await assert.rejects(client.startSpace('my-space', 'my-domain'), /Unsupported AppType "Studio"/)
+ })
+
+ it('uses fallback resource spec when none provided', async function () {
+ describeSpaceStub.resolves({
+ SpaceSettings: {
+ RemoteAccess: 'ENABLED',
+ AppType: 'JupyterLab',
+ JupyterLabAppSettings: {
+ DefaultResourceSpec: {
+ InstanceType: 'ml.t3.large',
+ },
},
+ },
+ })
+
+ createAppStub.resolves({})
+
+ await client.startSpace('my-space', 'my-domain')
+
+ sinon.assert.calledOnceWithExactly(
+ createAppStub,
+ sinon.match.hasNested('ResourceSpec', {
+ InstanceType: 'ml.t3.large',
+ SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:542918446943:image/sagemaker-distribution-cpu',
+ SageMakerImageVersionAlias: '3.2.0',
})
+ )
+ })
- await assert.rejects(client.startSpace('my-space', 'my-domain'), /Unsupported AppType "Studio"/)
- })
+ it('handles AccessDeniedException gracefully', async function () {
+ describeSpaceStub.rejects({ name: 'AccessDeniedException', message: 'no access' })
+
+ await assert.rejects(client.startSpace('my-space', 'my-domain'), /You do not have permission to start spaces/)
+ })
- it('uses fallback resource spec when none provided', async function () {
- describeSpaceStub.resolves({
- SpaceSettings: {
- RemoteAccess: 'ENABLED',
- AppType: 'JupyterLab',
- JupyterLabAppSettings: {
- DefaultResourceSpec: {
- InstanceType: 'ml.t3.large',
- },
+ it('prompts user for insufficient memory instance type', async function () {
+ describeSpaceStub.resolves({
+ SpaceName: 'my-space',
+ SpaceSettings: {
+ RemoteAccess: 'ENABLED',
+ AppType: 'CodeEditor',
+ CodeEditorAppSettings: {
+ DefaultResourceSpec: {
+ InstanceType: 'ml.t3.medium', // Insufficient memory type
},
},
- })
+ },
+ })
- createAppStub.resolves({})
+ createAppStub.resolves({})
- await client.startSpace('my-space', 'my-domain')
+ const promise = client.startSpace('my-space', 'my-domain')
- sinon.assert.calledOnceWithExactly(
- createAppStub,
- sinon.match.hasNested('ResourceSpec', {
- InstanceType: 'ml.t3.large',
- SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:542918446943:image/sagemaker-distribution-cpu',
- SageMakerImageVersionAlias: '3.2.0',
- })
- )
- })
+ // Wait for the error message to appear and select "Yes"
+ await getTestWindow().waitForMessage(/not supported for remote access/)
+ getTestWindow().getFirstMessage().selectItem('Yes')
- it('handles AccessDeniedException gracefully', async function () {
- describeSpaceStub.rejects({ name: 'AccessDeniedException', message: 'no access' })
+ await promise
+ sinon.assert.calledOnce(createAppStub)
+ })
- await assert.rejects(
- client.startSpace('my-space', 'my-domain'),
- /You do not have permission to start spaces/
- )
+ it('throws error when user declines insufficient memory upgrade', async function () {
+ describeSpaceStub.resolves({
+ SpaceName: 'my-space',
+ SpaceSettings: {
+ RemoteAccess: 'ENABLED',
+ AppType: 'CodeEditor',
+ CodeEditorAppSettings: {
+ DefaultResourceSpec: {
+ InstanceType: 'ml.t3.medium',
+ },
+ },
+ },
})
+
+ const promise = client.startSpace('my-space', 'my-domain')
+
+ // Wait for the error message to appear and select "No"
+ await getTestWindow().waitForMessage(/not supported for remote access/)
+ getTestWindow().getFirstMessage().selectItem('No')
+
+ await assert.rejects(promise, (err: ToolkitError) => err.message === 'InstanceType has insufficient memory.')
})
})
diff --git a/packages/toolkit/.changes/next-release/Feature-11b05698-7406-4ccd-afd2-a22d3adbfa67.json b/packages/toolkit/.changes/next-release/Feature-11b05698-7406-4ccd-afd2-a22d3adbfa67.json
new file mode 100644
index 00000000000..ba3abe449fd
--- /dev/null
+++ b/packages/toolkit/.changes/next-release/Feature-11b05698-7406-4ccd-afd2-a22d3adbfa67.json
@@ -0,0 +1,4 @@
+{
+ "type": "Feature",
+ "description": "Feature to support the access of SageMakerUnified Studio resources from the local VSCode IDE"
+}
diff --git a/packages/toolkit/package.json b/packages/toolkit/package.json
index e51656f6ae6..cd5ba5546ce 100644
--- a/packages/toolkit/package.json
+++ b/packages/toolkit/package.json
@@ -779,6 +779,11 @@
"name": "%AWS.codecatalyst.explorerTitle%",
"when": "(!isCloud9 && !aws.isSageMaker || isCloud9CodeCatalyst) && !aws.explorer.showAuthView"
},
+ {
+ "id": "aws.smus.rootView",
+ "name": "%AWS.sagemakerunifiedstudio.explorerTitle%",
+ "when": "!aws.explorer.showAuthView"
+ },
{
"type": "webview",
"id": "aws.toolkit.AmazonCommonAuth",
@@ -1257,6 +1262,14 @@
{
"command": "aws.sagemaker.filterSpaceApps",
"when": "false"
+ },
+ {
+ "command": "aws.smus.switchProject",
+ "when": "false"
+ },
+ {
+ "command": "aws.smus.refreshProject",
+ "when": "false"
}
],
"editor/title": [
@@ -1319,6 +1332,21 @@
}
],
"view/title": [
+ {
+ "command": "aws.smus.switchProject",
+ "when": "view == aws.smus.rootView && !aws.isWebExtHost && aws.smus.connected && !aws.smus.inSmusSpaceEnvironment",
+ "group": "smus@0"
+ },
+ {
+ "command": "aws.smus.refreshProject",
+ "when": "view == aws.smus.rootView && !aws.isWebExtHost && aws.smus.connected",
+ "group": "smus@1"
+ },
+ {
+ "command": "aws.smus.signOut",
+ "when": "view == aws.smus.rootView && !aws.isWebExtHost && aws.smus.connected && !aws.smus.inSmusSpaceEnvironment",
+ "group": "smus@2"
+ },
{
"command": "aws.toolkit.submitFeedback",
"when": "view == aws.explorer && !aws.isWebExtHost",
@@ -1460,18 +1488,38 @@
"command": "aws.stepfunctions.openWithWorkflowStudio",
"when": "isFileSystemResource && resourceFilename =~ /^.*\\.asl\\.(json|yml|yaml)$/",
"group": "z_aws@1"
+ },
+ {
+ "command": "aws.smus.notebookscheduling.createjob",
+ "when": "resourceExtname == .ipynb",
+ "group": "z_aws@1"
+ },
+ {
+ "command": "aws.smus.notebookscheduling.viewjobs",
+ "when": "resourceExtname == .ipynb",
+ "group": "z_aws@1"
}
],
"view/item/context": [
{
"command": "aws.sagemaker.stopSpace",
"group": "inline@0",
- "when": "viewItem =~ /^(awsSagemakerSpaceRunningRemoteEnabledNode|awsSagemakerSpaceRunningRemoteDisabledNode)$/"
+ "when": "view != aws.smus.rootView && viewItem =~ /^(awsSagemakerSpaceRunningRemoteEnabledNode|awsSagemakerSpaceRunningRemoteDisabledNode)$/"
+ },
+ {
+ "command": "aws.smus.stopSpace",
+ "group": "inline@0",
+ "when": "view == aws.smus.rootView && viewItem =~ /^(awsSagemakerSpaceRunningRemoteEnabledNode|awsSagemakerSpaceRunningRemoteDisabledNode|awsSagemakerSpaceRunningNode)$/"
},
{
"command": "aws.sagemaker.openRemoteConnection",
"group": "inline@1",
- "when": "viewItem =~ /^(awsSagemakerSpaceRunningRemoteEnabledNode|awsSagemakerSpaceStoppedRemoteEnabledNode|awsSagemakerSpaceStoppedRemoteDisabledNode)$/"
+ "when": "view != aws.smus.rootView && viewItem =~ /^(awsSagemakerSpaceRunningRemoteEnabledNode|awsSagemakerSpaceStoppedRemoteEnabledNode|awsSagemakerSpaceStoppedRemoteDisabledNode)$/"
+ },
+ {
+ "command": "aws.smus.openRemoteConnection",
+ "group": "inline@1",
+ "when": "view == aws.smus.rootView && viewItem =~ /^(awsSagemakerSpaceRunningRemoteEnabledNode|awsSagemakerSpaceStoppedRemoteEnabledNode|awsSagemakerSpaceStoppedRemoteDisabledNode)$/"
},
{
"command": "_aws.toolkit.notifications.dismiss",
@@ -1633,6 +1681,11 @@
"when": "view == aws.explorer && viewItem == awsSagemakerParentNode",
"group": "inline@1"
},
+ {
+ "command": "aws.smus.refreshProject",
+ "when": "view == aws.smus.rootView && viewItem == smusSelectedProject",
+ "group": "inline@1"
+ },
{
"command": "aws.toolkit.lambda.createServerlessLandProject",
"when": "view == aws.explorer && viewItem == awsLambdaNode || viewItem == awsRegionNode",
@@ -2389,6 +2442,11 @@
]
},
"commands": [
+ {
+ "command": "aws.smus.openSpaceRemoteConnection",
+ "title": "Connect to SageMaker-Unified-Studio Space",
+ "icon": "$(remote-explorer)"
+ },
{
"command": "_aws.toolkit.notifications.dismiss",
"title": "%AWS.generic.dismiss%",
@@ -2642,6 +2700,18 @@
}
}
},
+ {
+ "command": "aws.smus.openRemoteConnection",
+ "title": "Connect to SageMaker Space",
+ "icon": "$(remote-explorer)",
+ "category": "%AWS.title%",
+ "enablement": "isCloud9 || !aws.isWebExtHost",
+ "cloud9": {
+ "cn": {
+ "category": "%AWS.title.cn%"
+ }
+ }
+ },
{
"command": "aws.sagemaker.stopSpace",
"title": "Stop SageMaker Space",
@@ -2654,6 +2724,56 @@
}
}
},
+ {
+ "command": "aws.smus.stopSpace",
+ "title": "Stop SageMaker Space",
+ "icon": "$(debug-stop)",
+ "category": "%AWS.title%",
+ "enablement": "isCloud9 || !aws.isWebExtHost",
+ "cloud9": {
+ "cn": {
+ "category": "%AWS.title.cn%"
+ }
+ }
+ },
+ {
+ "command": "aws.smus.switchProject",
+ "title": "%AWS.command.smus.switchProject%",
+ "category": "%AWS.title%",
+ "enablement": "isCloud9 || !aws.isWebExtHost",
+ "cloud9": {
+ "cn": {
+ "category": "%AWS.title.cn%"
+ }
+ }
+ },
+ {
+ "command": "aws.smus.refreshProject",
+ "title": "%AWS.command.smus.refreshProject%",
+ "category": "%AWS.title%",
+ "enablement": "isCloud9 || !aws.isWebExtHost",
+ "icon": {
+ "dark": "resources/icons/vscode/dark/refresh.svg",
+ "light": "resources/icons/vscode/light/refresh.svg"
+ },
+ "cloud9": {
+ "cn": {
+ "category": "%AWS.title.cn%"
+ }
+ }
+ },
+ {
+ "command": "aws.smus.signOut",
+ "title": "%AWS.command.smus.signOut%",
+ "category": "%AWS.title%",
+ "enablement": "isCloud9 || !aws.isWebExtHost",
+ "icon": "$(sign-out)",
+ "cloud9": {
+ "cn": {
+ "category": "%AWS.title.cn%"
+ }
+ }
+ },
{
"command": "aws.ec2.startInstance",
"title": "%AWS.command.ec2.startInstance%",
@@ -4283,6 +4403,16 @@
"category": "%AWS.title.cn%"
}
}
+ },
+ {
+ "command": "aws.smus.notebookscheduling.createjob",
+ "title": "Create Notebook Job",
+ "category": "Job"
+ },
+ {
+ "command": "aws.smus.notebookscheduling.viewjobs",
+ "title": "View Notebook Jobs",
+ "category": "Job"
}
],
"jsonValidation": [
@@ -4823,26 +4953,61 @@
"fontCharacter": "\\f1e0"
}
},
- "aws-schemas-registry": {
+ "aws-sagemakerunifiedstudio-catalog": {
"description": "AWS Contributed Icon",
"default": {
"fontPath": "./resources/fonts/aws-toolkit-icons.woff",
"fontCharacter": "\\f1e1"
}
},
- "aws-schemas-schema": {
+ "aws-sagemakerunifiedstudio-spaces": {
"description": "AWS Contributed Icon",
"default": {
"fontPath": "./resources/fonts/aws-toolkit-icons.woff",
"fontCharacter": "\\f1e2"
}
},
- "aws-stepfunctions-preview": {
+ "aws-sagemakerunifiedstudio-spaces-dark": {
"description": "AWS Contributed Icon",
"default": {
"fontPath": "./resources/fonts/aws-toolkit-icons.woff",
"fontCharacter": "\\f1e3"
}
+ },
+ "aws-sagemakerunifiedstudio-symbol-int": {
+ "description": "AWS Contributed Icon",
+ "default": {
+ "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+ "fontCharacter": "\\f1e4"
+ }
+ },
+ "aws-sagemakerunifiedstudio-table": {
+ "description": "AWS Contributed Icon",
+ "default": {
+ "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+ "fontCharacter": "\\f1e5"
+ }
+ },
+ "aws-schemas-registry": {
+ "description": "AWS Contributed Icon",
+ "default": {
+ "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+ "fontCharacter": "\\f1e6"
+ }
+ },
+ "aws-schemas-schema": {
+ "description": "AWS Contributed Icon",
+ "default": {
+ "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+ "fontCharacter": "\\f1e7"
+ }
+ },
+ "aws-stepfunctions-preview": {
+ "description": "AWS Contributed Icon",
+ "default": {
+ "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+ "fontCharacter": "\\f1e8"
+ }
}
},
"notebooks": [