Add arguments in ArrowFunction/Return while removing promise() (#353)

Author: trivikr

.changeset/heavy-lemons-train.md

@@ -0,0 +1,5 @@
+---
+"aws-sdk-js-codemod": patch
+---
+
+Add callExpression arguments in ArrowFunction/Return while remove promise()

src/transforms/v2-to-v3/__fixtures__/api-promise/global-import.input.ts

@@ -1,6 +1,8 @@
 import AWS from "aws-sdk";
 
-export const listTables = (client: AWS.DynamoDB) => client.listTables().promise();
+export const listTables = async (client: AWS.DynamoDB) => client.listTables().promise();
+export const listTagsOfResource = async (client: AWS.DynamoDB) =>
+  client.listTagsOfResource({ ResourceArn: "STRING_VALUE" }).promise();
 
 // Client as class member
 class ClientClassMember {
@@ -11,6 +13,10 @@ class ClientClassMember {
   }
 
   async listTables() {
-    return await this.clientInClass.listTables().promise();
+    return this.clientInClass.listTables().promise();
+  }
+
+  async listTagsOfResource() {
+    return this.clientInClass.listTagsOfResource({ ResourceArn: "STRING_VALUE" }).promise();
   }
 }

src/transforms/v2-to-v3/__fixtures__/api-promise/global-import.output.ts

@@ -1,6 +1,8 @@
 import { DynamoDB } from "@aws-sdk/client-dynamodb";
 
-export const listTables = (client: DynamoDB) => client.listTables();
+export const listTables = async (client: DynamoDB) => client.listTables();
+export const listTagsOfResource = async (client: DynamoDB) =>
+  client.listTagsOfResource({ ResourceArn: "STRING_VALUE" });
 
 // Client as class member
 class ClientClassMember {
@@ -11,6 +13,10 @@ class ClientClassMember {
   }
 
   async listTables() {
-    return await this.clientInClass.listTables();
+    return this.clientInClass.listTables();
+  }
+
+  async listTagsOfResource() {
+    return this.clientInClass.listTagsOfResource({ ResourceArn: "STRING_VALUE" });
   }
 }

src/transforms/v2-to-v3/__fixtures__/api-promise/service-import.input.ts

@@ -1,6 +1,8 @@
 import DynamoDB from "aws-sdk/clients/dynamodb";
 
-export const listTables = (client: DynamoDB) => client.listTables().promise();
+export const listTables = async (client: DynamoDB) => client.listTables().promise();
+export const listTagsOfResource = async (client: DynamoDB) =>
+  client.listTagsOfResource({ ResourceArn: "STRING_VALUE" }).promise();
 
 // Client as class member
 class ClientClassMember {
@@ -11,6 +13,10 @@ class ClientClassMember {
   }
 
   async listTables() {
-    return await this.clientInClass.listTables().promise();
+    return this.clientInClass.listTables().promise();
+  }
+
+  async listTagsOfResource() {
+    return this.clientInClass.listTagsOfResource({ ResourceArn: "STRING_VALUE" }).promise();
   }
 }

src/transforms/v2-to-v3/__fixtures__/api-promise/service-import.output.ts

@@ -1,6 +1,8 @@
 import { DynamoDB } from "@aws-sdk/client-dynamodb";
 
-export const listTables = (client: DynamoDB) => client.listTables();
+export const listTables = async (client: DynamoDB) => client.listTables();
+export const listTagsOfResource = async (client: DynamoDB) =>
+  client.listTagsOfResource({ ResourceArn: "STRING_VALUE" });
 
 // Client as class member
 class ClientClassMember {
@@ -11,6 +13,10 @@ class ClientClassMember {
   }
 
   async listTables() {
-    return await this.clientInClass.listTables();
+    return this.clientInClass.listTables();
+  }
+
+  async listTagsOfResource() {
+    return this.clientInClass.listTagsOfResource({ ResourceArn: "STRING_VALUE" });
   }
 }

src/transforms/v2-to-v3/apis/removePromiseForCallExpression.ts

@@ -2,27 +2,34 @@ import { ASTPath, CallExpression, MemberExpression } from "jscodeshift";
 
 export const removePromiseForCallExpression = (callExpression: ASTPath<CallExpression>) => {
   switch (callExpression.parentPath.value.type) {
-    case "AwaitExpression":
+    case "AwaitExpression": {
       callExpression.parentPath.value.argument = (
         callExpression.value.callee as MemberExpression
       ).object;
       break;
-    case "MemberExpression":
+    }
+    case "MemberExpression": {
       callExpression.parentPath.value.object = (
         callExpression.value.callee as MemberExpression
       ).object;
       break;
-    case "VariableDeclarator":
+    }
+    case "VariableDeclarator": {
       callExpression.parentPath.value.init = (
         callExpression.value.callee as MemberExpression
       ).object;
       break;
+    }
     case "ArrowFunctionExpression":
-    case "ReturnStatement":
-      callExpression.value.callee = (
-        (callExpression.value.callee as MemberExpression).object as CallExpression
-      ).callee;
+    case "ReturnStatement": {
+      const currentCalleeObject = (callExpression.value.callee as MemberExpression)
+        .object as CallExpression;
+      if (currentCalleeObject.arguments.length > 0) {
+        callExpression.value.arguments = currentCalleeObject.arguments;
+      }
+      callExpression.value.callee = currentCalleeObject.callee;
       break;
+    }
     default:
       throw new Error(
         `Removal of .promise() not implemented for ${callExpression.parentPath.value.type}`