Fix AWS get_price action to use assumed role credentials

packages/blocks/aws/src/lib/actions/get-price-action.ts

@@ -3,6 +3,8 @@ import { Property, createAction } from '@openops/blocks-framework';
 import {
   amazonAuth,
   getAttributeValues,
+  getAwsAccountsSingleSelectDropdown,
+  getCredentialsForAccount,
   getCredentialsFromAuth,
   getPriceListWithCache,
   getServices,
@@ -17,22 +19,26 @@ export const getPriceAction = createAction({
   displayName: 'Get Price from Price Catalog',
   isWriteAction: false,
   props: {
+    account: getAwsAccountsSingleSelectDropdown().accounts,
     service: Property.Dropdown({
       displayName: 'Service Code',
       description: 'Service code for which to fetch the price',
-      refreshers: ['auth'],
+      refreshers: ['auth', 'account', 'account.accounts'],
       required: true,
-      options: async ({ auth }: any) => {
-        if (!auth) {
+      options: async ({ auth, account }: any) => {
+        if (!auth || !account) {
           return {
             disabled: true,
             options: [],
-            placeholder: 'Please authenticate first',
+            placeholder: 'Please authenticate first and select account',
           };
         }
 
-        const credentials = await getCredentialsFromAuth(auth);
         try {
+          const credentials = await getCredentialsForAccount(
+            auth,
+            account?.['accounts'],
+          );
           const services = await getServices(credentials, PRICING_REGION);
 
           if (!services.length) {
@@ -66,9 +72,9 @@ export const getPriceAction = createAction({
       displayName: '',
       description: '',
       required: true,
-      refreshers: ['auth', 'service'],
-      props: async ({ auth, service }, { input }) => {
-        if (!auth || !service) {
+      refreshers: ['auth', 'account', 'account.accounts', 'service'],
+      props: async ({ auth, account, service }, { input }) => {
+        if (!auth || !account || !service) {
           return {};
         }
 
@@ -101,7 +107,10 @@ export const getPriceAction = createAction({
                   };
                 }
 
-                const credentials = await getCredentialsFromAuth(auth);
+                const credentials = await getCredentialsForAccount(
+                  auth,
+                  account?.['accounts'],
+                );
                 try {
                   const attributeValues: AttributeValue[] =
                     await getAttributeValues(
@@ -139,17 +148,20 @@ export const getPriceAction = createAction({
   },
   async run(context) {
     try {
-      const { service, queryFilters } = context.propsValue;
+      const { account, service, queryFilters } = context.propsValue;
       const filters = queryFilters['queryFilters'].map((filter: any) => {
         return {
           Field: filter.attributeName,
           Type: FilterType.TERM_MATCH,
           Value: filter.attributeValue,
         };
       });
+      const credentials = account?.['accounts']
+        ? await getCredentialsForAccount(context.auth, account['accounts'])
+        : await getCredentialsFromAuth(context.auth);
 
       const priceList = getPriceListWithCache(
-        context.auth,
+        credentials,
         service.ServiceCode!,
         filters,
         PRICING_REGION,

packages/blocks/aws/test/get-price-action.test.ts

@@ -1,6 +1,7 @@
 const openopsCommonMock = {
   ...jest.requireActual('@openops/common'),
   getCredentialsFromAuth: jest.fn(),
+  getCredentialsForAccount: jest.fn(),
   getAttributeValues: jest.fn(),
   getServices: jest.fn(),
   getPriceListWithCache: jest.fn(),
@@ -18,6 +19,9 @@ describe('getPriceAction', () => {
     openopsCommonMock.getCredentialsFromAuth.mockResolvedValue({
       someCreds: 'some value',
     });
+    openopsCommonMock.getCredentialsForAccount.mockResolvedValue({
+      someCreds: 'some value',
+    });
   });
 
   const auth = {
@@ -34,12 +38,16 @@ describe('getPriceAction', () => {
     },
     auth: auth,
     propsValue: {
-      accountId: 'some account id',
+      account: { accounts: 'some account id' },
     },
   };
 
   test('should create action with correct properties', () => {
     expect(getPriceAction.props).toMatchObject({
+      account: {
+        type: 'DYNAMIC',
+        required: true,
+      },
       service: {
         type: 'DROPDOWN',
         required: true,
@@ -80,22 +88,40 @@ describe('getPriceAction', () => {
 
       expect(openopsCommonMock.getPriceListWithCache).toHaveBeenCalledTimes(1);
       expect(openopsCommonMock.getPriceListWithCache).toHaveBeenCalledWith(
-        auth,
+        { someCreds: 'some value' },
         'some service',
         expectedFilters,
         'us-east-1',
       );
+      expect(openopsCommonMock.getCredentialsForAccount).toHaveBeenCalledWith(
+        auth,
+        'some account id',
+      );
     },
   );
 
+  test('should fallback to getCredentialsFromAuth when account is missing in run', async () => {
+    openopsCommonMock.getPriceListWithCache.mockResolvedValue('mockResult');
+    context.propsValue = {
+      service: { ServiceCode: 'some service' },
+      queryFilters: { queryFilters: [] },
+    };
+
+    const result = (await getPriceAction.run(context)) as any;
+
+    expect(result).toEqual('mockResult');
+    expect(openopsCommonMock.getCredentialsFromAuth).toHaveBeenCalledWith(auth);
+    expect(openopsCommonMock.getCredentialsForAccount).not.toHaveBeenCalled();
+  });
+
   test('should return the list of services in property service', async () => {
     openopsCommonMock.getServices.mockResolvedValue([
       { ServiceCode: 'service1' },
       { ServiceCode: 'service2' },
     ]);
 
     const result = await getPriceAction.props['service'].options(
-      { auth },
+      { auth, account: { accounts: 'some account id' } },
       context,
     );
 
@@ -118,6 +144,7 @@ describe('getPriceAction', () => {
       ServiceCode: 'some service',
       AttributeNames: ['some attibute'],
     };
+    context.propsValue.account = { accounts: 'some account id' };
     context.propsValue.auth = auth;
     context.input = { attributeName: 'some attibute' };