Add unit tests for getVariableExprsFromSelection

src/lang/queryAst.test.ts

@@ -15,6 +15,7 @@ import {
   findAllPreviousVariables,
   findUsesOfTagInPipe,
   getNodeFromPath,
+  getVariableExprsFromSelection,
   hasSketchPipeBeenExtruded,
   isCursorInFunctionDefinition,
   isNodeSafeToReplace,
@@ -27,7 +28,7 @@ import { topLevelRange } from '@src/lang/util'
 import type { Identifier, PathToNode } from '@src/lang/wasm'
 import { assertParse, recast } from '@src/lang/wasm'
 import { initPromise } from '@src/lang/wasmUtils'
-import { type Selection } from '@src/lib/selections'
+import type { Selections, Selection } from '@src/lib/selections'
 import { enginelessExecutor } from '@src/lib/testHelpers'
 import { err } from '@src/lib/trap'
 
@@ -778,3 +779,184 @@ describe('Testing specific sketch getNodeFromPath workflow', () => {
     expect(result).toEqual(false)
   })
 })
+
+describe('Testing getVariableExprsFromSelection', () => {
+  it('should find the variable expr in a simple profile selection', async () => {
+    const circleProfileInVar = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+`
+    const ast = assertParse(circleProfileInVar)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifact = artifactGraph.values().find((a) => a.type === 'path')
+    if (!artifact) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: [
+        {
+          codeRef: artifact.codeRef,
+          artifact,
+        },
+      ],
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+
+    expect(variableExprs.exprs).toHaveLength(1)
+    if (variableExprs.exprs[0].type !== 'Name') {
+      throw new Error(`Expected Name, got ${variableExprs.exprs[0].type}`)
+    }
+
+    expect(variableExprs.exprs[0].name.name).toEqual('profile001')
+
+    expect(variableExprs.paths).toHaveLength(1)
+    expect(variableExprs.paths[0]).toEqual([
+      ['body', ''],
+      [1, 'index'],
+      ['declaration', 'VariableDeclaration'],
+      ['init', ''],
+    ])
+  })
+
+  it('should return the pipe substition symbol in a variable-less simple profile selection', async () => {
+    const circleProfileInVar = `startSketchOn(XY)
+  |> circle(center = [0, 0], radius = 1)
+`
+    const ast = assertParse(circleProfileInVar)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifact = artifactGraph.values().find((a) => a.type === 'path')
+    if (!artifact) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: [
+        {
+          codeRef: artifact.codeRef,
+          artifact,
+        },
+      ],
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+
+    expect(variableExprs.exprs).toHaveLength(1)
+    expect(variableExprs.exprs[0].type).toEqual('PipeSubstitution')
+
+    expect(variableExprs.paths).toHaveLength(1)
+    expect(variableExprs.paths[0]).toEqual([
+      ['body', ''],
+      [0, 'index'],
+      ['expression', 'ExpressionStatement'],
+      ['body', 'PipeExpression'],
+      [1, 'index'],
+    ])
+  })
+
+  it('should find the variable exprs in a multi profile selection ', async () => {
+    const circleProfileInVar = `sketch001 = startSketchOn(XY)
+profile001 = circle(sketch001, center = [0, 0], radius = 1)
+profile002 = circle(sketch001, center = [2, 2], radius = 1)
+`
+    const ast = assertParse(circleProfileInVar)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifacts = [...artifactGraph.values()].filter(
+      (a) => a.type === 'path'
+    )
+    if (!artifacts || artifacts.length !== 2) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: artifacts.map((artifact) => {
+        return {
+          codeRef: artifact.codeRef,
+          artifact,
+        }
+      }),
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+
+    expect(variableExprs.exprs).toHaveLength(2)
+    if (variableExprs.exprs[0].type !== 'Name') {
+      throw new Error(`Expected Name, got ${variableExprs.exprs[0].type}`)
+    }
+
+    if (variableExprs.exprs[1].type !== 'Name') {
+      throw new Error(`Expected Name, got ${variableExprs.exprs[1].type}`)
+    }
+
+    expect(variableExprs.exprs[0].name.name).toEqual('profile001')
+    expect(variableExprs.exprs[1].name.name).toEqual('profile002')
+
+    expect(variableExprs.paths).toHaveLength(2)
+    expect(variableExprs.paths[0]).toEqual([
+      ['body', ''],
+      [1, 'index'],
+      ['declaration', 'VariableDeclaration'],
+      ['init', ''],
+    ])
+    expect(variableExprs.paths[1]).toEqual([
+      ['body', ''],
+      [2, 'index'],
+      ['declaration', 'VariableDeclaration'],
+      ['init', ''],
+    ])
+  })
+
+  it('should return the pipe substition symbol and a variable name in a complex multi profile selection', async () => {
+    const circleProfileInVar = `startSketchOn(XY)
+  |> circle(center = [0, 0], radius = 1)
+profile002 = circle(sketch001, center = [2, 2], radius = 1)
+`
+    const ast = assertParse(circleProfileInVar)
+    const { artifactGraph } = await enginelessExecutor(ast)
+    const artifacts = [...artifactGraph.values()].filter(
+      (a) => a.type === 'path'
+    )
+    if (!artifacts || artifacts.length !== 2) {
+      throw new Error('Artifact not found in the graph')
+    }
+    const selections: Selections = {
+      graphSelections: artifacts.map((artifact) => {
+        return {
+          codeRef: artifact.codeRef,
+          artifact,
+        }
+      }),
+      otherSelections: [],
+    }
+    const variableExprs = getVariableExprsFromSelection(selections, ast)
+    if (err(variableExprs)) throw variableExprs
+
+    expect(variableExprs.exprs).toHaveLength(2)
+    if (variableExprs.exprs[0].type !== 'PipeSubstitution') {
+      throw new Error(
+        `Expected PipeSubstitution, got ${variableExprs.exprs[0].type}`
+      )
+    }
+
+    if (variableExprs.exprs[1].type !== 'Name') {
+      throw new Error(`Expected Name, got ${variableExprs.exprs[1].type}`)
+    }
+
+    expect(variableExprs.exprs[1].name.name).toEqual('profile002')
+
+    expect(variableExprs.paths).toHaveLength(2)
+    expect(variableExprs.paths[0]).toEqual([
+      ['body', ''],
+      [0, 'index'],
+      ['expression', 'ExpressionStatement'],
+      ['body', 'PipeExpression'],
+      [1, 'index'],
+    ])
+    expect(variableExprs.paths[1]).toEqual([
+      ['body', ''],
+      [1, 'index'],
+      ['declaration', 'VariableDeclaration'],
+      ['init', ''],
+    ])
+  })
+})