LEANN/tests/test_lmstudio_bridge.py

"""Unit tests for LM Studio TypeScript SDK bridge functionality.

This test suite defines the contract for the LM Studio SDK bridge that queries
model context length via Node.js subprocess. These tests verify:

1. Successful SDK query returns context length
2. Graceful fallback when Node.js not installed (FileNotFoundError)
3. Graceful fallback when SDK not installed (npm error)
4. Timeout handling (subprocess.TimeoutExpired)
5. Invalid JSON response handling

All tests are written in Red Phase - they should FAIL initially because the
`_query_lmstudio_context_limit` function does not exist yet.

The function contract:
- Inputs: model_name (str), base_url (str, WebSocket format "ws://localhost:1234")
- Outputs: context_length (int) or None on error
- Requirements:
  1. Call Node.js with inline JavaScript using @lmstudio/sdk
  2. 10-second timeout (accounts for Node.js startup)
  3. Graceful fallback on any error (returns None, doesn't raise)
  4. Parse JSON response with contextLength field
  5. Log errors at debug level (not warning/error)
"""

import subprocess
from unittest.mock import Mock

import pytest

# Try to import the function - if it doesn't exist, tests will fail as expected
try:
    from leann.embedding_compute import _query_lmstudio_context_limit
except ImportError:
    # Function doesn't exist yet (Red Phase) - create a placeholder that will fail
    def _query_lmstudio_context_limit(*args, **kwargs):
        raise NotImplementedError(
            "_query_lmstudio_context_limit not implemented yet - this is the Red Phase"
        )


