Add test for call/construct signatures

microsoft · andrewbranch · Jul 12, 2019 · Jul 9, 2019 · Jul 9, 2019 · Jul 10, 2019
commit 7be4e0aeda62537499dda27f5cb884ba1a8343b8
diff --git a/src/services/codefixes/addMissingAwait.ts b/src/services/codefixes/addMissingAwait.ts
@@ -2,6 +2,10 @@
 namespace ts.codefix {
     type ContextualTrackChangesFunction = (cb: (changeTracker: textChanges.ChangeTracker) => void) => FileTextChanges[];
     const fixId = "addMissingAwait";
+    const callableConstructableErrorCodes = [
+        Diagnostics.This_expression_is_not_callable.code,
+        Diagnostics.This_expression_is_not_constructable.code,
+    ];
     const errorCodes = [
         Diagnostics.An_arithmetic_operand_must_be_of_type_any_number_bigint_or_an_enum_type.code,
         Diagnostics.The_left_hand_side_of_an_arithmetic_operation_must_be_of_type_any_number_bigint_or_an_enum_type.code,
@@ -17,8 +21,7 @@ namespace ts.codefix {
         Diagnostics.Type_0_must_have_a_Symbol_iterator_method_that_returns_an_iterator.code,
         Diagnostics.Type_0_must_have_a_Symbol_asyncIterator_method_that_returns_an_async_iterator.code,
         Diagnostics.Argument_of_type_0_is_not_assignable_to_parameter_of_type_1.code,
-        Diagnostics.This_expression_is_not_callable.code,
-        Diagnostics.This_expression_is_not_constructable.code,
+        ...callableConstructableErrorCodes,
     ];
 
     registerCodeFix({
@@ -34,8 +37,8 @@ namespace ts.codefix {
                     return;
                 }
                 const trackChanges: ContextualTrackChangesFunction = cb => (cb(t), []);
-                return getDeclarationSiteFix(context, expression, checker, trackChanges)
-                    || getUseSiteFix(context, expression, checker, trackChanges);
+                return getDeclarationSiteFix(context, expression, diagnostic.code, checker, trackChanges)
+                    || getUseSiteFix(context, expression, diagnostic.code, checker, trackChanges);
             });
         },
     });
@@ -50,24 +53,24 @@ namespace ts.codefix {
         const checker = context.program.getTypeChecker();
         const trackChanges: ContextualTrackChangesFunction = cb => textChanges.ChangeTracker.with(context, cb);
         return compact([
-            getDeclarationSiteFix(context, expression, checker, trackChanges),
-            getUseSiteFix(context, expression, checker, trackChanges)]);
+            getDeclarationSiteFix(context, expression, errorCode, checker, trackChanges),
+            getUseSiteFix(context, expression, errorCode, checker, trackChanges)]);
     }
 
-    function getDeclarationSiteFix(context: CodeFixContext | CodeFixAllContext, expression: Expression, checker: TypeChecker, trackChanges: ContextualTrackChangesFunction) {
+    function getDeclarationSiteFix(context: CodeFixContext | CodeFixAllContext, expression: Expression, errorCode: number, checker: TypeChecker, trackChanges: ContextualTrackChangesFunction) {
         const { sourceFile } = context;
         const awaitableInitializer = findAwaitableInitializer(expression, sourceFile, checker);
         if (awaitableInitializer) {
-            const initializerChanges = trackChanges(t => makeChange(t, sourceFile, checker, awaitableInitializer));
+            const initializerChanges = trackChanges(t => makeChange(t, errorCode, sourceFile, checker, awaitableInitializer));
             return createCodeFixActionNoFixId(
                 "addMissingAwaitToInitializer",
                 initializerChanges,
                 [Diagnostics.Add_await_to_initializer_for_0, expression.getText(sourceFile)]);
         }
     }
 
-    function getUseSiteFix(context: CodeFixContext | CodeFixAllContext, expression: Expression, checker: TypeChecker, trackChanges: ContextualTrackChangesFunction) {
-        const changes = trackChanges(t => makeChange(t, context.sourceFile, checker, expression));
+    function getUseSiteFix(context: CodeFixContext | CodeFixAllContext, expression: Expression, errorCode: number, checker: TypeChecker, trackChanges: ContextualTrackChangesFunction) {
+        const changes = trackChanges(t => makeChange(t, errorCode, context.sourceFile, checker, expression));
         return createCodeFixAction(fixId, changes, Diagnostics.Add_await, fixId, Diagnostics.Fix_all_expressions_possibly_missing_await);
     }
 
@@ -126,7 +129,7 @@ namespace ts.codefix {
         return declaration.initializer;
     }
 
-    function makeChange(changeTracker: textChanges.ChangeTracker, sourceFile: SourceFile, checker: TypeChecker, insertionSite: Expression) {
+    function makeChange(changeTracker: textChanges.ChangeTracker, errorCode: number, sourceFile: SourceFile, checker: TypeChecker, insertionSite: Expression) {
         if (isBinaryExpression(insertionSite)) {
             const { left, right } = insertionSite;
             const leftType = checker.getTypeAtLocation(left);
@@ -135,9 +138,12 @@ namespace ts.codefix {
             const newRight = checker.getPromisedTypeOfPromise(rightType) ? createAwait(right) : right;
             changeTracker.replaceNode(sourceFile, left, newLeft);
             changeTracker.replaceNode(sourceFile, right, newRight);
-            return;
         }
-
-        changeTracker.replaceNode(sourceFile, insertionSite, createAwait(insertionSite));
+        else if (contains(callableConstructableErrorCodes, errorCode) && isCallOrNewExpression(insertionSite.parent)) {
+            changeTracker.replaceNode(sourceFile, insertionSite, createParen(createAwait(insertionSite)));
+        }
+        else {
+            changeTracker.replaceNode(sourceFile, insertionSite, createAwait(insertionSite));
+        }
     }
 }
diff --git a/tests/cases/fourslash/codeFixAddMissingAwait_signatures.ts b/tests/cases/fourslash/codeFixAddMissingAwait_signatures.ts
@@ -0,0 +1,50 @@
+/// <reference path="fourslash.ts" />
+////async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+////  a();
+////  b();
+////  new C();
+////}
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 0,
+  newFileContent:
+`async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+  (await a)();
+  b();
+  new C();
+}`
+});
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 1,
+  newFileContent:
+`async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+  a();
+  (await b)();
+  new C();
+}`
+});
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 2,
+  newFileContent:
+`async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+  a();
+  b();
+  new (await C)();
+}`
+});
+
+verify.codeFixAll({
+  fixAllDescription: ts.Diagnostics.Fix_all_expressions_possibly_missing_await.message,
+  fixId: "addMissingAwait",
+  newFileContent:
+`async function fn(a: Promise<() => void>, b: Promise<() => void> | (() => void), C: Promise<{ new(): any }>) {
+  (await a)();
+  (await b)();
+  new (await C)();
+}`
+});