feat(vscode): add test recognition to vscode

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, RunTestResponse
 from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
 from sqlmesh.lsp.uri import URI
 from lsprotocol import types
@@ -63,6 +63,56 @@ 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 get_document_tests(self, uri: URI) -> t.List[TestEntry]:
+        """Get tests for a specific document.
+
+        Args:
+            uri: The URI of the file to get tests for.
+
+        Returns:
+            List of TestEntry objects for the specified document.
+        """
+        tests = self.context.load_model_tests(tests=[str(uri.to_path())])
+        return [
+            TestEntry(
+                name=test.test_name,
+                # TODO NEED TO ADD RANGE
+                uri=URI.from_path(test.path).value,
+            )
+            for test in tests
+        ]
+
+    def run_test(self, uri: URI, test_name: str) -> RunTestResponse:
+        """Run a specific test for a model.
+
+        Args:
+            uri: The URI of the file containing the test.
+            test_name: The name of the test to run.
+
+        Returns:
+            List of annotated rule violations from the test run.
+        """
+        path = uri.to_path()
+        results = self.context.test(
+            tests=[str(path)],
+            match_patterns=[test_name],
+        )
+        if results.testsRun != 1:
+            raise ValueError(f"Expected to run 1 test, but ran {results.testsRun} tests.")
+        return RunTestResponse(success=len(results.successes) == 1)
+
     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,50 @@ 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]
+
+
+LIST_DOCUMENT_TESTS_FEATURE = "sqlmesh/list_document_tests"
+
+
+class ListDocumentTestsRequest(CustomMethodRequestBaseClass):
+    textDocument: types.TextDocumentIdentifier
+
+
+class ListDocumentTestsResponse(CustomMethodResponseBaseClass):
+    tests: t.List[TestEntry]
+
+
+RUN_TEST_FEATURE = "sqlmesh/run_test"
+
+
+class RunTestRequest(CustomMethodRequestBaseClass):
+    textDocument: types.TextDocumentIdentifier
+    testName: str
+
+
+class RunTestResponse(CustomMethodResponseBaseClass):
+    success: bool

sqlmesh/lsp/main.py

@@ -46,6 +46,15 @@
     FormatProjectRequest,
     FormatProjectResponse,
     CustomMethod,
+    LIST_WORKSPACE_TESTS_FEATURE,
+    ListWorkspaceTestsRequest,
+    ListWorkspaceTestsResponse,
+    LIST_DOCUMENT_TESTS_FEATURE,
+    ListDocumentTestsRequest,
+    ListDocumentTestsResponse,
+    RUN_TEST_FEATURE,
+    RunTestRequest,
+    RunTestResponse,
 )
 from sqlmesh.lsp.errors import ContextFailedError, context_error_to_diagnostic
 from sqlmesh.lsp.helpers import to_lsp_range, to_sqlmesh_position
@@ -127,11 +136,58 @@ 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,
+            LIST_DOCUMENT_TESTS_FEATURE: self._list_document_tests,
+            RUN_TEST_FEATURE: self._run_test,
         }
 
         # 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=[])
+
+    def _list_document_tests(
+        self,
+        ls: LanguageServer,
+        params: ListDocumentTestsRequest,
+    ) -> ListDocumentTestsResponse:
+        """List tests for a specific document."""
+        try:
+            uri = URI(params.textDocument.uri)
+            context = self._context_get_or_load(uri)
+            tests = context.get_document_tests(uri)
+            return ListDocumentTestsResponse(tests=tests)
+        except Exception as e:
+            ls.log_trace(f"Error listing document tests: {e}")
+            return ListDocumentTestsResponse(tests=[])
+
+    def _run_test(
+        self,
+        ls: LanguageServer,
+        params: RunTestRequest,
+    ) -> RunTestResponse:
+        """Run a specific test."""
+        try:
+            uri = URI(params.textDocument.uri)
+            context = self._context_get_or_load(uri)
+            result = context.run_test(uri, params.testName)
+            return result
+        except Exception as e:
+            ls.log_trace(f"Error running test: {e}")
+            return RunTestResponse(success=False, response_error=str(e))
+
     # 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)

tests/lsp/test_context.py

@@ -1,5 +1,8 @@
+from pathlib import Path
+
 from sqlmesh.core.context import Context
 from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.uri import URI
 
 
 def test_lsp_context():
