mariadb-columnstore-engine/cmapi/cmapi_server/test/mock_resolution.py

import socket
from contextlib import ExitStack, contextmanager
from typing import Dict, List, Optional, Tuple
from unittest.mock import patch

DEFAULT_LOCALHOST_HOSTNAME = 'localhost'
DEFAULT_LOCALHOST_IP = '127.0.0.1'

CUR_HOST_IP = '198.51.100.10'
CUR_HOST_HOSTNAME = 'current-host.test'

class MockResolutionBuilder:
    """Builder for creating complex name/IP resolution mocks as a context manager.
    We'll need all of its abilities later, when we add centralized address management.

    Usage:
        cm = (
            MockResolutionBuilder()
            .add_mapping('nodeA.local', '10.0.0.1')
            .add_mapping('nodeB.local', '10.0.0.2', bidirectional=False)
            .add_rev_mapping('10.0.0.2', 'nodeB.no_reverse')
            .build()
        )
        with cm:
            ...
    """

    def __init__(self):
        # Forward: hostname -> ip
        self._forward: Dict[str, str] = {}
        # Reverse: ip -> (primary_hostname, alias_list)
        self._reverse: Dict[str, Tuple[str, List[str]]] = {}
        # Defaults used when no explicit mapping is provided
        self._default_ip: Optional[str] = None
        self._default_hostname: Optional[str] = None

    def add_mapping(
        self,
        hostname: str,
        ip: str,
        bidirectional: bool = True,
    ):
        """Add a forward mapping hostname->ip.

        - If bidirectional is True (default), also add reverse ip->hostname.
        - If bidirectional is False, reverse resolution for this IP will fail
          unless explicitly set later via add_rev_mapping().
        """
        self._forward[hostname] = ip
        # Initialize defaults on first mapping (can be overwritten by set_default)
        if self._default_ip is None:
            self._default_ip = ip
        if self._default_hostname is None:
            self._default_hostname = hostname
        if bidirectional:
            # Always store reverse as a tuple: (primary, aliases)
            self._reverse[ip] = (hostname, [])
        return self

    def add_rev_mapping(self, ip: str, hostname: str, aliases: Optional[List[str]] = None):
        """Explicitly set reverse resolution mapping for IP -> hostname (and aliases).

        - Ensures reverse resolution succeeds for this IP.
        - Overrides any previous reverse mapping for this IP.
        - If aliases are provided, they will be returned by gethostbyaddr as the alias list.
        """
        self._reverse[ip] = (hostname, aliases or [])
        return self

    def set_default(self, ip: str, hostname: str):
        """Set default ip/hostname when a query is not present in mappings."""
        self._default_ip = ip
        self._default_hostname = hostname
        return self

    def _resolve_forward(self, host: str) -> Tuple[str, str]:
        """Resolve hostname or IP to (ip, hostname) using mappings/defaults."""
        # If input looks like an IP, return it with reverse or default hostname
        try:
            socket.inet_pton(socket.AF_INET, host)
            ip = host
            reverse = self._reverse.get(ip)
            if reverse is not None:
                primary, _aliases = reverse
                hostname = primary
            else:
                hostname = self._default_hostname or host
            return ip, hostname
        except OSError:
            pass

        # Treat as hostname
        if host in self._forward:
            ip = self._forward[host]
            hostname = host
            return ip, hostname

        if self._default_ip and self._default_hostname:
            return self._default_ip, self._default_hostname

        # As a last resort, echo back (host, host)
        return host, host

    def build(self):

        def _fake_getaddrinfo(host, port, family=socket.AF_UNSPEC, type=0, proto=0, flags=0):
            # Only handle AF_INET calls; otherwise, simulate failure
            if family not in (socket.AF_UNSPEC, socket.AF_INET):
                raise socket.gaierror
            # For localhost, return loopback first and include CUR_HOST_IP as secondary
            if host == DEFAULT_LOCALHOST_HOSTNAME:
                return [
                    (socket.AF_INET, socket.SOCK_STREAM, 6, '', (DEFAULT_LOCALHOST_IP, port)),
                    (socket.AF_INET, socket.SOCK_STREAM, 6, '', (CUR_HOST_IP, port)),
                ]
            ip, _ = self._resolve_forward(host)
            return [(socket.AF_INET, socket.SOCK_STREAM, 6, '', (ip, port))]

        def _fake_gethostbyname(name: str) -> str:
            ip, _ = self._resolve_forward(name)
            return ip

        def _fake_gethostbyaddr(addr: str):
            # If no reverse record was set, simulate reverse lookup failure
            if addr not in self._reverse:
                raise socket.herror
            primary, aliases = self._reverse[addr]
            return (primary, aliases, [addr])

        @contextmanager
        def _ctx():
            patches = [
                # Patch socket-level resolvers (NetworkManager uses these under the hood)
                patch('socket.getaddrinfo', side_effect=_fake_getaddrinfo),
                patch('socket.gethostbyname', side_effect=_fake_gethostbyname),
                patch('socket.gethostbyaddr', side_effect=_fake_gethostbyaddr),
                # Patch local identity to be synthetic; avoid real system calls
                patch('socket.gethostname', return_value=CUR_HOST_HOSTNAME),
                patch('socket.getfqdn', return_value=CUR_HOST_HOSTNAME),
            ]
            with ExitStack() as stack:
                for p in patches:
                    stack.enter_context(p)
                yield

        return _ctx()


def simple_resolution_mock(hostname: str, ip: str):
    """Return a context manager for simple name/IP resolution mocking.

    Behavior:
    - Unknown hostnames default to the synthetic current-host identity: (CUR_HOST_IP, CUR_HOST_HOSTNAME).
    - "localhost" resolves to 127.0.0.1, with reverse lookup to "localhost".
    - The synthetic current host name resolves to CUR_HOST_IP, with reverse lookup back to CUR_HOST_HOSTNAME.
    - The provided (hostname, ip) mapping is added bidirectionally.
    - socket.gethostname() / getfqdn() are patched to CUR_HOST_HOSTNAME to avoid touching the real system.
    """
    return (
        make_local_resolution_builder()
        .add_mapping(hostname, ip, bidirectional=True)
        .build()
    )


def make_local_resolution_builder() -> MockResolutionBuilder:
    """Return a MockResolutionBuilder pre-populated with local/current host
    resolution mappings. Tests can add additional mappings onto this builder
    before calling .build() to obtain a context manager.
    """
    return (
        MockResolutionBuilder()
        .set_default(CUR_HOST_IP, CUR_HOST_HOSTNAME)
        # Localhost forward and reverse
        .add_mapping(DEFAULT_LOCALHOST_HOSTNAME, DEFAULT_LOCALHOST_IP)
        .add_rev_mapping(DEFAULT_LOCALHOST_IP, DEFAULT_LOCALHOST_HOSTNAME)
        # Synthetic current host forward and reverse
        .add_mapping(CUR_HOST_HOSTNAME, CUR_HOST_IP)
        .add_rev_mapping(CUR_HOST_IP, CUR_HOST_HOSTNAME, [DEFAULT_LOCALHOST_HOSTNAME])
    )