class TestLMStudioBridge:
    """Tests for LM Studio TypeScript SDK bridge integration."""

    def test_query_lmstudio_success(self, monkeypatch):
        """Verify successful SDK query returns context length.

        When the Node.js subprocess successfully queries the LM Studio SDK,
        it should return a JSON response with contextLength field. The function
        should parse this and return the integer context length.
        """

        def mock_run(*args, **kwargs):
            # Verify timeout is set to 10 seconds
            assert kwargs.get("timeout") == 10, "Should use 10-second timeout for Node.js startup"

            # Verify capture_output and text=True are set
            assert kwargs.get("capture_output") is True, "Should capture stdout/stderr"
            assert kwargs.get("text") is True, "Should decode output as text"

            # Return successful JSON response
            mock_result = Mock()
            mock_result.returncode = 0
            mock_result.stdout = '{"contextLength": 8192, "identifier": "custom-model"}'
            mock_result.stderr = ""
            return mock_result

        monkeypatch.setattr("subprocess.run", mock_run)

        # Test with typical LM Studio model
        limit = _query_lmstudio_context_limit(
            model_name="custom-model", base_url="ws://localhost:1234"
        )

        assert limit == 8192, "Should return context length from SDK response"

    def test_query_lmstudio_nodejs_not_found(self, monkeypatch):
        """Verify graceful fallback when Node.js not installed.

        When Node.js is not installed, subprocess.run will raise FileNotFoundError.
        The function should catch this and return None (graceful fallback to registry).
        """

        def mock_run(*args, **kwargs):
            raise FileNotFoundError("node: command not found")

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="custom-model", base_url="ws://localhost:1234"
        )

        assert limit is None, "Should return None when Node.js not installed"

    def test_query_lmstudio_sdk_not_installed(self, monkeypatch):
        """Verify graceful fallback when @lmstudio/sdk not installed.

        When the SDK npm package is not installed, Node.js will return non-zero
        exit code with error message in stderr. The function should detect this
        and return None.
        """

        def mock_run(*args, **kwargs):
            mock_result = Mock()
            mock_result.returncode = 1
            mock_result.stdout = ""
            mock_result.stderr = (
                "Error: Cannot find module '@lmstudio/sdk'\nRequire stack:\n- /path/to/script.js"
            )
            return mock_result

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="custom-model", base_url="ws://localhost:1234"
        )

        assert limit is None, "Should return None when SDK not installed"

    def test_query_lmstudio_timeout(self, monkeypatch):
        """Verify graceful fallback when subprocess times out.

        When the Node.js process takes longer than 10 seconds (e.g., LM Studio
        not responding), subprocess.TimeoutExpired should be raised. The function
        should catch this and return None.
        """

        def mock_run(*args, **kwargs):
            raise subprocess.TimeoutExpired(cmd=["node", "lmstudio_bridge.js"], timeout=10)

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="custom-model", base_url="ws://localhost:1234"
        )

        assert limit is None, "Should return None on timeout"

    def test_query_lmstudio_invalid_json(self, monkeypatch):
        """Verify graceful fallback when response is invalid JSON.

        When the subprocess returns malformed JSON (e.g., due to SDK error),
        json.loads will raise ValueError/JSONDecodeError. The function should
        catch this and return None.
        """

        def mock_run(*args, **kwargs):
            mock_result = Mock()
            mock_result.returncode = 0
            mock_result.stdout = "This is not valid JSON"
            mock_result.stderr = ""
            return mock_result

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="custom-model", base_url="ws://localhost:1234"
        )

        assert limit is None, "Should return None when JSON parsing fails"

    def test_query_lmstudio_missing_context_length_field(self, monkeypatch):
        """Verify graceful fallback when JSON lacks contextLength field.

        When the SDK returns valid JSON but without the expected contextLength
        field (e.g., error response), the function should return None.
        """

        def mock_run(*args, **kwargs):
            mock_result = Mock()
            mock_result.returncode = 0
            mock_result.stdout = '{"identifier": "test-model", "error": "Model not found"}'
            mock_result.stderr = ""
            return mock_result

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="nonexistent-model", base_url="ws://localhost:1234"
        )

        assert limit is None, "Should return None when contextLength field missing"

    def test_query_lmstudio_null_context_length(self, monkeypatch):
        """Verify graceful fallback when contextLength is null.

        When the SDK returns contextLength: null (model couldn't be loaded),
        the function should return None for registry fallback.
        """

        def mock_run(*args, **kwargs):
            mock_result = Mock()
            mock_result.returncode = 0
            mock_result.stdout = '{"contextLength": null, "identifier": "test-model"}'
            mock_result.stderr = ""
            return mock_result

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="test-model", base_url="ws://localhost:1234"
        )

        assert limit is None, "Should return None when contextLength is null"

    def test_query_lmstudio_zero_context_length(self, monkeypatch):
        """Verify graceful fallback when contextLength is zero.

        When the SDK returns contextLength: 0 (invalid value), the function
        should return None to trigger registry fallback.
        """

        def mock_run(*args, **kwargs):
            mock_result = Mock()
            mock_result.returncode = 0
            mock_result.stdout = '{"contextLength": 0, "identifier": "test-model"}'
            mock_result.stderr = ""
            return mock_result

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="test-model", base_url="ws://localhost:1234"
        )

        assert limit is None, "Should return None when contextLength is zero"

    def test_query_lmstudio_with_custom_port(self, monkeypatch):
        """Verify SDK query works with non-default WebSocket port.

        LM Studio can run on custom ports. The function should pass the
        provided base_url to the Node.js subprocess.
        """

        def mock_run(*args, **kwargs):
            # Verify the base_url argument is passed correctly
            command = args[0] if args else kwargs.get("args", [])
            assert "ws://localhost:8080" in " ".join(command), (
                "Should pass custom port to subprocess"
            )

            mock_result = Mock()
            mock_result.returncode = 0
            mock_result.stdout = '{"contextLength": 4096, "identifier": "custom-model"}'
            mock_result.stderr = ""
            return mock_result

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="custom-model", base_url="ws://localhost:8080"
        )

        assert limit == 4096, "Should work with custom WebSocket port"

    @pytest.mark.parametrize(
        "context_length,expected",
        [
            (512, 512),  # Small context
            (2048, 2048),  # Common context
            (8192, 8192),  # Large context
            (32768, 32768),  # Very large context
        ],
    )
    def test_query_lmstudio_various_context_lengths(self, monkeypatch, context_length, expected):
        """Verify SDK query handles various context length values.

        Different models have different context lengths. The function should
        correctly parse and return any positive integer value.
        """

        def mock_run(*args, **kwargs):
            mock_result = Mock()
            mock_result.returncode = 0
            mock_result.stdout = f'{{"contextLength": {context_length}, "identifier": "test"}}'
            mock_result.stderr = ""
            return mock_result

        monkeypatch.setattr("subprocess.run", mock_run)

        limit = _query_lmstudio_context_limit(
            model_name="test-model", base_url="ws://localhost:1234"
        )

        assert limit == expected, f"Should return {expected} for context length {context_length}"

    def test_query_lmstudio_logs_at_debug_level(self, monkeypatch, caplog):
        """Verify errors are logged at DEBUG level, not WARNING/ERROR.

        Following the graceful fallback pattern from Ollama implementation,
        errors should be logged at debug level to avoid alarming users when
        fallback to registry works fine.
        """
        import logging

        caplog.set_level(logging.DEBUG, logger="leann.embedding_compute")

        def mock_run(*args, **kwargs):
            raise FileNotFoundError("node: command not found")

        monkeypatch.setattr("subprocess.run", mock_run)

        _query_lmstudio_context_limit(model_name="test-model", base_url="ws://localhost:1234")

        # Check that debug logging occurred (not warning/error)
        debug_logs = [record for record in caplog.records if record.levelname == "DEBUG"]
        assert len(debug_logs) > 0, "Should log error at DEBUG level"

        # Verify no WARNING or ERROR logs
        warning_or_error_logs = [
            record for record in caplog.records if record.levelname in ["WARNING", "ERROR"]
        ]
        assert len(warning_or_error_logs) == 0, (
            "Should not log at WARNING/ERROR level for expected failures"
        )