From e00f17927c5094af5c4a2ef94cdd886b084a9f65 Mon Sep 17 00:00:00 2001
From: Alex Carney <alcarneyme@gmail.com>
Date: Mon, 7 Oct 2024 20:57:41 +0100
Subject: [PATCH] test: add the option to run the end-to-end tests over TCP

This extends pygls' test suite with a `--lsp-transport` command line
argument.

- `--lsp-transport stdio` (the default), runs the end-to-end tests
  over stdin/stdout.
- `--lsp-transport tcp`, runs the end-to-end tests over a TCP
  connection

The `poe test` task used in CI has been updated to run both style of
end-to-end tests.
---
 pyproject.toml    |  1 +
 tests/conftest.py | 49 ++++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7aeb4247..a8b6ea03 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -67,6 +67,7 @@ poetry_lock_check = "poetry check"
 [tool.poe.tasks.test]
 sequence = [
      { cmd = "pytest --cov" },
+     { cmd = "pytest tests/e2e --lsp-transport tcp" },
 ]
 ignore_fail = "return_non_zero"
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 2e89d2eb..74c661ad 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -16,6 +16,8 @@
 # See the License for the specific language governing permissions and      #
 # limitations under the License.                                           #
 ############################################################################
+from __future__ import annotations
+
 import asyncio
 import pathlib
 import sys
@@ -126,6 +128,15 @@ def pytest_addoption(parser):
         help="Choose the runtime in which to run servers under test.",
     )
 
+    group.addoption(
+        "--lsp-transport",
+        dest="lsp_transport",
+        action="store",
+        default="stdio",
+        choices=("stdio", "tcp"),
+        help="Choose the transport to use with servers under test.",
+    )
+
 
 @pytest.fixture(scope="session")
 def runtime(request):
@@ -134,6 +145,13 @@ def runtime(request):
     return request.config.getoption("lsp_runtime")
 
 
+@pytest.fixture(scope="session")
+def transport(request):
+    """This fixture is the source of truth for the transport we should run the
+    end-to-end tests with."""
+    return request.config.getoption("lsp_transport")
+
+
 @pytest.fixture(scope="session")
 def path_for():
     """Returns the path corresponding to a file in the example workspace"""
@@ -171,14 +189,31 @@ def server_dir():
     return path.resolve()
 
 
-def get_client_for_cpython_server(uri_fixture):
+def get_client_for_cpython_server(transport, uri_fixture):
     """Return a client configured to communicate with a server running under cpython."""
 
     async def fn(
         server_name: str, capabilities: Optional[types.ClientCapabilities] = None
     ):
         client = LanguageClient("pygls-test-suite", "v1")
-        await client.start_io(sys.executable, str(SERVER_DIR / server_name))
+
+        server_cmd = [sys.executable, str(SERVER_DIR / server_name)]
+        server: asyncio.subprocess.Process | None = None
+
+        if transport == "stdio":
+            await client.start_io(*server_cmd)
+
+        elif transport == "tcp":
+            # TODO: Make host/port configurable?
+            host, port = 'localhost', 8888
+            server_cmd.extend(["--tcp", "--host", host, "--port", f"{port}"])
+
+            server = await asyncio.create_subprocess_exec(*server_cmd)
+            await asyncio.sleep(1)
+            await client.start_tcp(host, port)
+
+        else:
+            raise NotImplementedError(f"Unsupported transport: {transport!r}")
 
         response = await client.initialize_async(
             types.InitializeParams(
@@ -193,19 +228,19 @@ async def fn(
         client.exit(None)
 
         await client.stop()
+        if server is not None and server.returncode is None:
+            server.terminate()
 
     return fn
 
 
 @pytest.fixture(scope="session")
-def get_client_for(runtime, uri_for):
+def get_client_for(runtime, transport, uri_for):
     """Return a client configured to communicate with the specified server.
 
-    Takes into account the current runtime.
-
-    It's the consuming fixture's responsibility to stop the client.
+    Takes into account the current runtime and transport.
     """
     if runtime not in {"cpython"}:
         raise NotImplementedError(f"get_client_for: {runtime=}")
 
-    return get_client_for_cpython_server(uri_for)
+    return get_client_for_cpython_server(transport, uri_for)