chore: import upstream snapshot with attribution

tests/python_client/common/mock_tei_server.py

@@ -0,0 +1,408 @@
+"""
+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!")