address the comments and add unit tests

Author: PotatoWKY

packages/core/src/awsService/sagemaker/sagemakerSpace.ts

@@ -11,6 +11,9 @@ import { getIcon, IconPath } from '../../shared/icons'
 import { generateSpaceStatus, updateIdleFile, startMonitoringTerminalActivity, ActivityCheckInterval } from './utils'
 import { UserActivity } from '../../shared/extensionUtilities'
 import { getLogger } from '../../shared/logger/logger'
+import { ToolkitError } from '../../shared/errors'
+
+const logger = getLogger('sagemaker')
 
 export class SagemakerSpace {
     public label: string = ''
@@ -109,6 +112,8 @@ export class SagemakerSpace {
             DomainId: this.spaceApp.DomainId,
             SpaceName: this.spaceApp.SpaceName,
         })
+        // If the App is not null and AppName exists -> use DescribeApp API (requires AppName as a parameter)
+        // otherwise -> use ListApp API, with given DomainId and SpaceName
         let app
         if (this.spaceApp.App?.AppName) {
             app = await this.client.describeApp({
@@ -118,7 +123,14 @@ export class SagemakerSpace {
                 SpaceName: this.spaceApp.SpaceName,
             })
         } else if (this.spaceApp.DomainId && this.spaceApp.SpaceName) {
-            app = await this.client.getAppBySpace(this.spaceApp.DomainId, this.spaceApp.SpaceName)
+            app = await this.client.listAppForSpace(this.spaceApp.DomainId, this.spaceApp.SpaceName)
+        } else {
+            logger.error(
+                `updateSpaceAppStatus: unable to get app, [AppName: ${this.spaceApp.App?.AppName}], [DomainId: ${this.spaceApp.DomainId}], [SpaceName: ${this.spaceApp.SpaceName}]`
+            )
+            throw new ToolkitError(
+                `Cannot update app status without [AppName: ${this.spaceApp.App?.AppName}] or [DomainId: ${this.spaceApp.DomainId} and SpaceName: ${this.spaceApp.SpaceName}]`
+            )
         }
 
         // AWS DescribeSpace API returns full details with property names like 'SpaceSettings'
@@ -199,7 +211,6 @@ export class SagemakerSpace {
  * Sets up user activity monitoring for SageMaker spaces
  */
 export async function setupUserActivityMonitoring(extensionContext: vscode.ExtensionContext): Promise<void> {
-    const logger = getLogger()
     logger.info('setupUserActivityMonitoring: Starting user activity monitoring setup')
 
     const tmpDirectory = '/tmp/'

packages/core/src/shared/clients/sagemaker.ts

@@ -126,11 +126,11 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
         return this.makeRequest(DeleteAppCommand, request)
     }
 
-    public async getAppBySpace(domainId: string, spaceName: string): Promise<AppDetails | undefined> {
+    public async listAppForSpace(domainId: string, spaceName: string): Promise<AppDetails | undefined> {
         const appsList = await this.listApps({ DomainIdEquals: domainId, SpaceNameEquals: spaceName })
             .flatten()
             .promise()
-        return appsList[0]
+        return appsList[0] // At most one App for one SagemakerSpace
     }
 
     public async startSpace(spaceName: string, domainId: string) {

packages/core/src/test/awsService/sagemaker/sagemakerSpace.test.ts

@@ -63,7 +63,7 @@ describe('SagemakerSpace', function () {
 
             mockClient.describeSpace.resolves(mockDescribeSpaceResponse)
             mockClient.describeApp.resolves(mockDescribeAppResponse)
-            mockClient.getAppBySpace.resolves(mockDescribeAppResponse)
+            mockClient.listAppForSpace.resolves(mockDescribeAppResponse)
 
             const space = new SagemakerSpace(mockClient as any, 'us-east-1', mockSpaceApp)
             const updateSpaceSpy = sinon.spy(space, 'updateSpace')
@@ -126,5 +126,37 @@ describe('SagemakerSpace', function () {
             assert.strictEqual(updateSpaceArgs.OwnershipSettingsSummary, undefined)
             assert.strictEqual(updateSpaceArgs.SpaceSharingSettingsSummary, undefined)
         })
+
+        it('should use listAppForSpace when AppName is not available', async function () {
+            const mockDescribeSpaceResponse = {
+                SpaceName: 'test-space',
+                Status: 'InService',
+                DomainId: 'test-domain',
+                $metadata: { requestId: 'test-request-id' },
+            }
+
+            const mockAppFromList = {
+                AppName: 'listed-app',
+                Status: 'InService',
+                $metadata: { requestId: 'test-request-id' },
+            }
+
+            mockClient.describeSpace.resolves(mockDescribeSpaceResponse)
+            mockClient.listAppForSpace.resolves(mockAppFromList)
+
+            // Create space without App.AppName
+            const spaceWithoutAppName: SagemakerSpaceApp = {
+                ...mockSpaceApp,
+                App: undefined,
+            }
+
+            const space = new SagemakerSpace(mockClient as any, 'us-east-1', spaceWithoutAppName)
+            await space.updateSpaceAppStatus()
+
+            // Verify listAppForSpace was called instead of describeApp
+            assert.ok(mockClient.listAppForSpace.calledOnce)
+            assert.ok(mockClient.listAppForSpace.calledWith('test-domain', 'test-space'))
+            assert.ok(mockClient.describeApp.notCalled)
+        })
     })
 })

packages/core/src/test/shared/clients/sagemakerClient.test.ts

@@ -173,6 +173,41 @@ describe('SagemakerClient.listSpaceApps', function () {
     })
 })
 
+describe('SagemakerClient.listAppForSpace', function () {
+    const region = 'test-region'
+    let client: SagemakerClient
+    let listAppsStub: sinon.SinonStub
+
+    beforeEach(function () {
+        client = new SagemakerClient(region)
+        listAppsStub = sinon.stub(client, 'listApps')
+    })
+
+    afterEach(function () {
+        sinon.restore()
+    })
+
+    it('returns first app for given domain and space', async function () {
+        const appDetails: AppDetails[] = [
+            { AppName: 'app1', DomainId: 'domain1', SpaceName: 'space1', AppType: AppType.CodeEditor },
+        ]
+        listAppsStub.returns(intoCollection([appDetails]))
+
+        const result = await client.listAppForSpace('domain1', 'space1')
+
+        assert.strictEqual(result?.AppName, 'app1')
+        sinon.assert.calledWith(listAppsStub, { DomainIdEquals: 'domain1', SpaceNameEquals: 'space1' })
+    })
+
+    it('returns undefined when no apps found', async function () {
+        listAppsStub.returns(intoCollection([[]]))
+
+        const result = await client.listAppForSpace('domain1', 'space1')
+
+        assert.strictEqual(result, undefined)
+    })
+})
+
 describe('SagemakerClient.waitForAppInService', function () {
     const region = 'test-region'
     let client: SagemakerClient