@@ -18,3 +21,43 @@ def test_lsp_context():
     # Check that the value is a ModelInfo with the expected model name
     assert isinstance(lsp_context.map[active_customers_key], ModelTarget)
     assert "sushi.active_customers" in lsp_context.map[active_customers_key].names
+
+
+def test_lsp_context_list_workspace_tests():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # List workspace tests
+    tests = lsp_context.list_workspace_tests()
+
+    # Check that the tests are returned correctly
+    assert len(tests) == 0
+    assert any(test.name == "test_order_items" for test in tests)
+
+
+def test_lsp_context_get_deocument_tests():
+    test_path = Path.cwd() / "examples/sushi/tests/test_order_items.yaml"
+    uri = URI.from_path(test_path)
+
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+    tests = lsp_context.get_document_tests(uri)
+
+    assert len(tests) == 1
+    assert tests[0].uri == uri.value
+    assert tests[0].name == "test_order_items"
+
+
+def test_lsp_context_run_test():
+    test_path = Path.cwd() / "examples/sushi/tests/test_order_items.yaml"
+    uri = URI.from_path(test_path)
+
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Run the test
+    result = lsp_context.run_test(uri, "test_order_items")
+
+    # Check that the result is not None and has the expected properties
+    assert result is not None
+    assert result.success is True

vscode/extension/src/extension.ts

@@ -21,6 +21,10 @@ import { selector, completionProvider } from './completion/completion'
 import { LineagePanel } from './webviews/lineagePanel'
 import { RenderedModelProvider } from './providers/renderedModelProvider'
 import { sleep } from './utilities/sleep'
+import {
+  controller as testController,
+  setupTestController,
+} from './tests/tests'
 
 let lspClient: LSPClient | undefined
 
@@ -128,6 +132,7 @@ export async function activate(context: vscode.ExtensionContext) {
         )
       }
       context.subscriptions.push(lspClient)
+      context.subscriptions.push(setupTestController(lspClient))
     } else {
       lspClient = new LSPClient()
       const result = await lspClient.start()
@@ -140,6 +145,7 @@ export async function activate(context: vscode.ExtensionContext) {
         )
       } else {
         context.subscriptions.push(lspClient)
+        context.subscriptions.push(setupTestController(lspClient))
       }
     }
   }
@@ -175,6 +181,8 @@ export async function activate(context: vscode.ExtensionContext) {
     )
   } else {
     context.subscriptions.push(lspClient)
+    setupTestController(lspClient)
+    context.subscriptions.push(testController)
   }
 
   if (lspClient && !lspClient.hasCompletionCapability()) {

vscode/extension/src/lsp/custom.ts

@@ -32,6 +32,9 @@ export type CustomLSPMethods =
   | AllModelsForRenderMethod
   | SupportedMethodsMethod
   | FormatProjectMethod
+  | ListWorkspaceTests
+  | ListDocumentTests
+  | RunTest
 
 interface AllModelsRequest {
   textDocument: {
@@ -111,3 +114,55 @@ 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 extends BaseResponse {
+  tests: TestEntry[]
+}
+
+export interface ListDocumentTests {
+  method: 'sqlmesh/list_document_tests'
+  request: ListDocumentTestsRequest
+  response: ListDocumentTestsResponse
+}
+
+export interface DocumentIdentifier {
+  uri: string
+}
+
+export interface ListDocumentTestsRequest {
+  textDocument: DocumentIdentifier
+}
+
+export interface ListDocumentTestsResponse extends BaseResponse {
+  tests: TestEntry[]
+}
+
+export interface RunTest {
+  method: 'sqlmesh/run_test'
+  request: RunTestRequest
+  response: RunTestResponse
+}
+
+export interface RunTestRequest {
+  textDocument: DocumentIdentifier
+  testName: string
+}
+
+export interface RunTestResponse extends BaseResponse {
+  success: boolean
+  // TODO Add message of why not passed
+}

vscode/extension/src/lsp/lsp.ts

@@ -252,8 +252,8 @@ export class LSPClient implements Disposable {
 
     try {
       const result = await this.client.sendRequest<Response>(method, request)
-      if (result.response_error) {
-        return err(result.response_error)
+      if ((result as any).response_error) {
+        return err((result as any).response_error)
       }
       return ok(result)
     } catch (error) {