fix(smus): Look for profile in both config and credentials files when adding region to profile

packages/core/src/sagemakerunifiedstudio/auth/ui/iamProfileSelection.ts

@@ -5,9 +5,10 @@
 
 import * as vscode from 'vscode'
 import * as path from 'path'
+import * as os from 'os'
 import { getLogger } from '../../../shared/logger/logger'
 import { ToolkitError } from '../../../shared/errors'
-import { loadSharedCredentialsProfiles } from '../../../auth/credentials/sharedCredentials'
+import { loadSharedCredentialsProfiles, parseIni } from '../../../auth/credentials/sharedCredentials'
 import { getCredentialsFilename, getConfigFilename } from '../../../auth/credentials/sharedCredentialsFile'
 import { SmusErrorCodes, DataZoneServiceId } from '../../shared/smusUtils'
 import globals from '../../../shared/extensionGlobals'
@@ -1189,7 +1190,7 @@ export class SmusIamProfileSelector {
         }
 
         // Parse the file line by line to handle profile replacement properly
-        const lines = content.split('\n')
+        const lines = content.split(os.EOL)
         const newLines: string[] = []
         let inTargetProfile = false
         let profileFound = false
@@ -1230,7 +1231,7 @@ export class SmusIamProfileSelector {
         }
 
         // Update content with the new lines
-        content = newLines.join('\n')
+        content = newLines.join(os.EOL)
 
         // Write back to file
         await fs.writeFile(credentialsPath, content)
