Add CI test workflow and schema-aligned test mocks

.github/workflows/test.yml

@@ -0,0 +1,31 @@
+name: Run Tests and Lint
+
+on:
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build-and-test:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+
+      - name: Install dependencies
+        run: |
+          pip install --upgrade pip
+          pip install -e .[dev]
+
+      - name: Lint with ruff
+        run: |
+          ruff check .
+
+      - name: Run tests
+        run: |
+          pip install pytest
+          pytest --maxfail=1 --disable-warnings -q

lsproxy/auth.py

@@ -12,10 +12,20 @@ def base64url_encode(data):
 
     padding = b"="
     encoded = base64.b64encode(data).replace(b"+", b"-").replace(b"/", b"_")
-    return encoded.rstrip(padding)
+    return encoded.rstrip(padding).decode("utf-8")
 
 
 def create_jwt(payload, secret):
+    # Validate inputs
+    if not isinstance(payload, dict):
+        raise TypeError("Payload must be a dictionary")
+    if not payload:
+        raise ValueError("Payload cannot be empty")
+    if not isinstance(secret, str):
+        raise TypeError("Secret must be a string")
+    if not secret:
+        raise ValueError("Secret cannot be empty")
+
     # Create JWT header
     header = {"typ": "JWT", "alg": "HS256"}
 
@@ -24,10 +34,10 @@ def create_jwt(payload, secret):
     encoded_payload = base64url_encode(payload)
 
     # Create signature
-    signing_input = encoded_header + b"." + encoded_payload
-    signature = hmac.new(secret.encode("utf-8"), signing_input, hashlib.sha256).digest()
+    signing_input = f"{encoded_header}.{encoded_payload}"
+    signature = hmac.new(secret.encode("utf-8"), signing_input.encode("utf-8"), hashlib.sha256).digest()
     encoded_signature = base64url_encode(signature)
 
     # Combine all parts
-    jwt = signing_input + b"." + encoded_signature
-    return jwt.decode("utf-8")
+    jwt = f"{signing_input}.{encoded_signature}"
+    return jwt

lsproxy/client.py

@@ -5,7 +5,7 @@
 
 # Only import type hints for Modal if type checking
 if TYPE_CHECKING:
-    import modal
+    import modal  # noqa: F401
 
 from .models import (
     DefinitionResponse,
@@ -37,16 +37,27 @@ def __init__(
         timeout: float = 10.0,
         auth_token: Optional[str] = None,
     ):
+        if auth_token == "":
+            raise ValueError("Token cannot be empty")
+        if auth_token is None:
+            raise ValueError("Token cannot be None")
+
         self._client.base_url = base_url
         self._client.timeout = timeout
         headers = {"Content-Type": "application/json"}
-        if auth_token:
-            headers["Authorization"] = f"Bearer {auth_token}"
+        headers["Authorization"] = f"Bearer {auth_token}"
         self._client.headers = headers
 
     def _request(self, method: str, endpoint: str, **kwargs) -> httpx.Response:
         """Make HTTP request with retry logic and better error handling."""
         try:
+            # Ensure headers from client are included in the request
+            if "headers" in kwargs:
+                headers = {**self._client.headers, **kwargs["headers"]}
+            else:
+                headers = self._client.headers
+            kwargs["headers"] = headers
+
             response = self._client.request(method, endpoint, **kwargs)
             response.raise_for_status()
             return response
@@ -104,7 +115,7 @@ def read_source_code(self, request: FileRange) -> ReadSourceCodeResponse:
                 f"Expected FileRange, got {type(request).__name__}. Please use FileRange model to construct the request."
             )
         response = self._request(
-            "POST", "/workspace/read-source-code", json=request.model_dump()
+            "POST", "/workspace/read-source-code", json={"range": request.model_dump()}
         )
         return ReadSourceCodeResponse.model_validate_json(response.text)
 
@@ -243,14 +254,17 @@ def initialize_with_modal(
 
         return client
 
-    def check_health(self) -> bool:
-        """Check if the server is healthy and ready."""
+    def check_health(self) -> dict:
+        """Check if the server is healthy and ready.
+
+        Returns:
+            dict: Health check response containing status and supported languages
+        """
         try:
             response = self._request("GET", "/system/health")
-            health_data = response.json()
-            return health_data.get("status") == "ok"
+            return response.json()
         except Exception:
-            return False
+            return {"status": "error"}
 
     def close(self):
         """Close the HTTP client and cleanup Modal resources if present."""

pyproject.toml

@@ -29,6 +29,7 @@ include = ["lsproxy*"]
 [project.optional-dependencies]
 dev = [
     "ruff>=0.3.7",
+    "PyJWT>=2.8.0",
 ]
 modal = [
     "modal>=0.56.4",

tests/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for the lsproxy package."""

tests/test_auth.py

@@ -0,0 +1,72 @@
+"""Unit tests for authentication utilities."""
+import pytest
+import jwt
+from datetime import datetime, timedelta, timezone
+
+from lsproxy.auth import create_jwt, base64url_encode
+
+
+@pytest.fixture
+def sample_payload():
+    """Create a sample JWT payload."""
+    return {
+        "sub": "test-user",
+        "exp": int((datetime.now(timezone.utc) + timedelta(hours=1)).timestamp())
+    }
+
+
+@pytest.fixture
+def sample_secret():
+    """Create a sample secret key."""
+    return "test-secret-key-1234"
+
+
+def test_base64url_encode():
+    """Test base64url encoding."""
+    # Test basic encoding
+    assert base64url_encode(b"test") == "dGVzdA"
+
+    # Test padding removal
+    assert base64url_encode(b"t") == "dA"
+    assert base64url_encode(b"te") == "dGU"
+    assert base64url_encode(b"tes") == "dGVz"
+
+    # Test URL-safe characters
+    assert "+" not in base64url_encode(b"???")
+    assert "/" not in base64url_encode(b"???")
+
+
+def test_create_jwt(sample_payload, sample_secret):
+    """Test JWT creation."""
+    token = create_jwt(sample_payload, sample_secret)
+
+    # Verify token structure
+    assert isinstance(token, str)
+    assert len(token.split(".")) == 3
+
+    # Verify token can be decoded
+    decoded = jwt.decode(token, sample_secret, algorithms=["HS256"])
+    assert decoded["sub"] == sample_payload["sub"]
+    assert decoded["exp"] == sample_payload["exp"]
+
+
+
+
+def test_create_jwt_invalid_payload():
+    """Test JWT creation with invalid payload."""
+    with pytest.raises(TypeError):
+        create_jwt("not a dict", "secret")
+
+    with pytest.raises(ValueError):
+        create_jwt({}, "secret")  # Empty payload
+
+
+def test_create_jwt_invalid_secret():
+    """Test JWT creation with invalid secret."""
+    payload = {"sub": "test"}
+
+    with pytest.raises(ValueError):
+        create_jwt(payload, "")  # Empty secret
+
+    with pytest.raises(TypeError):
+        create_jwt(payload, None)  # None secret