🐛 Fix bug with function return type check

elastic · Oct 5, 2022 · f26e9b7 · f26e9b7
1 parent 125ecdd
commit f26e9b7
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 17 deletions.
diff --git a/...ugins/lens/public/indexpattern_datasource/operations/definitions/formula/formula.test.tsx b/...ugins/lens/public/indexpattern_datasource/operations/definitions/formula/formula.test.tsx
@@ -1612,7 +1612,53 @@ invalid: "
           });
         });
       } else {
-        it(`[${fn}] returns an error if the argument is of the wrong type`, () => {
+        const indexReverseMap = {
+          cond: [0],
+          left: [1],
+          right: [2],
+          all: [0, 1, 2],
+        };
+        it.each`
+          cond             | left                    | right                   | expectedFail
+          ${'1'}           | ${'2'}                  | ${'3'}                  | ${'cond'}
+          ${'1 > 1'}       | ${'2 > 2'}              | ${'3'}                  | ${'left'}
+          ${'1 > 1'}       | ${'2'}                  | ${'3 > 3'}              | ${'right'}
+          ${'1'}           | ${'2 > 2'}              | ${'3 > 3'}              | ${'all'}
+          ${'count()'}     | ${'average(bytes)'}     | ${'average(bytes)'}     | ${'cond'}
+          ${'count() > 1'} | ${'average(bytes) > 2'} | ${'average(bytes)'}     | ${'left'}
+          ${'count() > 1'} | ${'average(bytes)'}     | ${'average(bytes) > 3'} | ${'right'}
+          ${'count()'}     | ${'average(bytes) > 2'} | ${'average(bytes) > 3'} | ${'all'}
+        `(
+          `[${fn}] returns an error if $expectedFail argument is/are of the wrong type: ${fn}($cond, $left, $right)`,
+          ({
+            cond,
+            left,
+            right,
+            expectedFail,
+          }: {
+            cond: string;
+            left: string;
+            right: string;
+            expectedFail: keyof typeof indexReverseMap;
+          }) => {
+            const argsSorted = [cond, left, right];
+            expect(
+              formulaOperation.getErrorMessage!(
+                getNewLayerWithFormula(`${fn}(${cond}, ${left}, ${right})`),
+                'col1',
+                indexPattern,
+                operationDefinitionMap
+              )
+            ).toEqual(
+              indexReverseMap[expectedFail].map((i) => {
+                const arg = tinymathFunctions[fn].positionalArguments[i];
+                const passedValue = />/.test(argsSorted[i]) ? 'boolean' : 'number';
+                return `The ${arg.name} argument for the operation ${fn} in the Formula is of the wrong type: ${passedValue} instead of ${arg.type}`;
+              })
+            );
+          }
+        );
+        it(`[${fn}] returns an error if the condition argument is of the wrong type`, () => {
           const [firstArg] = tinymathFunctions[fn].positionalArguments;
           expect(
             formulaOperation.getErrorMessage!(

diff --git a/.../plugins/lens/public/indexpattern_datasource/operations/definitions/formula/validation.ts b/.../plugins/lens/public/indexpattern_datasource/operations/definitions/formula/validation.ts
@@ -118,6 +118,22 @@ function getNodeLocation(node: TinymathFunction): TinymathLocation[] {
   return [node.location].filter(nonNullable);
 }
 
+function getArgumentType(arg: TinymathAST, operations: Record<string, GenericOperationDefinition>) {
+  if (!isObject(arg)) {
+    return typeof arg;
+  }
+  if (arg.type === 'function') {
+    if (tinymathFunctions[arg.name]) {
+      return tinymathFunctions[arg.name].outputType ?? 'number';
+    }
+    // Assume it's a number for now
+    if (operations[arg.name]) {
+      return 'number';
+    }
+  }
+  // leave for now other argument types
+}
+
 export function isParsingError(message: string) {
   return message.includes('Failed to parse expression');
 }
@@ -688,7 +704,7 @@ function runFullASTValidation(
     const [firstArg] = node?.args || [];
 
     if (!nodeOperation) {
-      errors.push(...validateMathNodes(node, missingVariablesSet));
+      errors.push(...validateMathNodes(node, missingVariablesSet, operations));
       // carry on with the validation for all the functions within the math operation
       if (functions?.length) {
         return errors.concat(functions.flatMap((fn) => validateNode(fn)));
@@ -785,7 +801,8 @@ function runFullASTValidation(
         // In general this should be handled down the Esaggs route rather than here
         const isFirstArgumentNotValid = Boolean(
           !isArgumentValidType(firstArg, 'function') ||
-            (isMathNode(firstArg) && validateMathNodes(firstArg, missingVariablesSet).length)
+            (isMathNode(firstArg) &&
+              validateMathNodes(firstArg, missingVariablesSet, operations).length)
         );
         // First field has a special handling
         if (isFirstArgumentNotValid) {
@@ -989,7 +1006,11 @@ export function isArgumentValidType(arg: TinymathAST | string, type: TinymathNod
   return isObject(arg) && arg.type === type;
 }
 
-export function validateMathNodes(root: TinymathAST, missingVariableSet: Set<string>) {
+export function validateMathNodes(
+  root: TinymathAST,
+  missingVariableSet: Set<string>,
+  operations: Record<string, GenericOperationDefinition>
+) {
   const mathNodes = findMathNodes(root);
   const errors: ErrorWrapper[] = [];
   mathNodes.forEach((node: TinymathFunction) => {
@@ -1083,26 +1104,26 @@ export function validateMathNodes(root: TinymathAST, missingVariableSet: Set<str
         );
       }
     }
-    const wrongTypeArgumentIndex = positionalArguments.findIndex(({ type }, index) => {
-      const arg = node.args[index];
-      if (arg != null) {
-        if (!isObject(arg)) {
-          return typeof arg !== type;
-        }
-        if (arg.type === 'function' && tinymathFunctions[arg.name]) {
-          const { outputType = 'number' } = tinymathFunctions[arg.name];
-          return outputType !== type;
+    const wrongTypeArgumentIndexes = positionalArguments
+      .map(({ type }, index) => {
+        const arg = node.args[index];
+        if (arg != null) {
+          const argType = getArgumentType(arg, operations);
+          if (argType && argType !== type) {
+            return index;
+          }
         }
-      }
-    });
-    if (wrongTypeArgumentIndex > -1) {
+      })
+      .filter(nonNullable);
+    for (const wrongTypeArgumentIndex of wrongTypeArgumentIndexes) {
+      const arg = node.args[wrongTypeArgumentIndex];
       errors.push(
         getMessageFromId({
           messageId: 'wrongTypeArgument',
           values: {
             operation: node.name,
             name: positionalArguments[wrongTypeArgumentIndex].name,
-            type: typeof node.args[wrongTypeArgumentIndex],
+            type: getArgumentType(arg, operations) || 'number',
             expectedType: positionalArguments[wrongTypeArgumentIndex].type || '',
           },
           locations: getNodeLocation(node),