@@ -1245,62 +1246,85 @@ export class SmusIamProfileSelector {
         try {
             logger.debug(`Updating profile ${profileName} with region ${region}`)
 
-            const credentialsPath = getCredentialsFilename()
+            // Check both config and credential files
+            const filepathsToCheck = [getCredentialsFilename(), getConfigFilename()]
 
-            if (!(await fs.existsFile(credentialsPath))) {
-                throw new ToolkitError('Credentials file not found', { code: 'CredentialsFileNotFound' })
-            }
+            let profileUpdated = false
 
-            // Read the current credentials file
-            const content = await fs.readFileText(credentialsPath)
+            for (const filePath of filepathsToCheck) {
+                // File does not exist, try next file
+                if (!(await fs.existsFile(filePath))) {
+                    continue
+                }
 
-            // Find the profile section
-            const profileSectionRegex = new RegExp(`^\\[${profileName}\\]$`, 'm')
-            const profileMatch = content.match(profileSectionRegex)
+                const content = await fs.readFileText(filePath)
+                const sections = parseIni(content, vscode.Uri.file(filePath))
 
-            if (!profileMatch) {
-                throw new ToolkitError(`Profile ${profileName} not found in credentials file`, {
-                    code: 'ProfileNotFound',
-                })
-            }
+                // Find the profile section in this file
+                const profileSection = sections.find(
+                    (section) => section.type === 'profile' && section.name === profileName
+                )
 
-            // Find the next profile section or end of file
-            const profileStartIndex = profileMatch.index!
-            const nextProfileMatch = content.slice(profileStartIndex + 1).match(/^\[.*\]$/m)
-            const profileEndIndex = nextProfileMatch ? profileStartIndex + 1 + nextProfileMatch.index! : content.length
-
-            // Extract the profile section
-            const profileSection = content.slice(profileStartIndex, profileEndIndex)
-
-            // Check if region already exists in the profile
-            let updatedProfileSection: string
-
-            if (this.regionLinePattern.test(profileSection)) {
-                // Replace existing region
-                updatedProfileSection = profileSection.replace(this.regionLinePattern, `region = ${region}`)
-            } else {
-                // Add region to the profile (before any empty lines at the end)
-                const lines = profileSection.split('\n')
-                // Find the last non-empty line index (compatible with older JS versions)
-                let lastNonEmptyIndex = -1
-                for (let i = lines.length - 1; i >= 0; i--) {
-                    if (lines[i].trim() !== '') {
-                        lastNonEmptyIndex = i
+                // Profile not in this file, try next file
+                if (!profileSection) {
+                    continue
+                }
+
+                // Find the profile section boundaries using the startLines from parsed section
+                const profileStartLine = profileSection.startLines[0]
+                const lines = content.split(os.EOL)
+
+                // Find the next profile section or end of file
+                let profileEndLine = lines.length
+                for (let i = profileStartLine + 1; i < lines.length; i++) {
+                    if (lines[i].match(/^\s*\[([^\[\]]+)]\s*$/)) {
+                        profileEndLine = i
                         break
                     }
                 }
-                lines.splice(lastNonEmptyIndex + 1, 0, `region = ${region}`)
-                updatedProfileSection = lines.join('\n')
-            }
 
-            // Replace the profile section in the content
-            const updatedContent =
-                content.slice(0, profileStartIndex) + updatedProfileSection + content.slice(profileEndIndex)
+                // Extract the profile section lines
+                const profileLines = lines.slice(profileStartLine, profileEndLine)
+
+                // Check if region already exists in the profile
+                const regionLineIndex = profileLines.findIndex((line) => this.regionLinePattern.test(line))
 
-            // Write back to file
-            await fs.writeFile(credentialsPath, updatedContent)
+                if (regionLineIndex !== -1) {
+                    // Replace existing region
+                    profileLines[regionLineIndex] = `region = ${region}`
+                } else {
+                    // Add region to the profile (after the last non-empty line)
+                    let lastNonEmptyIndex = -1
+                    for (let i = profileLines.length - 1; i >= 0; i--) {
+                        if (profileLines[i].trim() !== '') {
+                            lastNonEmptyIndex = i
+                            break
+                        }
+                    }
+                    profileLines.splice(lastNonEmptyIndex + 1, 0, `region = ${region}`)
+                }
+
+                // Reconstruct the file content
+                const updatedLines = [
+                    ...lines.slice(0, profileStartLine),
+                    ...profileLines,
+                    ...lines.slice(profileEndLine),
+                ]
+                const updatedContent = updatedLines.join(os.EOL)
 
-            logger.debug(`Successfully updated profile ${profileName} with region ${region}`)
+                // Write back to file
+                await fs.writeFile(filePath, updatedContent)
+
+                logger.debug(`Successfully updated profile ${profileName} with region ${region} in ${filePath}`)
+                profileUpdated = true
+                break
+            }
+
+            if (!profileUpdated) {
+                throw new ToolkitError(`Profile ${profileName} not found in credentials or config file`, {
+                    code: 'ProfileNotFound',
+                })
+            }
         } catch (error) {
             logger.error('Failed to update profile region: %s', error)
             throw new ToolkitError(`Failed to update profile region: ${(error as Error).message}`, {

packages/core/src/test/sagemakerunifiedstudio/auth/ui/iamProfileSelection.test.ts

@@ -4,7 +4,13 @@
  */
 
 import assert from 'assert'
+import * as path from 'path'
+import * as os from 'os'
+import * as sinon from 'sinon'
 import { SmusIamProfileSelector } from '../../../../sagemakerunifiedstudio/auth/ui/iamProfileSelection'
+import { makeTemporaryToolkitFolder } from '../../../../shared/filesystemUtilities'
+import { fs } from '../../../../shared'
+import { EnvironmentVariables } from '../../../../shared/environmentVariables'
 
 describe('SmusIamProfileSelector', function () {
     describe('showRegionSelection', function () {
@@ -18,4 +24,134 @@ describe('SmusIamProfileSelector', function () {
             assert.strictEqual(typeof SmusIamProfileSelector.showIamProfileSelection, 'function')
         })
     })
+
+    describe('updateProfileRegion', function () {
+        let tempFolder: string
+        let credentialsPath: string
+        let configPath: string
+
+        beforeEach(async function () {
+            tempFolder = await makeTemporaryToolkitFolder()
+            credentialsPath = path.join(tempFolder, 'credentials')
+            configPath = path.join(tempFolder, 'config')
+
+            // Stub environment variables to use temp files
+            sinon.stub(process, 'env').value({
+                AWS_SHARED_CREDENTIALS_FILE: credentialsPath,
+                AWS_CONFIG_FILE: configPath,
+            } as EnvironmentVariables)
+        })
+
+        afterEach(async function () {
+            await fs.delete(tempFolder, { recursive: true })
+            sinon.restore()
+        })
+
+        it('should update region in credentials file when profile exists there', async function () {
+            // Create credentials file with a profile without region
+            const credentialsContent = [
+                '[test-profile]',
+                'aws_access_key_id = XYZ',
+                'aws_secret_access_key = XYZ',
+                '',
+            ].join(os.EOL)
+            await fs.writeFile(credentialsPath, credentialsContent)
+
+            // Call the private method using bracket notation
+            await (SmusIamProfileSelector as any).updateProfileRegion('test-profile', 'us-west-2')
+
+            // Verify the region was added
+            const updatedContent = await fs.readFileText(credentialsPath)
+            assert.ok(updatedContent.includes('region = us-west-2'))
+            assert.ok(updatedContent.includes('[test-profile]'))
+            assert.ok(updatedContent.includes('aws_access_key_id = XYZ'))
+        })
+
+        it('should update region in config file when profile exists there', async function () {
+            // Create config file with a profile without region
+            const configContent = ['[profile test-profile]', 'output = json', ''].join(os.EOL)
+            await fs.writeFile(configPath, configContent)
+
+            // Call the private method
+            await (SmusIamProfileSelector as any).updateProfileRegion('test-profile', 'eu-west-1')
+
+            // Verify the region was added
+            const updatedContent = await fs.readFileText(configPath)
+            assert.ok(updatedContent.includes('region = eu-west-1'))
+            assert.ok(updatedContent.includes('[profile test-profile]'))
+            assert.ok(updatedContent.includes('output = json'))
+        })
+
+        it('should handle multiple profiles in credentials file', async function () {
+            // Create credentials file with multiple profiles
+            const credentialsContent = [
+                '[default]',
+                'aws_access_key_id = XYZ',
+                'aws_secret_access_key = XYZ',
+                '',
+                '[test-profile]',
+                'aws_access_key_id = XYZ',
+                'aws_secret_access_key = XYZ',
+                '',
+                '[another-profile]',
+                'aws_access_key_id = XYZ',
+                'aws_secret_access_key = XYZ',
+                '',
+            ].join(os.EOL)
+            await fs.writeFile(credentialsPath, credentialsContent)
+
+            // Update the region for test-profile
+            await (SmusIamProfileSelector as any).updateProfileRegion('test-profile', 'us-west-2')
+
+            // Verify the region was added only to test-profile
+            const updatedContent = await fs.readFileText(credentialsPath)
+            const lines = updatedContent.split(os.EOL)
+
+            // Find test-profile section
+            const testProfileIndex = lines.findIndex((line) => line.includes('[test-profile]'))
+            const anotherProfileIndex = lines.findIndex((line) => line.includes('[another-profile]'))
+
+            // Check that region is between test-profile and another-profile
+            const testProfileSection = lines.slice(testProfileIndex, anotherProfileIndex).join(os.EOL)
+            assert.ok(testProfileSection.includes('region = us-west-2'))
+
+            // Check that other profiles are unchanged
+            assert.ok(updatedContent.includes('[default]'))
+            assert.ok(updatedContent.includes('[another-profile]'))
+        })
+
+        it('should throw error when profile does not exist in either file', async function () {
+            // Create both files without the target profile
+            const credentialsContent = ['[default]', 'aws_access_key_id = XYZ', 'aws_secret_access_key = XYZ'].join(
+                os.EOL
+            )
+            const configContent = ['[profile default]', 'region = us-east-1'].join(os.EOL)
+            await fs.writeFile(credentialsPath, credentialsContent)
+            await fs.writeFile(configPath, configContent)
+
+            // Attempt to update non-existent profile
+            await assert.rejects(
+                async () => {
+                    await (SmusIamProfileSelector as any).updateProfileRegion('non-existent-profile', 'us-west-2')
+                },
+                (error: Error) => {
+                    assert.ok(error.message.includes('not found'))
+                    return true
+                }
+            )
+        })
+
+        it('should throw error when neither file exists', async function () {
+            // Attempt to update profile
+            await assert.rejects(
+                async () => {
+                    await (SmusIamProfileSelector as any).updateProfileRegion('test-profile', 'us-west-2')
+                },
+                (error: Error) => {
+                    assert.ok(error.message.includes('not found'))
+                    return true
+                }
+            )
+        })
+    })
 })