This is page 6 of 6. Use http://codebase.md/dbt-labs/dbt-mcp?lines=true&page={x} to view the full context.
# Directory Structure
```
├── .changes
│ ├── header.tpl.md
│ ├── unreleased
│ │ └── .gitkeep
│ ├── v0.1.3.md
│ ├── v0.10.0.md
│ ├── v0.10.1.md
│ ├── v0.10.2.md
│ ├── v0.10.3.md
│ ├── v0.2.0.md
│ ├── v0.2.1.md
│ ├── v0.2.10.md
│ ├── v0.2.11.md
│ ├── v0.2.12.md
│ ├── v0.2.13.md
│ ├── v0.2.14.md
│ ├── v0.2.15.md
│ ├── v0.2.16.md
│ ├── v0.2.17.md
│ ├── v0.2.18.md
│ ├── v0.2.19.md
│ ├── v0.2.2.md
│ ├── v0.2.20.md
│ ├── v0.2.3.md
│ ├── v0.2.4.md
│ ├── v0.2.5.md
│ ├── v0.2.6.md
│ ├── v0.2.7.md
│ ├── v0.2.8.md
│ ├── v0.2.9.md
│ ├── v0.3.0.md
│ ├── v0.4.0.md
│ ├── v0.4.1.md
│ ├── v0.4.2.md
│ ├── v0.5.0.md
│ ├── v0.6.0.md
│ ├── v0.6.1.md
│ ├── v0.6.2.md
│ ├── v0.7.0.md
│ ├── v0.8.0.md
│ ├── v0.8.1.md
│ ├── v0.8.2.md
│ ├── v0.8.3.md
│ ├── v0.8.4.md
│ ├── v0.9.0.md
│ ├── v0.9.1.md
│ ├── v1.0.0.md
│ └── v1.1.0.md
├── .changie.yaml
├── .github
│ ├── actions
│ │ └── setup-python
│ │ └── action.yml
│ ├── CODEOWNERS
│ ├── ISSUE_TEMPLATE
│ │ ├── bug_report.yml
│ │ └── feature_request.yml
│ ├── pull_request_template.md
│ └── workflows
│ ├── changelog-check.yml
│ ├── codeowners-check.yml
│ ├── create-release-pr.yml
│ ├── release.yml
│ └── run-checks-pr.yaml
├── .gitignore
├── .pre-commit-config.yaml
├── .task
│ └── checksum
│ └── d2
├── .tool-versions
├── .vscode
│ ├── launch.json
│ └── settings.json
├── CHANGELOG.md
├── CONTRIBUTING.md
├── docs
│ ├── d2.png
│ └── diagram.d2
├── evals
│ └── semantic_layer
│ └── test_eval_semantic_layer.py
├── examples
│ ├── .DS_Store
│ ├── aws_strands_agent
│ │ ├── __init__.py
│ │ ├── .DS_Store
│ │ ├── dbt_data_scientist
│ │ │ ├── __init__.py
│ │ │ ├── .env.example
│ │ │ ├── agent.py
│ │ │ ├── prompts.py
│ │ │ ├── quick_mcp_test.py
│ │ │ ├── test_all_tools.py
│ │ │ └── tools
│ │ │ ├── __init__.py
│ │ │ ├── dbt_compile.py
│ │ │ ├── dbt_mcp.py
│ │ │ └── dbt_model_analyzer.py
│ │ ├── LICENSE
│ │ ├── README.md
│ │ └── requirements.txt
│ ├── google_adk_agent
│ │ ├── __init__.py
│ │ ├── main.py
│ │ ├── pyproject.toml
│ │ └── README.md
│ ├── langgraph_agent
│ │ ├── __init__.py
│ │ ├── .python-version
│ │ ├── main.py
│ │ ├── pyproject.toml
│ │ ├── README.md
│ │ └── uv.lock
│ ├── openai_agent
│ │ ├── __init__.py
│ │ ├── .gitignore
│ │ ├── .python-version
│ │ ├── main_streamable.py
│ │ ├── main.py
│ │ ├── pyproject.toml
│ │ ├── README.md
│ │ └── uv.lock
│ ├── openai_responses
│ │ ├── __init__.py
│ │ ├── .gitignore
│ │ ├── .python-version
│ │ ├── main.py
│ │ ├── pyproject.toml
│ │ ├── README.md
│ │ └── uv.lock
│ ├── pydantic_ai_agent
│ │ ├── __init__.py
│ │ ├── .gitignore
│ │ ├── .python-version
│ │ ├── main.py
│ │ ├── pyproject.toml
│ │ └── README.md
│ └── remote_mcp
│ ├── .python-version
│ ├── main.py
│ ├── pyproject.toml
│ ├── README.md
│ └── uv.lock
├── LICENSE
├── pyproject.toml
├── README.md
├── src
│ ├── client
│ │ ├── __init__.py
│ │ ├── main.py
│ │ └── tools.py
│ ├── dbt_mcp
│ │ ├── __init__.py
│ │ ├── .gitignore
│ │ ├── config
│ │ │ ├── config_providers.py
│ │ │ ├── config.py
│ │ │ ├── dbt_project.py
│ │ │ ├── dbt_yaml.py
│ │ │ ├── headers.py
│ │ │ ├── settings.py
│ │ │ └── transport.py
│ │ ├── dbt_admin
│ │ │ ├── __init__.py
│ │ │ ├── client.py
│ │ │ ├── constants.py
│ │ │ ├── run_results_errors
│ │ │ │ ├── __init__.py
│ │ │ │ ├── config.py
│ │ │ │ └── parser.py
│ │ │ └── tools.py
│ │ ├── dbt_cli
│ │ │ ├── binary_type.py
│ │ │ └── tools.py
│ │ ├── dbt_codegen
│ │ │ ├── __init__.py
│ │ │ └── tools.py
│ │ ├── discovery
│ │ │ ├── client.py
│ │ │ └── tools.py
│ │ ├── errors
│ │ │ ├── __init__.py
│ │ │ ├── admin_api.py
│ │ │ ├── base.py
│ │ │ ├── cli.py
│ │ │ ├── common.py
│ │ │ ├── discovery.py
│ │ │ ├── semantic_layer.py
│ │ │ └── sql.py
│ │ ├── gql
│ │ │ └── errors.py
│ │ ├── lsp
│ │ │ ├── __init__.py
│ │ │ ├── lsp_binary_manager.py
│ │ │ ├── lsp_client.py
│ │ │ ├── lsp_connection.py
│ │ │ ├── providers
│ │ │ │ ├── __init__.py
│ │ │ │ ├── local_lsp_client_provider.py
│ │ │ │ ├── local_lsp_connection_provider.py
│ │ │ │ ├── lsp_client_provider.py
│ │ │ │ └── lsp_connection_provider.py
│ │ │ └── tools.py
│ │ ├── main.py
│ │ ├── mcp
│ │ │ ├── create.py
│ │ │ └── server.py
│ │ ├── oauth
│ │ │ ├── client_id.py
│ │ │ ├── context_manager.py
│ │ │ ├── dbt_platform.py
│ │ │ ├── fastapi_app.py
│ │ │ ├── logging.py
│ │ │ ├── login.py
│ │ │ ├── refresh_strategy.py
│ │ │ ├── token_provider.py
│ │ │ └── token.py
│ │ ├── prompts
│ │ │ ├── __init__.py
│ │ │ ├── admin_api
│ │ │ │ ├── cancel_job_run.md
│ │ │ │ ├── get_job_details.md
│ │ │ │ ├── get_job_run_artifact.md
│ │ │ │ ├── get_job_run_details.md
│ │ │ │ ├── get_job_run_error.md
│ │ │ │ ├── list_job_run_artifacts.md
│ │ │ │ ├── list_jobs_runs.md
│ │ │ │ ├── list_jobs.md
│ │ │ │ ├── retry_job_run.md
│ │ │ │ └── trigger_job_run.md
│ │ │ ├── dbt_cli
│ │ │ │ ├── args
│ │ │ │ │ ├── full_refresh.md
│ │ │ │ │ ├── limit.md
│ │ │ │ │ ├── resource_type.md
│ │ │ │ │ ├── selectors.md
│ │ │ │ │ ├── sql_query.md
│ │ │ │ │ └── vars.md
│ │ │ │ ├── build.md
│ │ │ │ ├── compile.md
│ │ │ │ ├── docs.md
│ │ │ │ ├── list.md
│ │ │ │ ├── parse.md
│ │ │ │ ├── run.md
│ │ │ │ ├── show.md
│ │ │ │ └── test.md
│ │ │ ├── dbt_codegen
│ │ │ │ ├── args
│ │ │ │ │ ├── case_sensitive_cols.md
│ │ │ │ │ ├── database_name.md
│ │ │ │ │ ├── generate_columns.md
│ │ │ │ │ ├── include_data_types.md
│ │ │ │ │ ├── include_descriptions.md
│ │ │ │ │ ├── leading_commas.md
│ │ │ │ │ ├── materialized.md
│ │ │ │ │ ├── model_name.md
│ │ │ │ │ ├── model_names.md
│ │ │ │ │ ├── schema_name.md
│ │ │ │ │ ├── source_name.md
│ │ │ │ │ ├── table_name.md
│ │ │ │ │ ├── table_names.md
│ │ │ │ │ ├── tables.md
│ │ │ │ │ └── upstream_descriptions.md
│ │ │ │ ├── generate_model_yaml.md
│ │ │ │ ├── generate_source.md
│ │ │ │ └── generate_staging_model.md
│ │ │ ├── discovery
│ │ │ │ ├── get_all_models.md
│ │ │ │ ├── get_all_sources.md
│ │ │ │ ├── get_exposure_details.md
│ │ │ │ ├── get_exposures.md
│ │ │ │ ├── get_mart_models.md
│ │ │ │ ├── get_model_children.md
│ │ │ │ ├── get_model_details.md
│ │ │ │ ├── get_model_health.md
│ │ │ │ └── get_model_parents.md
│ │ │ ├── lsp
│ │ │ │ ├── args
│ │ │ │ │ ├── column_name.md
│ │ │ │ │ └── model_id.md
│ │ │ │ └── get_column_lineage.md
│ │ │ ├── prompts.py
│ │ │ └── semantic_layer
│ │ │ ├── get_dimensions.md
│ │ │ ├── get_entities.md
│ │ │ ├── get_metrics_compiled_sql.md
│ │ │ ├── list_metrics.md
│ │ │ ├── list_saved_queries.md
│ │ │ └── query_metrics.md
│ │ ├── py.typed
│ │ ├── semantic_layer
│ │ │ ├── client.py
│ │ │ ├── gql
│ │ │ │ ├── gql_request.py
│ │ │ │ └── gql.py
│ │ │ ├── levenshtein.py
│ │ │ ├── tools.py
│ │ │ └── types.py
│ │ ├── sql
│ │ │ └── tools.py
│ │ ├── telemetry
│ │ │ └── logging.py
│ │ ├── tools
│ │ │ ├── annotations.py
│ │ │ ├── definitions.py
│ │ │ ├── policy.py
│ │ │ ├── register.py
│ │ │ ├── tool_names.py
│ │ │ └── toolsets.py
│ │ └── tracking
│ │ └── tracking.py
│ └── remote_mcp
│ ├── __init__.py
│ └── session.py
├── Taskfile.yml
├── tests
│ ├── __init__.py
│ ├── conftest.py
│ ├── env_vars.py
│ ├── integration
│ │ ├── __init__.py
│ │ ├── dbt_codegen
│ │ │ ├── __init__.py
│ │ │ └── test_dbt_codegen.py
│ │ ├── discovery
│ │ │ └── test_discovery.py
│ │ ├── initialization
│ │ │ ├── __init__.py
│ │ │ └── test_initialization.py
│ │ ├── lsp
│ │ │ └── test_lsp_connection.py
│ │ ├── remote_mcp
│ │ │ └── test_remote_mcp.py
│ │ ├── remote_tools
│ │ │ └── test_remote_tools.py
│ │ ├── semantic_layer
│ │ │ └── test_semantic_layer.py
│ │ └── tracking
│ │ └── test_tracking.py
│ ├── mocks
│ │ └── config.py
│ └── unit
│ ├── __init__.py
│ ├── config
│ │ ├── __init__.py
│ │ ├── test_config.py
│ │ └── test_transport.py
│ ├── dbt_admin
│ │ ├── __init__.py
│ │ ├── test_client.py
│ │ ├── test_error_fetcher.py
│ │ └── test_tools.py
│ ├── dbt_cli
│ │ ├── __init__.py
│ │ ├── test_cli_integration.py
│ │ └── test_tools.py
│ ├── dbt_codegen
│ │ ├── __init__.py
│ │ └── test_tools.py
│ ├── discovery
│ │ ├── __init__.py
│ │ ├── conftest.py
│ │ ├── test_exposures_fetcher.py
│ │ └── test_sources_fetcher.py
│ ├── lsp
│ │ ├── __init__.py
│ │ ├── test_local_lsp_client_provider.py
│ │ ├── test_local_lsp_connection_provider.py
│ │ ├── test_lsp_client.py
│ │ ├── test_lsp_connection.py
│ │ └── test_lsp_tools.py
│ ├── oauth
│ │ ├── test_credentials_provider.py
│ │ ├── test_fastapi_app_pagination.py
│ │ └── test_token.py
│ ├── semantic_layer
│ │ ├── __init__.py
│ │ └── test_saved_queries.py
│ ├── tools
│ │ ├── test_disable_tools.py
│ │ ├── test_tool_names.py
│ │ ├── test_tool_policies.py
│ │ └── test_toolsets.py
│ └── tracking
│ └── test_tracking.py
├── ui
│ ├── .gitignore
│ ├── assets
│ │ ├── dbt_logo BLK.svg
│ │ └── dbt_logo WHT.svg
│ ├── eslint.config.js
│ ├── index.html
│ ├── package.json
│ ├── pnpm-lock.yaml
│ ├── pnpm-workspace.yaml
│ ├── README.md
│ ├── src
│ │ ├── App.css
│ │ ├── App.tsx
│ │ ├── global.d.ts
│ │ ├── index.css
│ │ ├── main.tsx
│ │ └── vite-env.d.ts
│ ├── tsconfig.app.json
│ ├── tsconfig.json
│ ├── tsconfig.node.json
│ └── vite.config.ts
└── uv.lock
```
# Files
--------------------------------------------------------------------------------
/tests/unit/lsp/test_lsp_connection.py:
--------------------------------------------------------------------------------
```python
1 | """Unit tests for the LSP connection module."""
2 |
3 | import asyncio
4 | import socket
5 | import subprocess
6 | from unittest.mock import AsyncMock, MagicMock, patch
7 |
8 | import pytest
9 |
10 | from dbt_mcp.lsp.lsp_connection import (
11 | SocketLSPConnection,
12 | LspConnectionState,
13 | LspEventName,
14 | JsonRpcMessage,
15 | event_name_from_string,
16 | )
17 |
18 |
19 | class TestJsonRpcMessage:
20 | """Test JsonRpcMessage dataclass."""
21 |
22 | def test_to_dict_with_request(self):
23 | """Test converting a request message to dictionary."""
24 | msg = JsonRpcMessage(id=1, method="initialize", params={"processId": None})
25 | result = msg.to_dict()
26 |
27 | assert result == {
28 | "jsonrpc": "2.0",
29 | "id": 1,
30 | "method": "initialize",
31 | "params": {"processId": None},
32 | }
33 |
34 | def test_to_dict_with_response(self):
35 | """Test converting a response message to dictionary."""
36 | msg = JsonRpcMessage(id=1, result={"capabilities": {}})
37 | result = msg.to_dict()
38 |
39 | assert result == {"jsonrpc": "2.0", "id": 1, "result": {"capabilities": {}}}
40 |
41 | def test_to_dict_with_error(self):
42 | """Test converting an error message to dictionary."""
43 | msg = JsonRpcMessage(
44 | id=1, error={"code": -32601, "message": "Method not found"}
45 | )
46 | result = msg.to_dict()
47 |
48 | assert result == {
49 | "jsonrpc": "2.0",
50 | "id": 1,
51 | "error": {"code": -32601, "message": "Method not found"},
52 | }
53 |
54 | def test_to_dict_notification(self):
55 | """Test converting a notification message to dictionary."""
56 | msg = JsonRpcMessage(
57 | method="window/logMessage", params={"type": 3, "message": "Server started"}
58 | )
59 | result = msg.to_dict()
60 |
61 | assert result == {
62 | "jsonrpc": "2.0",
63 | "method": "window/logMessage",
64 | "params": {"type": 3, "message": "Server started"},
65 | }
66 |
67 | def test_from_dict(self):
68 | """Test creating message from dictionary."""
69 | data = {
70 | "jsonrpc": "2.0",
71 | "id": 42,
72 | "method": "textDocument/completion",
73 | "params": {"textDocument": {"uri": "file:///test.sql"}},
74 | }
75 | msg = JsonRpcMessage(**data)
76 |
77 | assert msg.jsonrpc == "2.0"
78 | assert msg.id == 42
79 | assert msg.method == "textDocument/completion"
80 | assert msg.params == {"textDocument": {"uri": "file:///test.sql"}}
81 |
82 |
83 | class TestLspEventName:
84 | """Test LspEventName enum and helpers."""
85 |
86 | def test_event_name_from_string_valid(self):
87 | """Test converting valid string to event name."""
88 | assert (
89 | event_name_from_string("dbt/lspCompileComplete")
90 | == LspEventName.compileComplete
91 | )
92 | assert event_name_from_string("window/logMessage") == LspEventName.logMessage
93 | assert event_name_from_string("$/progress") == LspEventName.progress
94 |
95 | def test_event_name_from_string_invalid(self):
96 | """Test converting invalid string returns None."""
97 | assert event_name_from_string("invalid/event") is None
98 | assert event_name_from_string("") is None
99 |
100 |
101 | class TestLspConnectionState:
102 | """Test LspConnectionState dataclass."""
103 |
104 | def test_initial_state(self):
105 | """Test initial state values."""
106 | state = LspConnectionState()
107 |
108 | assert state.initialized is False
109 | assert state.shutting_down is False
110 | assert state.capabilities is not None
111 | assert len(state.capabilities) == 0
112 | assert state.pending_requests == {}
113 | assert state.pending_notifications == {}
114 | assert state.compiled is False
115 |
116 | def test_get_next_request_id(self):
117 | """Test request ID generation."""
118 | state = LspConnectionState()
119 |
120 | # Should start at 20 to avoid collisions
121 | id1 = state.get_next_request_id()
122 | id2 = state.get_next_request_id()
123 | id3 = state.get_next_request_id()
124 |
125 | assert id1 == 20
126 | assert id2 == 21
127 | assert id3 == 22
128 |
129 |
130 | class TestLSPConnectionInitialization:
131 | """Test LSP connection initialization and validation."""
132 |
133 | def test_init_valid_binary(self, tmp_path):
134 | """Test initialization with valid binary path."""
135 | # Create a dummy binary file
136 | binary_path = tmp_path / "lsp-server"
137 | binary_path.touch()
138 |
139 | conn = SocketLSPConnection(
140 | binary_path=str(binary_path),
141 | cwd="/test/dir",
142 | args=["--arg1", "--arg2"],
143 | connection_timeout=15,
144 | default_request_timeout=60,
145 | )
146 |
147 | assert conn.binary_path == binary_path
148 | assert conn.cwd == "/test/dir"
149 | assert conn.args == ["--arg1", "--arg2"]
150 | assert conn.host == "127.0.0.1"
151 | assert conn.port == 0
152 | assert conn.connection_timeout == 15
153 | assert conn.default_request_timeout == 60
154 | assert conn.process is None
155 | assert isinstance(conn.state, LspConnectionState)
156 |
157 |
158 | class TestSocketSetup:
159 | """Test socket setup and lifecycle."""
160 |
161 | def test_setup_socket_success(self, tmp_path):
162 | """Test successful socket setup."""
163 | binary_path = tmp_path / "lsp"
164 | binary_path.touch()
165 |
166 | conn = SocketLSPConnection(str(binary_path), "/test")
167 |
168 | with patch("socket.socket") as mock_socket_class:
169 | mock_socket = MagicMock()
170 | mock_socket.getsockname.return_value = ("127.0.0.1", 54321)
171 | mock_socket_class.return_value = mock_socket
172 |
173 | conn.setup_socket()
174 |
175 | # Verify socket setup
176 | mock_socket_class.assert_called_once_with(
177 | socket.AF_INET, socket.SOCK_STREAM
178 | )
179 | mock_socket.setsockopt.assert_called_once_with(
180 | socket.SOL_SOCKET, socket.SO_REUSEADDR, 1
181 | )
182 | mock_socket.bind.assert_called_once_with(("127.0.0.1", 0))
183 | mock_socket.listen.assert_called_once_with(1)
184 |
185 | assert conn.port == 54321
186 | assert conn._socket == mock_socket
187 |
188 |
189 | class TestProcessLaunching:
190 | """Test LSP process launching and termination."""
191 |
192 | @pytest.mark.asyncio
193 | async def test_launch_lsp_process_success(self, tmp_path):
194 | """Test successful process launch."""
195 | binary_path = tmp_path / "lsp"
196 | binary_path.touch()
197 |
198 | conn = SocketLSPConnection(str(binary_path), "/test/dir")
199 | conn.port = 12345
200 |
201 | with patch("asyncio.create_subprocess_exec") as mock_create_subprocess:
202 | mock_process = MagicMock()
203 | mock_process.pid = 9999
204 | mock_create_subprocess.return_value = mock_process
205 |
206 | await conn.launch_lsp_process()
207 |
208 | # Verify process was started with correct arguments
209 | mock_create_subprocess.assert_called_once_with(
210 | str(binary_path), "--socket", "12345", "--project-dir", "/test/dir"
211 | )
212 |
213 | assert conn.process == mock_process
214 |
215 |
216 | class TestStartStop:
217 | """Test start/stop lifecycle."""
218 |
219 | @pytest.mark.asyncio
220 | async def test_start_success(self, tmp_path):
221 | """Test successful server start."""
222 | binary_path = tmp_path / "lsp"
223 | binary_path.touch()
224 |
225 | conn = SocketLSPConnection(str(binary_path), "/test")
226 |
227 | # Mock socket setup
228 | mock_socket = MagicMock()
229 | mock_connection = MagicMock()
230 | mock_socket.getsockname.return_value = ("127.0.0.1", 54321)
231 |
232 | # Mock process
233 | mock_process = MagicMock()
234 | mock_process.pid = 9999
235 |
236 | with (
237 | patch("socket.socket", return_value=mock_socket),
238 | patch("asyncio.create_subprocess_exec", return_value=mock_process),
239 | patch.object(conn, "_read_loop", new_callable=AsyncMock),
240 | patch.object(conn, "_write_loop", new_callable=AsyncMock),
241 | ):
242 | # Mock socket accept
243 | async def mock_accept_wrapper():
244 | return mock_connection, ("127.0.0.1", 12345)
245 |
246 | with patch("asyncio.get_running_loop") as mock_loop:
247 | mock_loop.return_value.run_in_executor.return_value = (
248 | mock_accept_wrapper()
249 | )
250 | mock_loop.return_value.create_task.side_effect = (
251 | lambda coro: asyncio.create_task(coro)
252 | )
253 |
254 | await conn.start()
255 |
256 | assert conn.process == mock_process
257 | assert conn._connection == mock_connection
258 | assert conn._reader_task is not None
259 | assert conn._writer_task is not None
260 |
261 | @pytest.mark.asyncio
262 | async def test_start_already_running(self, tmp_path):
263 | """Test starting when already running."""
264 | binary_path = tmp_path / "lsp"
265 | binary_path.touch()
266 |
267 | conn = SocketLSPConnection(str(binary_path), "/test")
268 | conn.process = MagicMock() # Simulate already running
269 |
270 | with (
271 | patch("socket.socket"),
272 | patch("asyncio.create_subprocess_exec") as mock_create_subprocess,
273 | ):
274 | await conn.start()
275 |
276 | # Should not create a new process
277 | mock_create_subprocess.assert_not_called()
278 |
279 | @pytest.mark.asyncio
280 | async def test_start_timeout(self, tmp_path):
281 | """Test start timeout when server doesn't connect."""
282 | binary_path = tmp_path / "lsp"
283 | binary_path.touch()
284 |
285 | conn = SocketLSPConnection(str(binary_path), "/test", connection_timeout=0.1)
286 |
287 | mock_socket = MagicMock()
288 | mock_socket.getsockname.return_value = ("127.0.0.1", 54321)
289 | mock_process = MagicMock()
290 |
291 | with (
292 | patch("socket.socket", return_value=mock_socket),
293 | patch("asyncio.create_subprocess_exec", return_value=mock_process),
294 | ):
295 | # Simulate timeout in socket.accept
296 | mock_socket.accept.side_effect = TimeoutError
297 |
298 | with patch("asyncio.get_running_loop") as mock_loop:
299 | mock_loop.return_value.run_in_executor.side_effect = TimeoutError
300 |
301 | with pytest.raises(
302 | RuntimeError, match="Timeout waiting for LSP server to connect"
303 | ):
304 | await conn.start()
305 |
306 | @pytest.mark.asyncio
307 | async def test_stop_complete_cleanup(self, tmp_path):
308 | """Test complete cleanup on stop."""
309 | binary_path = tmp_path / "lsp"
310 | binary_path.touch()
311 |
312 | conn = SocketLSPConnection(str(binary_path), "/test")
313 |
314 | # Setup mocks for running state
315 | conn.process = MagicMock()
316 | conn.process.terminate = MagicMock()
317 | conn.process.wait = AsyncMock()
318 | conn.process.kill = MagicMock()
319 |
320 | conn._socket = MagicMock()
321 | conn._connection = MagicMock()
322 |
323 | # Create mock tasks with proper async behavior
324 | async def mock_task():
325 | pass
326 |
327 | conn._reader_task = asyncio.create_task(mock_task())
328 | conn._writer_task = asyncio.create_task(mock_task())
329 |
330 | # Let tasks complete
331 | await asyncio.sleep(0.01)
332 |
333 | # Store references before they are set to None
334 | mock_connection = conn._connection
335 | mock_socket = conn._socket
336 | mock_process = conn.process
337 |
338 | with patch.object(conn, "_send_shutdown_request") as mock_shutdown:
339 | await conn.stop()
340 |
341 | # Verify cleanup methods were called
342 | mock_shutdown.assert_called_once()
343 | mock_connection.close.assert_called_once()
344 | mock_socket.close.assert_called_once()
345 | mock_process.terminate.assert_called_once()
346 |
347 | # Verify everything was set to None
348 | assert conn.process is None
349 | assert conn._socket is None
350 | assert conn._connection is None
351 |
352 | @pytest.mark.asyncio
353 | async def test_stop_force_kill(self, tmp_path):
354 | """Test force kill when process doesn't terminate."""
355 | binary_path = tmp_path / "lsp"
356 | binary_path.touch()
357 |
358 | conn = SocketLSPConnection(str(binary_path), "/test")
359 |
360 | # Setup mock process that doesn't terminate
361 | mock_process = MagicMock()
362 | mock_process.terminate = MagicMock()
363 | mock_process.wait = AsyncMock(
364 | side_effect=[subprocess.TimeoutExpired("cmd", 1), None]
365 | )
366 | mock_process.kill = MagicMock()
367 | conn.process = mock_process
368 |
369 | await conn.stop()
370 |
371 | # Verify force kill was called
372 | mock_process.terminate.assert_called_once()
373 | mock_process.kill.assert_called_once()
374 |
375 |
376 | class TestInitializeMethod:
377 | """Test LSP initialization handshake."""
378 |
379 | @pytest.mark.asyncio
380 | async def test_initialize_success(self, tmp_path):
381 | """Test successful initialization."""
382 | binary_path = tmp_path / "lsp"
383 | binary_path.touch()
384 |
385 | conn = SocketLSPConnection(str(binary_path), "/test")
386 | conn.process = MagicMock() # Simulate running
387 |
388 | # Mock send_request to return capabilities
389 | mock_result = {
390 | "capabilities": {
391 | "textDocumentSync": 2,
392 | "completionProvider": {"triggerCharacters": [".", ":"]},
393 | }
394 | }
395 |
396 | with (
397 | patch.object(
398 | conn, "send_request", new_callable=AsyncMock
399 | ) as mock_send_request,
400 | patch.object(conn, "send_notification") as mock_send_notification,
401 | ):
402 | mock_send_request.return_value = mock_result
403 |
404 | await conn.initialize(timeout=5)
405 |
406 | # Verify initialize request was sent
407 | mock_send_request.assert_called_once()
408 | call_args = mock_send_request.call_args
409 | assert call_args[0][0] == "initialize"
410 | assert call_args[1]["timeout"] == 5
411 |
412 | params = call_args[0][1]
413 | assert params["rootUri"] is None # currently not using cwd
414 | assert params["clientInfo"]["name"] == "dbt-mcp"
415 |
416 | # Verify initialized notification was sent
417 | mock_send_notification.assert_called_once_with("initialized", {})
418 |
419 | # Verify state was updated
420 | assert conn.state.initialized is True
421 | assert conn.state.capabilities == mock_result["capabilities"]
422 |
423 | @pytest.mark.asyncio
424 | async def test_initialize_already_initialized(self, tmp_path):
425 | """Test initialization when already initialized."""
426 | binary_path = tmp_path / "lsp"
427 | binary_path.touch()
428 |
429 | conn = SocketLSPConnection(str(binary_path), "/test")
430 | conn.process = MagicMock()
431 | conn.state.initialized = True
432 |
433 | with pytest.raises(RuntimeError, match="LSP server is already initialized"):
434 | await conn.initialize()
435 |
436 |
437 | class TestMessageParsing:
438 | """Test JSON-RPC message parsing."""
439 |
440 | def test_parse_message_complete(self, tmp_path):
441 | """Test parsing a complete message."""
442 | binary_path = tmp_path / "lsp"
443 | binary_path.touch()
444 |
445 | conn = SocketLSPConnection(str(binary_path), "/test")
446 |
447 | # Create a valid LSP message
448 | content = '{"jsonrpc":"2.0","id":1,"result":{"test":true}}'
449 | header = f"Content-Length: {len(content)}\r\n\r\n"
450 | buffer = (header + content).encode("utf-8")
451 |
452 | message, remaining = conn._parse_message(buffer)
453 |
454 | assert message is not None
455 | assert message.id == 1
456 | assert message.result == {"test": True}
457 | assert remaining == b""
458 |
459 | def test_parse_message_incomplete_header(self, tmp_path):
460 | """Test parsing with incomplete header."""
461 | binary_path = tmp_path / "lsp"
462 | binary_path.touch()
463 |
464 | conn = SocketLSPConnection(str(binary_path), "/test")
465 |
466 | buffer = b"Content-Length: 50\r\n" # Missing \r\n\r\n
467 |
468 | message, remaining = conn._parse_message(buffer)
469 |
470 | assert message is None
471 | assert remaining == buffer
472 |
473 | def test_parse_message_incomplete_content(self, tmp_path):
474 | """Test parsing with incomplete content."""
475 | binary_path = tmp_path / "lsp"
476 | binary_path.touch()
477 |
478 | conn = SocketLSPConnection(str(binary_path), "/test")
479 |
480 | content = '{"jsonrpc":"2.0","id":1,"result":{"test":true}}'
481 | header = f"Content-Length: {len(content)}\r\n\r\n"
482 | # Only include part of the content
483 | buffer = (header + content[:10]).encode("utf-8")
484 |
485 | message, remaining = conn._parse_message(buffer)
486 |
487 | assert message is None
488 | assert remaining == buffer
489 |
490 | def test_parse_message_invalid_json(self, tmp_path):
491 | """Test parsing with invalid JSON content."""
492 | binary_path = tmp_path / "lsp"
493 | binary_path.touch()
494 |
495 | conn = SocketLSPConnection(str(binary_path), "/test")
496 |
497 | content = '{"invalid json'
498 | header = f"Content-Length: {len(content)}\r\n\r\n"
499 | buffer = (header + content).encode("utf-8")
500 |
501 | message, remaining = conn._parse_message(buffer)
502 |
503 | assert message is None
504 | assert remaining == b"" # Invalid message is discarded
505 |
506 | def test_parse_message_missing_content_length(self, tmp_path):
507 | """Test parsing with missing Content-Length header."""
508 | binary_path = tmp_path / "lsp"
509 | binary_path.touch()
510 |
511 | conn = SocketLSPConnection(str(binary_path), "/test")
512 |
513 | buffer = b'Some-Header: value\r\n\r\n{"test":true}'
514 |
515 | message, remaining = conn._parse_message(buffer)
516 |
517 | assert message is None
518 | assert remaining == b'{"test":true}' # Header consumed, content remains
519 |
520 | def test_parse_message_multiple_messages(self, tmp_path):
521 | """Test parsing multiple messages from buffer."""
522 | binary_path = tmp_path / "lsp"
523 | binary_path.touch()
524 |
525 | conn = SocketLSPConnection(str(binary_path), "/test")
526 |
527 | # Create two messages
528 | content1 = '{"jsonrpc":"2.0","id":1,"result":true}'
529 | content2 = '{"jsonrpc":"2.0","id":2,"result":false}'
530 | header1 = f"Content-Length: {len(content1)}\r\n\r\n"
531 | header2 = f"Content-Length: {len(content2)}\r\n\r\n"
532 | buffer = (header1 + content1 + header2 + content2).encode("utf-8")
533 |
534 | # Parse first message
535 | message1, remaining1 = conn._parse_message(buffer)
536 | assert message1 is not None
537 | assert message1.id == 1
538 | assert message1.result is True
539 |
540 | # Parse second message
541 | message2, remaining2 = conn._parse_message(remaining1)
542 | assert message2 is not None
543 | assert message2.id == 2
544 | assert message2.result is False
545 | assert remaining2 == b""
546 |
547 |
548 | class TestMessageHandling:
549 | """Test incoming message handling."""
550 |
551 | def test_handle_response_message(self, tmp_path):
552 | """Test handling response to a request."""
553 | binary_path = tmp_path / "lsp"
554 | binary_path.touch()
555 |
556 | conn = SocketLSPConnection(str(binary_path), "/test")
557 |
558 | # Create a pending request
559 | future = asyncio.Future()
560 | conn.state.pending_requests[42] = future
561 |
562 | # Handle response message
563 | message = JsonRpcMessage(id=42, result={"success": True})
564 |
565 | with patch.object(future, "get_loop") as mock_get_loop:
566 | mock_loop = MagicMock()
567 | mock_get_loop.return_value = mock_loop
568 |
569 | conn._handle_incoming_message(message)
570 |
571 | # Verify future was resolved
572 | mock_loop.call_soon_threadsafe.assert_called_once()
573 | args = mock_loop.call_soon_threadsafe.call_args[0]
574 | assert args[0] == future.set_result
575 | assert args[1] == {"success": True}
576 |
577 | # Verify request was removed from pending
578 | assert 42 not in conn.state.pending_requests
579 |
580 | def test_handle_error_response(self, tmp_path):
581 | """Test handling error response."""
582 | binary_path = tmp_path / "lsp"
583 | binary_path.touch()
584 |
585 | conn = SocketLSPConnection(str(binary_path), "/test")
586 |
587 | # Create a pending request
588 | future = asyncio.Future()
589 | conn.state.pending_requests[42] = future
590 |
591 | # Handle error response
592 | message = JsonRpcMessage(
593 | id=42, error={"code": -32601, "message": "Method not found"}
594 | )
595 |
596 | with patch.object(future, "get_loop") as mock_get_loop:
597 | mock_loop = MagicMock()
598 | mock_get_loop.return_value = mock_loop
599 |
600 | conn._handle_incoming_message(message)
601 |
602 | # Verify future was rejected
603 | mock_loop.call_soon_threadsafe.assert_called_once()
604 | args = mock_loop.call_soon_threadsafe.call_args[0]
605 | assert args[0] == future.set_exception
606 |
607 | def test_handle_unknown_response(self, tmp_path):
608 | """Test handling response for unknown request ID."""
609 | binary_path = tmp_path / "lsp"
610 | binary_path.touch()
611 |
612 | conn = SocketLSPConnection(str(binary_path), "/test")
613 |
614 | # Handle response with unknown ID
615 | message = JsonRpcMessage(id=999, result={"test": True})
616 |
617 | with patch.object(conn, "_send_message") as mock_send:
618 | conn._handle_incoming_message(message)
619 |
620 | # Should send empty response back
621 | mock_send.assert_called_once()
622 | sent_msg = mock_send.call_args[0][0]
623 | assert isinstance(sent_msg, JsonRpcMessage)
624 | assert sent_msg.id == 999
625 | assert sent_msg.result is None
626 |
627 | def test_handle_notification(self, tmp_path):
628 | """Test handling notification messages."""
629 | binary_path = tmp_path / "lsp"
630 | binary_path.touch()
631 |
632 | conn = SocketLSPConnection(str(binary_path), "/test")
633 |
634 | # Create futures waiting for compile complete event
635 | future1 = asyncio.Future()
636 | future2 = asyncio.Future()
637 | conn.state.pending_notifications[LspEventName.compileComplete] = [
638 | future1,
639 | future2,
640 | ]
641 |
642 | # Handle compile complete notification
643 | message = JsonRpcMessage(
644 | method="dbt/lspCompileComplete", params={"success": True}
645 | )
646 |
647 | with (
648 | patch.object(future1, "get_loop") as mock_get_loop1,
649 | patch.object(future2, "get_loop") as mock_get_loop2,
650 | ):
651 | mock_loop1 = MagicMock()
652 | mock_loop2 = MagicMock()
653 | mock_get_loop1.return_value = mock_loop1
654 | mock_get_loop2.return_value = mock_loop2
655 |
656 | conn._handle_incoming_message(message)
657 |
658 | # Verify futures were resolved
659 | mock_loop1.call_soon_threadsafe.assert_called_once_with(
660 | future1.set_result, {"success": True}
661 | )
662 | mock_loop2.call_soon_threadsafe.assert_called_once_with(
663 | future2.set_result, {"success": True}
664 | )
665 |
666 | # Verify compile state was set
667 | assert conn.state.compiled is True
668 |
669 | def test_handle_unknown_notification(self, tmp_path):
670 | """Test handling unknown notification."""
671 | binary_path = tmp_path / "lsp"
672 | binary_path.touch()
673 |
674 | conn = SocketLSPConnection(str(binary_path), "/test")
675 |
676 | # Handle unknown notification
677 | message = JsonRpcMessage(method="unknown/notification", params={"data": "test"})
678 |
679 | # Should not raise, just log
680 | conn._handle_incoming_message(message)
681 |
682 |
683 | class TestSendRequest:
684 | """Test sending requests to LSP server."""
685 |
686 | @pytest.mark.asyncio
687 | async def test_send_request_success(self, tmp_path):
688 | """Test successful request sending."""
689 | binary_path = tmp_path / "lsp"
690 | binary_path.touch()
691 |
692 | conn = SocketLSPConnection(str(binary_path), "/test")
693 | conn.process = MagicMock() # Simulate running
694 |
695 | with (
696 | patch.object(conn, "_send_message") as mock_send,
697 | patch("asyncio.wait_for", new_callable=AsyncMock) as mock_wait_for,
698 | ):
699 | mock_wait_for.return_value = {"result": "success"}
700 |
701 | result = await conn.send_request(
702 | "testMethod", {"param": "value"}, timeout=5
703 | )
704 |
705 | # Verify message was sent
706 | mock_send.assert_called_once()
707 | sent_msg = mock_send.call_args[0][0]
708 | assert isinstance(sent_msg, JsonRpcMessage)
709 | assert sent_msg.method == "testMethod"
710 | assert sent_msg.params == {"param": "value"}
711 | assert sent_msg.id is not None
712 |
713 | # Verify result
714 | assert result == {"result": "success"}
715 |
716 | @pytest.mark.asyncio
717 | async def test_send_request_not_running(self, tmp_path):
718 | """Test sending request when server not running."""
719 | binary_path = tmp_path / "lsp"
720 | binary_path.touch()
721 |
722 | conn = SocketLSPConnection(str(binary_path), "/test")
723 | # process is None - not running
724 |
725 | with pytest.raises(RuntimeError, match="LSP server is not running"):
726 | await conn.send_request("testMethod")
727 |
728 | @pytest.mark.asyncio
729 | async def test_send_request_timeout(self, tmp_path):
730 | """Test request timeout."""
731 | binary_path = tmp_path / "lsp"
732 | binary_path.touch()
733 |
734 | conn = SocketLSPConnection(str(binary_path), "/test", default_request_timeout=1)
735 | conn.process = MagicMock()
736 |
737 | with patch.object(conn, "_send_message"):
738 | # Create a future that never resolves
739 | future = asyncio.Future()
740 | conn.state.pending_requests[20] = future
741 |
742 | # Use real wait_for to test timeout
743 | result = await conn.send_request("testMethod", timeout=0.01)
744 |
745 | assert "error" in result
746 |
747 |
748 | class TestSendNotification:
749 | """Test sending notifications to LSP server."""
750 |
751 | def test_send_notification_success(self, tmp_path):
752 | """Test successful notification sending."""
753 | binary_path = tmp_path / "lsp"
754 | binary_path.touch()
755 |
756 | conn = SocketLSPConnection(str(binary_path), "/test")
757 | conn.process = MagicMock()
758 |
759 | with patch.object(conn, "_send_message") as mock_send:
760 | conn.send_notification(
761 | "window/showMessage", {"type": 3, "message": "Hello"}
762 | )
763 |
764 | # Verify message was sent
765 | mock_send.assert_called_once()
766 | sent_msg = mock_send.call_args[0][0]
767 | assert isinstance(sent_msg, JsonRpcMessage)
768 | assert sent_msg.method == "window/showMessage"
769 | assert sent_msg.params == {"type": 3, "message": "Hello"}
770 | assert sent_msg.id is None # Notifications have no ID
771 |
772 | def test_send_notification_not_running(self, tmp_path):
773 | """Test sending notification when server not running."""
774 | binary_path = tmp_path / "lsp"
775 | binary_path.touch()
776 |
777 | conn = SocketLSPConnection(str(binary_path), "/test")
778 | # process is None - not running
779 |
780 | with pytest.raises(RuntimeError, match="LSP server is not running"):
781 | conn.send_notification("testMethod")
782 |
783 |
784 | class TestWaitForNotification:
785 | """Test waiting for notifications."""
786 |
787 | def test_wait_for_notification(self, tmp_path):
788 | """Test registering to wait for a notification."""
789 | binary_path = tmp_path / "lsp"
790 | binary_path.touch()
791 |
792 | conn = SocketLSPConnection(str(binary_path), "/test")
793 |
794 | with patch("asyncio.get_running_loop") as mock_get_loop:
795 | mock_loop = MagicMock()
796 | mock_future = MagicMock()
797 | mock_loop.create_future.return_value = mock_future
798 | mock_get_loop.return_value = mock_loop
799 |
800 | result = conn.wait_for_notification(LspEventName.compileComplete)
801 |
802 | # Verify future was created and registered
803 | assert result == mock_future
804 | assert LspEventName.compileComplete in conn.state.pending_notifications
805 | assert (
806 | mock_future
807 | in conn.state.pending_notifications[LspEventName.compileComplete]
808 | )
809 |
810 |
811 | class TestSendMessage:
812 | """Test low-level message sending."""
813 |
814 | def test_send_message_with_jsonrpc_message(self, tmp_path):
815 | """Test sending JsonRpcMessage."""
816 | binary_path = tmp_path / "lsp"
817 | binary_path.touch()
818 |
819 | conn = SocketLSPConnection(str(binary_path), "/test")
820 | conn._outgoing_queue = MagicMock()
821 |
822 | message = JsonRpcMessage(id=1, method="test", params={"key": "value"})
823 |
824 | conn._send_message(message)
825 |
826 | # Verify message was queued
827 | conn._outgoing_queue.put_nowait.assert_called_once()
828 | data = conn._outgoing_queue.put_nowait.call_args[0][0]
829 |
830 | # Parse the data to verify format
831 | assert b"Content-Length:" in data
832 | assert b"\r\n\r\n" in data
833 | # JSON might have spaces after colons, check for both variants
834 | assert b'"jsonrpc"' in data and b'"2.0"' in data
835 | assert b'"method"' in data and b'"test"' in data
836 |
837 |
838 | class TestShutdown:
839 | """Test shutdown sequence."""
840 |
841 | def test_send_shutdown_request(self, tmp_path):
842 | """Test sending shutdown and exit messages."""
843 | binary_path = tmp_path / "lsp"
844 | binary_path.touch()
845 |
846 | conn = SocketLSPConnection(str(binary_path), "/test")
847 |
848 | with patch.object(conn, "_send_message") as mock_send:
849 | conn._send_shutdown_request()
850 |
851 | # Verify two messages were sent
852 | assert mock_send.call_count == 2
853 |
854 | # First should be shutdown request
855 | shutdown_msg = mock_send.call_args_list[0][0][0]
856 | assert isinstance(shutdown_msg, JsonRpcMessage)
857 | assert shutdown_msg.method == "shutdown"
858 | assert shutdown_msg.id is not None
859 |
860 | # Second should be exit notification
861 | exit_msg = mock_send.call_args_list[1][0][0]
862 | assert isinstance(exit_msg, JsonRpcMessage)
863 | assert exit_msg.method == "exit"
864 | assert exit_msg.id is None
865 |
866 |
867 | class TestIsRunning:
868 | """Test is_running method."""
869 |
870 | def test_is_running_true(self, tmp_path):
871 | """Test when process is running."""
872 | binary_path = tmp_path / "lsp"
873 | binary_path.touch()
874 |
875 | conn = SocketLSPConnection(str(binary_path), "/test")
876 | conn.process = MagicMock()
877 | conn.process.returncode = None
878 |
879 | assert conn.is_running() is True
880 |
881 | def test_is_running_false_no_process(self, tmp_path):
882 | """Test when no process."""
883 | binary_path = tmp_path / "lsp"
884 | binary_path.touch()
885 |
886 | conn = SocketLSPConnection(str(binary_path), "/test")
887 |
888 | assert conn.is_running() is False
889 |
890 | def test_is_running_false_process_exited(self, tmp_path):
891 | """Test when process has exited."""
892 | binary_path = tmp_path / "lsp"
893 | binary_path.touch()
894 |
895 | conn = SocketLSPConnection(str(binary_path), "/test")
896 | conn.process = MagicMock()
897 | conn.process.returncode = 0
898 |
899 | assert conn.is_running() is False
900 |
901 |
902 | class TestReadWriteLoops:
903 | """Test async I/O loops."""
904 |
905 | @pytest.mark.asyncio
906 | async def test_read_loop_processes_messages(self, tmp_path):
907 | """Test read loop processes incoming messages."""
908 | binary_path = tmp_path / "lsp"
909 | binary_path.touch()
910 |
911 | conn = SocketLSPConnection(str(binary_path), "/test")
912 |
913 | # Setup mock connection
914 | mock_connection = MagicMock()
915 | conn._connection = mock_connection
916 |
917 | # Create test data
918 | content = '{"jsonrpc":"2.0","id":1,"result":true}'
919 | header = f"Content-Length: {len(content)}\r\n\r\n"
920 | test_data = (header + content).encode("utf-8")
921 |
922 | # Mock recv to return data once then empty
923 | recv_calls = [test_data, b""]
924 |
925 | async def mock_recv_wrapper(size):
926 | if recv_calls:
927 | return recv_calls.pop(0)
928 | return b""
929 |
930 | with (
931 | patch("asyncio.get_running_loop") as mock_get_loop,
932 | patch.object(conn, "_handle_incoming_message") as mock_handle,
933 | ):
934 | mock_loop = MagicMock()
935 | mock_get_loop.return_value = mock_loop
936 | mock_loop.run_in_executor.side_effect = (
937 | lambda _, func, *args: mock_recv_wrapper(*args)
938 | )
939 |
940 | # Run read loop (will exit when recv returns empty)
941 | await conn._read_loop()
942 |
943 | # Verify message was handled
944 | mock_handle.assert_called_once()
945 | handled_msg = mock_handle.call_args[0][0]
946 | assert handled_msg.id == 1
947 | assert handled_msg.result is True
948 |
949 | @pytest.mark.asyncio
950 | async def test_write_loop_sends_messages(self, tmp_path):
951 | """Test write loop sends queued messages."""
952 | binary_path = tmp_path / "lsp"
953 | binary_path.touch()
954 |
955 | conn = SocketLSPConnection(str(binary_path), "/test")
956 |
957 | # Setup mock connection
958 | mock_connection = MagicMock()
959 | conn._connection = mock_connection
960 |
961 | # Queue test data
962 | test_data = b"test message data"
963 | conn._outgoing_queue.put_nowait(test_data)
964 |
965 | # Set stop event after first iteration
966 | async def stop_after_one():
967 | await asyncio.sleep(0.01)
968 | conn._stop_event.set()
969 |
970 | with patch("asyncio.get_running_loop") as mock_get_loop:
971 | mock_loop = MagicMock()
972 | mock_get_loop.return_value = mock_loop
973 | mock_loop.run_in_executor.return_value = asyncio.sleep(0)
974 |
975 | # Run both coroutines
976 | await asyncio.gather(
977 | conn._write_loop(), stop_after_one(), return_exceptions=True
978 | )
979 |
980 | # Verify data was sent
981 | mock_loop.run_in_executor.assert_called()
982 | call_args = mock_loop.run_in_executor.call_args_list[-1]
983 | assert call_args[0][1] == mock_connection.sendall
984 | assert call_args[0][2] == test_data
985 |
986 |
987 | class TestEdgeCases:
988 | """Test edge cases and error conditions."""
989 |
990 | @pytest.mark.asyncio
991 | async def test_concurrent_requests(self, tmp_path):
992 | """Test handling concurrent requests."""
993 | binary_path = tmp_path / "lsp"
994 | binary_path.touch()
995 |
996 | conn = SocketLSPConnection(str(binary_path), "/test")
997 | conn.process = MagicMock()
998 |
999 | # Track sent messages
1000 | sent_messages = []
1001 |
1002 | def track_message(msg):
1003 | sent_messages.append(msg)
1004 |
1005 | with patch.object(conn, "_send_message", side_effect=track_message):
1006 | # Create futures for multiple requests
1007 | future1 = asyncio.create_task(
1008 | conn.send_request("method1", JsonRpcMessage(id=1))
1009 | )
1010 | future2 = asyncio.create_task(
1011 | conn.send_request("method2", JsonRpcMessage(id=2))
1012 | )
1013 | future3 = asyncio.create_task(
1014 | conn.send_request("method3", JsonRpcMessage(id=3))
1015 | )
1016 |
1017 | # Let tasks start
1018 | await asyncio.sleep(0.01)
1019 |
1020 | # Verify all messages were sent with unique IDs
1021 | assert len(sent_messages) == 3
1022 | ids = [msg.id for msg in sent_messages]
1023 | assert len(set(ids)) == 3 # All IDs are unique
1024 |
1025 | # Simulate responses
1026 | for msg in sent_messages:
1027 | if msg.id in conn.state.pending_requests:
1028 | future = conn.state.pending_requests[msg.id]
1029 | future.set_result({"response": msg.id})
1030 |
1031 | # Wait for all requests
1032 | results = await asyncio.gather(future1, future2, future3)
1033 |
1034 | # Verify each got correct response
1035 | assert all("response" in r for r in results)
1036 |
1037 | @pytest.mark.asyncio
1038 | async def test_stop_with_pending_requests(self, tmp_path):
1039 | """Test stopping with pending requests."""
1040 | binary_path = tmp_path / "lsp"
1041 | binary_path.touch()
1042 |
1043 | conn = SocketLSPConnection(str(binary_path), "/test")
1044 | conn.process = MagicMock()
1045 | conn.process.terminate = MagicMock()
1046 | conn.process.wait = AsyncMock()
1047 |
1048 | # Add pending requests
1049 | future1 = asyncio.Future()
1050 | future2 = asyncio.Future()
1051 | conn.state.pending_requests[1] = future1
1052 | conn.state.pending_requests[2] = future2
1053 |
1054 | await conn.stop()
1055 |
1056 | # Verify state was cleared
1057 | assert len(conn.state.pending_requests) == 0
1058 |
1059 | def test_message_with_unicode(self, tmp_path):
1060 | """Test handling messages with unicode content."""
1061 | binary_path = tmp_path / "lsp"
1062 | binary_path.touch()
1063 |
1064 | conn = SocketLSPConnection(str(binary_path), "/test")
1065 |
1066 | # Create message with unicode
1067 | content = '{"jsonrpc":"2.0","method":"test","params":{"text":"Hello 世界 🚀"}}'
1068 | header = f"Content-Length: {len(content.encode('utf-8'))}\r\n\r\n"
1069 | buffer = header.encode("utf-8") + content.encode("utf-8")
1070 |
1071 | message, remaining = conn._parse_message(buffer)
1072 |
1073 | assert message is not None
1074 | assert message.method == "test"
1075 | assert message.params["text"] == "Hello 世界 🚀"
1076 | assert remaining == b""
1077 |
```