Skip to content

Commit cee9380

Browse files
parameja1parameja1
committed
update isSageMaker logic to use env vars properly and add unit tests
1 parent 022c81c commit cee9380

File tree

4 files changed

+217
-6
lines changed

4 files changed

+217
-6
lines changed

‎packages/core/src/shared/extensionUtilities.ts‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,15 @@ export function isSageMakerUnifiedStudio(): boolean {
158158
return isSMUS
159159
}
160160

161+
/**
162+
* Reset cached SageMaker state - for testing purposes only
163+
* @internal
164+
*/
165+
export function resetSageMakerState(): void {
166+
serviceName = notInitialized
167+
isSMUS = false
168+
}
169+
161170
/**
162171
* Decides if the current system is (the specified flavor of) Cloud9.
163172
*/
@@ -177,17 +186,17 @@ export function isCloud9(flavor: 'classic' | 'codecatalyst' | 'any' = 'any'): bo
177186
*/
178187
export function isSageMaker(appName: 'SMAI' | 'SMUS' = 'SMAI'): boolean {
179188
// Check for SageMaker-specific environment variables first
189+
let hasSMEnvVars: boolean = false
180190
if (hasSageMakerEnvVars()) {
181191
getLogger().debug('SageMaker environment detected via environment variables')
182-
return true
192+
hasSMEnvVars = true
183193
}
184194

185-
// Fall back to app name checks
186195
switch (appName) {
187196
case 'SMAI':
188-
return vscode.env.appName === sageMakerAppname
197+
return vscode.env.appName === sageMakerAppname && hasSMEnvVars
189198
case 'SMUS':
190-
return vscode.env.appName === sageMakerAppname && isSageMakerUnifiedStudio()
199+
return vscode.env.appName === sageMakerAppname && isSageMakerUnifiedStudio() && hasSMEnvVars
191200
default:
192201
return false
193202
}

‎packages/core/src/shared/telemetry/util.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import { mapMetadata, MetadataObj } from './telemetryLogger'
2323
import { Result } from './telemetry.gen'
2424
import { MetricDatum } from './clienttelemetry'
2525
import { isValidationExemptMetric } from './exemptMetrics'
26-
import { isAmazonQ, isCloud9, isSageMaker, isSageMakerUnifiedStudio } from '../../shared/extensionUtilities'
26+
import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities'
2727
import { isUuid, randomUUID } from '../crypto'
2828
import { ClassToInterfaceType } from '../utilities/tsUtils'
2929
import { asStringifiedStack, FunctionEntry } from './spans'
@@ -488,7 +488,7 @@ export function withTelemetryContext(opts: TelemetryContextArgs) {
488488
* Returns default value of vscode appName or AmazonQ-For-SMUS-CE in case of a sagemaker unified studio environment
489489
*/
490490
export function getClientName(): string {
491-
if (isSageMaker() && isSageMakerUnifiedStudio()) {
491+
if (isSageMaker('SMUS')) {
492492
return 'AmazonQ-For-SMUS-CE'
493493
}
494494
return env.appName

‎packages/core/src/test/shared/extensionUtilities.test.ts‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ import globals from '../../shared/extensionGlobals'
1818
import { maybeShowMinVscodeWarning } from '../../shared/extensionStartup'
1919
import { getTestWindow } from './vscode/window'
2020
import { assertTelemetry } from '../testUtil'
21+
import { isSageMaker } from '../../shared/extensionUtilities'
22+
import { hasSageMakerEnvVars } from '../../shared/vscode/env'
2123

2224
describe('extensionUtilities', function () {
2325
it('maybeShowMinVscodeWarning', async () => {
@@ -361,3 +363,146 @@ describe('UserActivity', function () {
361363
return event.event
362364
}
363365
})
366+
367+
describe('isSageMaker', function () {
368+
let sandbox: sinon.SinonSandbox
369+
const env = require('../../shared/vscode/env')
370+
const utils = require('../../shared/extensionUtilities')
371+
372+
beforeEach(function () {
373+
sandbox = sinon.createSandbox()
374+
utils.resetSageMakerState()
375+
})
376+
377+
afterEach(function () {
378+
sandbox.restore()
379+
delete process.env.SERVICE_NAME
380+
})
381+
382+
describe('SMAI detection', function () {
383+
it('returns true when both app name and env vars match', function () {
384+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
385+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(true)
386+
387+
assert.strictEqual(isSageMaker('SMAI'), true)
388+
})
389+
390+
it('returns false when app name is different', function () {
391+
sandbox.stub(vscode.env, 'appName').value('Visual Studio Code')
392+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(true)
393+
394+
assert.strictEqual(isSageMaker('SMAI'), false)
395+
})
396+
397+
it('returns false when env vars are missing', function () {
398+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
399+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(false)
400+
401+
assert.strictEqual(isSageMaker('SMAI'), false)
402+
})
403+
404+
it('defaults to SMAI when no parameter provided', function () {
405+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
406+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(true)
407+
408+
assert.strictEqual(isSageMaker(), true)
409+
})
410+
})
411+
412+
describe('SMUS detection', function () {
413+
it('returns true when all conditions are met', function () {
414+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
415+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(true)
416+
process.env.SERVICE_NAME = 'SageMakerUnifiedStudio'
417+
418+
assert.strictEqual(isSageMaker('SMUS'), true)
419+
})
420+
421+
it('returns false when unified studio is missing', function () {
422+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
423+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(true)
424+
process.env.SERVICE_NAME = 'SomeOtherService'
425+
426+
assert.strictEqual(isSageMaker('SMUS'), false)
427+
})
428+
429+
it('returns false when env vars are missing', function () {
430+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
431+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(false)
432+
process.env.SERVICE_NAME = 'SageMakerUnifiedStudio'
433+
434+
assert.strictEqual(isSageMaker('SMUS'), false)
435+
})
436+
437+
it('returns false when app name is different', function () {
438+
sandbox.stub(vscode.env, 'appName').value('Visual Studio Code')
439+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(true)
440+
process.env.SERVICE_NAME = 'SageMakerUnifiedStudio'
441+
442+
assert.strictEqual(isSageMaker('SMUS'), false)
443+
})
444+
})
445+
446+
it('returns false for invalid appName parameter', function () {
447+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
448+
sandbox.stub(env, 'hasSageMakerEnvVars').returns(true)
449+
450+
// @ts-ignore - Testing invalid input
451+
assert.strictEqual(isSageMaker('INVALID'), false)
452+
})
453+
})
454+
455+
describe('hasSageMakerEnvVars', function () {
456+
let originalEnv: NodeJS.ProcessEnv
457+
458+
beforeEach(function () {
459+
originalEnv = { ...process.env }
460+
// Clear all SageMaker-related env vars
461+
delete process.env.SAGEMAKER_APP_TYPE
462+
delete process.env.SAGEMAKER_INTERNAL_IMAGE_URI
463+
delete process.env.STUDIO_LOGGING_DIR
464+
delete process.env.SM_APP_TYPE
465+
delete process.env.SM_INTERNAL_IMAGE_URI
466+
delete process.env.SERVICE_NAME
467+
})
468+
469+
afterEach(function () {
470+
process.env = originalEnv
471+
})
472+
473+
const testCases = [
474+
{ env: 'SAGEMAKER_APP_TYPE', value: 'JupyterServer', expected: true },
475+
{ env: 'SAGEMAKER_INTERNAL_IMAGE_URI', value: 'some-uri', expected: true },
476+
{ env: 'STUDIO_LOGGING_DIR', value: '/var/log/studio/app.log', expected: true },
477+
{ env: 'STUDIO_LOGGING_DIR', value: '/var/log/other/app.log', expected: false },
478+
{ env: 'SM_APP_TYPE', value: 'JupyterServer', expected: true },
479+
{ env: 'SM_INTERNAL_IMAGE_URI', value: 'some-uri', expected: true },
480+
{ env: 'SERVICE_NAME', value: 'SageMakerUnifiedStudio', expected: true },
481+
{ env: 'SERVICE_NAME', value: 'SomeOtherService', expected: false },
482+
]
483+
484+
for (const { env, value, expected } of testCases) {
485+
it(`returns ${expected} when ${env} is set to "${value}"`, function () {
486+
process.env[env] = value
487+
488+
const result = hasSageMakerEnvVars()
489+
490+
assert.strictEqual(result, expected)
491+
})
492+
}
493+
494+
it('returns true when multiple SageMaker env vars are set', function () {
495+
process.env.SAGEMAKER_APP_TYPE = 'JupyterServer'
496+
process.env.SM_APP_TYPE = 'CodeEditor'
497+
498+
const result = hasSageMakerEnvVars()
499+
500+
assert.strictEqual(result, true)
501+
})
502+
503+
it('returns false when no SageMaker env vars are set', function () {
504+
const result = hasSageMakerEnvVars()
505+
506+
assert.strictEqual(result, false)
507+
})
508+
})

‎packages/core/src/test/shared/telemetry/util.test.ts‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,10 @@ import { randomUUID } from 'crypto'
2424
import { isUuid } from '../../../shared/crypto'
2525
import { MetricDatum } from '../../../shared/telemetry/clienttelemetry'
2626
import { assertLogsContain } from '../../globalSetup.test'
27+
import { getClientName } from '../../../shared/telemetry/util'
28+
import * as extensionUtilities from '../../../shared/extensionUtilities'
29+
import * as sinon from 'sinon'
30+
import * as vscode from 'vscode'
2731

2832
describe('TelemetryConfig', function () {
2933
const settingKey = 'aws.telemetry'
@@ -391,3 +395,56 @@ describe('validateMetricEvent', function () {
391395
assertLogsContain('invalid Metric', false, 'warn')
392396
})
393397
})
398+
399+
describe('getClientName', function () {
400+
let sandbox: sinon.SinonSandbox
401+
let isSageMakerStub: sinon.SinonStub
402+
403+
beforeEach(function () {
404+
sandbox = sinon.createSandbox()
405+
isSageMakerStub = sandbox.stub(extensionUtilities, 'isSageMaker')
406+
})
407+
408+
afterEach(function () {
409+
sandbox.restore()
410+
})
411+
412+
it('returns "AmazonQ-For-SMUS-CE" when in SMUS environment', function () {
413+
isSageMakerStub.withArgs('SMUS').returns(true)
414+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
415+
416+
const result = getClientName()
417+
418+
assert.strictEqual(result, 'AmazonQ-For-SMUS-CE')
419+
assert.ok(isSageMakerStub.calledOnceWith('SMUS'))
420+
})
421+
422+
it('returns vscode app name when not in SMUS environment', function () {
423+
const mockAppName = 'Visual Studio Code'
424+
isSageMakerStub.withArgs('SMUS').returns(false)
425+
sandbox.stub(vscode.env, 'appName').value(mockAppName)
426+
427+
const result = getClientName()
428+
429+
assert.strictEqual(result, mockAppName)
430+
assert.ok(isSageMakerStub.calledOnceWith('SMUS'))
431+
})
432+
433+
it('handles undefined app name gracefully', function () {
434+
isSageMakerStub.withArgs('SMUS').returns(false)
435+
sandbox.stub(vscode.env, 'appName').value(undefined)
436+
437+
const result = getClientName()
438+
439+
assert.strictEqual(result, undefined)
440+
})
441+
442+
it('prioritizes SMUS detection over app name', function () {
443+
isSageMakerStub.withArgs('SMUS').returns(true)
444+
sandbox.stub(vscode.env, 'appName').value('SageMaker Code Editor')
445+
446+
const result = getClientName()
447+
448+
assert.strictEqual(result, 'AmazonQ-For-SMUS-CE')
449+
})
450+
})

0 commit comments

Comments
 (0)