milvus-io--milvus/tests/python_client/common/mock_tei_server.py

"""
Mock TEI (Text Embeddings Inference) Server for testing.

This module provides utilities to mock TEI API using pytest-httpserver.
It can be used to test scenarios where the embedding service becomes unavailable
after a collection function has been created.

TEI API Reference:
- POST /embed: Generate embeddings for input texts
  - Request: {"inputs": ["text1", "text2"], "truncate": true, "truncation_direction": "Left"}
  - Response: [[0.1, 0.2, ...], [0.3, 0.4, ...]]

Usage with pytest-httpserver (recommended):
    @pytest.fixture
    def mock_tei(httpserver):
        return MockTEIHandler(httpserver, dim=768)

    def test_example(mock_tei):
        mock_tei.setup_embed()
        endpoint = mock_tei.endpoint
        # use endpoint...

        # Simulate error
        mock_tei.setup_error(503, "Service unavailable")

Usage with standalone server (for environments without pytest-httpserver):
    server = MockTEIServer(dim=768)
    server.start()
    endpoint = server.endpoint
    server.set_error_mode(True)
    server.stop()
"""

import json
import threading
import time
from http.server import HTTPServer, BaseHTTPRequestHandler
from typing import Optional, Callable, Any
import socket


def generate_mock_embedding(text: str, dim: int) -> list:
    """Generate a deterministic mock embedding based on text content."""
    hash_val = hash(text) & 0xFFFFFFFF
    embedding = []
    for i in range(dim):
        val = ((hash_val * (i + 1)) % 10000) / 10000.0 * 2 - 1
        embedding.append(round(val, 6))
    return embedding


# =============================================================================
# pytest-httpserver based implementation (recommended)
# =============================================================================

class MockTEIHandler:
    """
    TEI mock handler for pytest-httpserver.

    This is the recommended way to mock TEI in pytest tests.

    Example:
        def test_with_tei(httpserver):
            tei = MockTEIHandler(httpserver, dim=768)
            tei.setup_embed()

            # Your test code using tei.endpoint
            ...

            # Simulate service failure
            tei.setup_error(503, "Model integration is not active")
    """

    def __init__(self, httpserver, dim: int = 768):
        """
        Initialize TEI handler.

        Args:
            httpserver: pytest-httpserver's HTTPServer fixture
            dim: Embedding dimension
        """
        self.httpserver = httpserver
        self.dim = dim

    @property
    def endpoint(self) -> str:
        """Get the server endpoint URL."""
        return self.httpserver.url_for("")

    def setup_embed(self):
        """Setup /embed endpoint to return mock embeddings."""
        def handle_embed(request):
            data = request.json
            inputs = data.get("inputs", [])
            embeddings = [generate_mock_embedding(text, self.dim) for text in inputs]
            return json.dumps(embeddings)

        self.httpserver.expect_request(
            "/embed",
            method="POST"
        ).respond_with_handler(handle_embed)

        return self

    def setup_error(self, status_code: int = 500, message: str = "Service unavailable"):
        """
        Setup server to return errors for all requests.

        Args:
            status_code: HTTP status code
            message: Error message
        """
        self.httpserver.clear()

        error_response = json.dumps({"error": message})

        self.httpserver.expect_request(
            "/embed",
            method="POST"
        ).respond_with_data(
            error_response,
            status=status_code,
            content_type="application/json"
        )

        return self

    def setup_health(self):
        """Setup /health endpoint."""
        self.httpserver.expect_request(
            "/health",
            method="GET"
        ).respond_with_json({"status": "ok"})

        return self

    def clear(self):
        """Clear all handlers."""
        self.httpserver.clear()
        return self


# =============================================================================
# Standalone server implementation (fallback for environments without pytest-httpserver)
# =============================================================================

def create_handler_class(server_state: dict):
    """Create a handler class with instance-specific state."""

    class _StandaloneHandler(BaseHTTPRequestHandler):
        """HTTP request handler for standalone mock TEI server."""

        def log_message(self, format, *args):
            pass

        def _send_json(self, data, status: int = 200):
            self.send_response(status)
            self.send_header('Content-Type', 'application/json')
            self.end_headers()
            self.wfile.write(json.dumps(data).encode('utf-8'))

        def do_POST(self):
            if server_state.get('error_mode', False):
                self._send_json(
                    {"error": server_state.get('error_message', 'Service unavailable')},
                    server_state.get('error_status_code', 500)
                )
                return

            if self.path == '/embed':
                content_length = int(self.headers.get('Content-Length', 0))
                body = json.loads(self.rfile.read(content_length).decode('utf-8'))
                inputs = body.get('inputs', [])
                dim = server_state.get('dim', 768)
                embeddings = [generate_mock_embedding(text, dim) for text in inputs]
                self._send_json(embeddings)
            else:
                self._send_json({"error": "Not found"}, 404)

        def do_GET(self):
            if server_state.get('error_mode', False):
                self._send_json(
                    {"error": server_state.get('error_message', 'Service unavailable')},
                    server_state.get('error_status_code', 500)
                )
                return

            if self.path == '/health':
                self._send_json({"status": "ok"})
            else:
                self._send_json({"error": "Not found"}, 404)

    return _StandaloneHandler


