From fc64ac4d3f1eae034083235ee7a18dd48e4d1e19 Mon Sep 17 00:00:00 2001
From: vadimkibana <vadimkibana@gmail.com>
Date: Wed, 31 Jul 2024 14:35:52 +0200
Subject: [PATCH] add early walk termination test

---
 .../src/visitor/__tests__/expressions.test.ts | 27 +++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts b/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts
index 586c4cce5d73a..efd30f035e7ca 100644
--- a/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts
+++ b/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts
@@ -33,6 +33,33 @@ test('"visitExpression" captures all non-captured expressions', () => {
   );
 });
 
+test('can terminate walk early, does not visit all literals', () => {
+  const numbers: number[] = [];
+  const { ast } = getAstAndSyntaxErrors(`
+    FROM index
+      | STATS 0, 1, 2, 3
+      | LIMIT 123
+  `);
+  const result = new Visitor()
+    .on('visitExpression', (ctx) => {
+      return 0;
+    })
+    .on('visitLiteralExpression', (ctx) => {
+      numbers.push(ctx.node.value as number);
+      return ctx.node.value;
+    })
+    .on('visitCommand', (ctx) => {
+      for (const res of ctx.visitArguments()) if (res) return res;
+    })
+    .on('visitQuery', (ctx) => {
+      for (const res of ctx.visitCommands()) if (res) return res;
+    })
+    .visitQuery(ast);
+
+  expect(result).toBe(1);
+  expect(numbers).toEqual([0, 1]);
+});
+
 test('"visitColumnExpression" takes over all column visits', () => {
   const { ast } = getAstAndSyntaxErrors(`
     FROM index | STATS a