feat(vscode): add test recognition to vscode

[ci skip]

Author: benfdking

sqlmesh/lsp/context.py

@@ -5,7 +5,7 @@
 
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
-from sqlmesh.lsp.custom import ModelForRendering
+from sqlmesh.lsp.custom import ModelForRendering, TestEntry
 from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
 from sqlmesh.lsp.uri import URI
 from lsprotocol import types
@@ -63,6 +63,18 @@ def __init__(self, context: Context) -> None:
             **audit_map,
         }
 
+    def list_workspace_tests(self) -> t.List[TestEntry]:
+        """List all tests in the workspace."""
+        tests = self.context.load_model_tests()
+        # TODO Probably want to get all the positions for the tests
+        return [
+            TestEntry(
+                name=test.test_name,
+                uri=URI.from_path(test.path).value,
+            )
+            for test in tests
+        ]
+
     def render_model(self, uri: URI) -> t.List[RenderModelEntry]:
         """Get rendered models for a file, using cache when available.
 

sqlmesh/lsp/custom.py

@@ -143,3 +143,27 @@ class FormatProjectResponse(CustomMethodResponseBaseClass):
     """
 
     pass
+
+
+LIST_WORKSPACE_TESTS_FEATURE = "sqlmesh/list_workspace_tests"
+
+
+class ListWorkspaceTestsRequest(CustomMethodRequestBaseClass):
+    """
+    Request to list all tests in the current project.
+    """
+
+    pass
+
+
+class TestEntry(PydanticModel):
+    """
+    An entry representing a test in the workspace.
+    """
+
+    name: str
+    uri: str
+
+
+class ListWorkspaceTestsResponse(CustomMethodResponseBaseClass):
+    tests: t.List[TestEntry]

sqlmesh/lsp/main.py

@@ -46,6 +46,9 @@
     FormatProjectRequest,
     FormatProjectResponse,
     CustomMethod,
+    LIST_WORKSPACE_TESTS_FEATURE,
+    ListWorkspaceTestsRequest,
+    ListWorkspaceTestsResponse,
 )
 from sqlmesh.lsp.errors import ContextFailedError, context_error_to_diagnostic
 from sqlmesh.lsp.hints import get_hints
@@ -126,11 +129,26 @@ def __init__(
             API_FEATURE: self._custom_api,
             SUPPORTED_METHODS_FEATURE: self._custom_supported_methods,
             FORMAT_PROJECT_FEATURE: self._custom_format_project,
+            LIST_WORKSPACE_TESTS_FEATURE: self._list_workspace_tests,
         }
 
         # Register LSP features (e.g., formatting, hover, etc.)
         self._register_features()
 
+    def _list_workspace_tests(
+        self,
+        ls: LanguageServer,
+        params: ListWorkspaceTestsRequest,
+    ) -> ListWorkspaceTestsResponse:
+        """List all tests in the current workspace."""
+        try:
+            context = self._context_get_or_load()
+            tests = context.list_workspace_tests()
+            return ListWorkspaceTestsResponse(tests=tests)
+        except Exception as e:
+            ls.log_trace(f"Error listing workspace tests: {e}")
+            return ListWorkspaceTestsResponse(tests=[])
+
     # All the custom LSP methods are registered here and prefixed with _custom
     def _custom_all_models(self, ls: LanguageServer, params: AllModelsRequest) -> AllModelsResponse:
         uri = URI(params.textDocument.uri)

vscode/extension/src/lsp/custom.ts

@@ -111,3 +111,21 @@ interface FormatProjectResponse extends BaseResponse {}
 interface BaseResponse {
   response_error?: string
 }
+
+export interface ListWorkspaceTests {
+  method: 'sqlmesh/list_workspace_tests'
+  request: ListWorkspaceTestsRequest
+  response: ListWorkspaceTestsResponse
+}
+
+type ListWorkspaceTestsRequest = object
+
+interface TestEntry {
+  name: string
+  uri: string
+  //  TODO Probably want to add position at some point
+}
+
+interface ListWorkspaceTestsResponse {
+  tests: TestEntry[]
+}

vscode/extension/src/tests/tests.ts

@@ -0,0 +1,66 @@
+import * as vscode from 'vscode'
+import { LSPClient } from '../lsp/lsp'
+
+export const controller = vscode.tests.createTestController(
+  'sqlmeshTests',
+  'SQLMesh Tests',
+)
+
+export const setupTestController = (lsp: LSPClient) => {
+  // First, create the `resolveHandler`. This may initially be called with
+  // "undefined" to ask for all tests in the workspace to be discovered, usually
+  // when the user opens the Test Explorer for the first time.
+  controller.resolveHandler = async test => {
+    if (!test) {
+      await discoverAllFilesInWorkspace()
+    } else {
+      await parseTestsInFileContents(test)
+    }
+  }
+
+  // When text documents are open, parse tests in them.
+  vscode.workspace.onDidOpenTextDocument(parseTestsInDocument)
+  // We could also listen to document changes to re-parse unsaved changes:
+  vscode.workspace.onDidChangeTextDocument(e =>
+    parseTestsInDocument(e.document),
+  )
+
+  // In this function, we'll get the file TestItem if we've already found it,
+  // otherwise we'll create it with `canResolveChildren = true` to indicate it
+  // can be passed to the `controller.resolveHandler` to gets its children.
+  function getOrCreateFile(uri: vscode.Uri) {
+    const existing = controller.items.get(uri.toString())
+    if (existing) {
+      return existing
+    }
+
+    const file = controller.createTestItem(
+      uri.toString(),
+      uri.path.split('/').pop()!,
+      uri,
+    )
+    file.canResolveChildren = true
+    return file
+  }
+
+  function parseTestsInDocument(e: vscode.TextDocument) {
+    if (e.uri.scheme === 'file' && e.uri.path.endsWith('.md')) {
+      parseTestsInFileContents(getOrCreateFile(e.uri), e.getText())
+    }
+  }
+
+  async function parseTestsInFileContents(
+    file: vscode.TestItem,
+    contents?: string,
+  ) {
+    // If a document is open, VS Code already knows its contents. If this is being
+    // called from the resolveHandler when a document isn't open, we'll need to
+    // read them from disk ourselves.
+    if (contents === undefined) {
+      const rawContent = await vscode.workspace.fs.readFile(file.uri)
+      contents = new TextDecoder().decode(rawContent)
+    }
+
+    // some custom logic to fill in test.children from the contents...
+  }
+}