def get_local_ip() -> str:
    """
    Get the local IP address that can be accessed from external hosts.
    Returns the first non-loopback IPv4 address.
    """
    # Method 1: get IP from socket connection to external host
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
        s.close()
        if not ip.startswith('127.'):
            return ip
    except (OSError, socket.error):
        # Socket connection failed, try next method
        pass

    # Method 2: get from hostname
    try:
        hostname = socket.gethostname()
        ip = socket.gethostbyname(hostname)
        if not ip.startswith('127.'):
            return ip
    except (OSError, socket.error, socket.gaierror):
        # Hostname resolution failed, fall back to localhost
        pass

    return '127.0.0.1'


def get_docker_host() -> str:
    """
    Get the hostname that Docker containers can use to access the host machine.

    - macOS/Windows Docker Desktop: host.docker.internal
    - Linux: returns the host's IP address (containers need --add-host or host network)
    """
    import platform
    system = platform.system().lower()

    if system in ('darwin', 'windows'):
        # Docker Desktop provides this special DNS name
        return 'host.docker.internal'
    else:
        # Linux: use host IP
        return get_local_ip()


class MockTEIServer:
    """
    Standalone mock TEI server.

    Use this when pytest-httpserver is not available.
    For pytest tests, prefer using MockTEIHandler with httpserver fixture.

    Example:
        with MockTEIServer(dim=768) as server:
            endpoint = server.endpoint
            # use endpoint...

            server.set_error_mode(True, 503, "Service unavailable")

    For remote Milvus access, use external_host parameter:
        # Auto-detect external IP
        server = MockTEIServer(dim=768, host='0.0.0.0', external_host='auto')

        # For Docker container access (macOS/Windows)
        server = MockTEIServer(dim=768, host='0.0.0.0', external_host='docker')

        # Or specify explicit IP
        server = MockTEIServer(dim=768, host='0.0.0.0', external_host='192.168.1.100')
    """

    def __init__(self, port: int = 0, dim: int = 768, host: str = '127.0.0.1', external_host: str = None):
        self.host = host
        self.port = port
        self.dim = dim
        self._external_host = external_host
        self._server: Optional[HTTPServer] = None
        self._thread: Optional[threading.Thread] = None
        self._running = False
        # Instance-specific state (not shared between servers)
        self._state = {
            'dim': dim,
            'error_mode': False,
            'error_status_code': 500,
            'error_message': 'Service unavailable'
        }

    @property
    def endpoint(self) -> str:
        if self._server is None:
            raise RuntimeError("Server not started")
        # Use external_host for endpoint URL if specified
        if self._external_host:
            if self._external_host == 'auto':
                host = get_local_ip()
            elif self._external_host == 'docker':
                host = get_docker_host()
            else:
                host = self._external_host
        else:
            host = self.host
        return f"http://{host}:{self._server.server_address[1]}"

    def start(self) -> str:
        if self._running:
            return self.endpoint

        # Create handler class with instance-specific state
        handler_class = create_handler_class(self._state)

        self._server = HTTPServer((self.host, self.port), handler_class)
        self.port = self._server.server_address[1]

        self._thread = threading.Thread(target=self._server.serve_forever)
        self._thread.daemon = True
        self._thread.start()
        self._running = True

        self._wait_for_server()
        return self.endpoint

    def _wait_for_server(self, timeout: float = 5.0):
        start = time.time()
        while time.time() - start < timeout:
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(1)
                if sock.connect_ex((self.host, self.port)) == 0:
                    sock.close()
                    return
                sock.close()
            except (OSError, socket.error):
                # Connection not ready yet, retry
                pass
            time.sleep(0.1)
        raise RuntimeError(f"Server failed to start within {timeout}s")

    def stop(self):
        if self._server:
            self._server.shutdown()
            self._server.server_close()
            self._server = None
        if self._thread:
            self._thread.join(timeout=10)
            self._thread = None
        self._running = False

    def set_error_mode(self, enabled: bool, status_code: int = 500, message: str = "Service unavailable"):
        self._state['error_mode'] = enabled
        self._state['error_status_code'] = status_code
        self._state['error_message'] = message

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, *args):
        self.stop()
        return False


# =============================================================================
# Pytest fixtures
# =============================================================================

def pytest_httpserver_fixture(dim: int = 768):
    """
    Create a pytest fixture for MockTEIHandler.

    Usage in conftest.py:
        from common.mock_tei_server import pytest_httpserver_fixture

        @pytest.fixture
        def mock_tei(httpserver):
            handler = MockTEIHandler(httpserver, dim=768)
            handler.setup_embed()
            yield handler
    """
    def fixture(httpserver):
        handler = MockTEIHandler(httpserver, dim=dim)
        handler.setup_embed()
        yield handler
    return fixture


if __name__ == '__main__':
    import urllib.request
    import urllib.error

    print("Testing standalone MockTEIServer...")
    with MockTEIServer(port=8080, dim=768) as server:
        print(f"Server: {server.endpoint}")

        # Test embed
        req = urllib.request.Request(
            f"{server.endpoint}/embed",
            data=json.dumps({"inputs": ["Hello", "World"]}).encode(),
            headers={'Content-Type': 'application/json'}
        )
        with urllib.request.urlopen(req) as resp:
            result = json.loads(resp.read())
            print(f"Embed: {len(result)} vectors, dim={len(result[0])}")

        # Test error mode
        server.set_error_mode(True, 503, "Model integration is not active")
        try:
            urllib.request.urlopen(req)
        except urllib.error.HTTPError as e:
            print(f"Error mode: {e.code} - {json.loads(e.read())}")

    print("Done!")