| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| """ |
| Shared infrastructure for the image-server test suite (stdlib unittest only). |
| |
| Provides: |
| - A singleton image server process started once for the entire test run. |
| - Server stdout/stderr appended to ``<tmp>/imageserver.log``. |
| - On shutdown: stop the child process, close the log handle, unlink the control socket; |
| the temp directory and ``imageserver.log`` are left on disk. |
| - Control-socket helpers using pure-Python AF_UNIX. |
| - qemu-nbd server management. |
| - Transfer registration / teardown helpers. |
| - HTTP helper functions. |
| """ |
| |
| import atexit |
| import functools |
| import json |
| import logging |
| import os |
| import random |
| import signal |
| import socket |
| import subprocess |
| import sys |
| import tempfile |
| import time |
| import unittest |
| import uuid |
| from pathlib import Path |
| from typing import Any, Dict, Optional, TextIO |
| |
| IMAGE_SIZE = 1 * 1024 * 1024 # 1 MiB |
| SERVER_STARTUP_TIMEOUT = 10 |
| QEMU_NBD_STARTUP_TIMEOUT = 5 |
| HTTP_TIMEOUT = 30 # seconds per HTTP request |
| |
| logging.basicConfig( |
| level=logging.INFO, |
| stream=sys.stderr, |
| format="%(asctime)s [TEST] %(message)s", |
| ) |
| log = logging.getLogger(__name__) |
| |
| |
| def randbytes(seed, n): |
| """Generate n deterministic pseudo-random bytes (works on Python 3.6+).""" |
| rng = random.Random(seed) |
| return rng.getrandbits(8 * n).to_bytes(n, "big") |
| |
| |
| def test_timeout(seconds): |
| """Decorator that fails a test if it exceeds *seconds* (SIGALRM, Unix only).""" |
| def decorator(func): |
| @functools.wraps(func) |
| def wrapper(*args, **kwargs): |
| def _alarm(signum, frame): |
| raise TimeoutError( |
| "{} timed out after {}s".format(func.__qualname__, seconds) |
| ) |
| prev = signal.signal(signal.SIGALRM, _alarm) |
| signal.alarm(seconds) |
| try: |
| return func(*args, **kwargs) |
| finally: |
| signal.alarm(0) |
| signal.signal(signal.SIGALRM, prev) |
| return wrapper |
| return decorator |
| |
| # ── Singleton state shared across all test modules ────────────────────── |
| |
| _tmp_dir: Optional[str] = None |
| _server_proc: Optional[subprocess.Popen] = None |
| _server_info: Optional[Dict[str, Any]] = None |
| _server_log_fp: Optional[TextIO] = None |
| _server_log_path: Optional[str] = None |
| _atexit_registered: bool = False |
| |
| |
| def _free_port() -> int: |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| s.bind(("127.0.0.1", 0)) |
| return s.getsockname()[1] |
| |
| |
| def control_socket_send(sock_path: str, message: dict, retries: int = 5) -> dict: |
| """Send a JSON message to the control socket and return the parsed response.""" |
| payload = (json.dumps(message) + "\n").encode("utf-8") |
| last_err = None |
| for attempt in range(retries): |
| try: |
| with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s: |
| s.settimeout(5) |
| s.connect(sock_path) |
| s.sendall(payload) |
| s.shutdown(socket.SHUT_WR) |
| data = b"" |
| while True: |
| chunk = s.recv(4096) |
| if not chunk: |
| break |
| data += chunk |
| return json.loads(data.strip()) |
| except (BlockingIOError, ConnectionRefusedError, OSError) as e: |
| last_err = e |
| time.sleep(0.1 * (attempt + 1)) |
| raise last_err |
| |
| |
| def _wait_for_control_socket(sock_path: str, timeout: float = SERVER_STARTUP_TIMEOUT) -> None: |
| deadline = time.monotonic() + timeout |
| while time.monotonic() < deadline: |
| try: |
| resp = control_socket_send(sock_path, {"action": "status"}) |
| if resp.get("status") == "ok": |
| return |
| except (ConnectionRefusedError, FileNotFoundError, OSError): |
| pass |
| time.sleep(0.2) |
| raise RuntimeError( |
| f"Image server control socket at {sock_path} not ready within {timeout}s" |
| ) |
| |
| |
| def _wait_for_nbd_socket(sock_path: str, timeout: float = QEMU_NBD_STARTUP_TIMEOUT) -> None: |
| deadline = time.monotonic() + timeout |
| while time.monotonic() < deadline: |
| if os.path.exists(sock_path): |
| try: |
| with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s: |
| s.settimeout(1) |
| s.connect(sock_path) |
| return |
| except (ConnectionRefusedError, OSError): |
| pass |
| time.sleep(0.2) |
| raise RuntimeError( |
| f"qemu-nbd socket at {sock_path} not ready within {timeout}s" |
| ) |
| |
| |
| def get_tmp_dir() -> str: |
| global _tmp_dir |
| if _tmp_dir is None: |
| _tmp_dir = tempfile.mkdtemp(prefix="imageserver_test_") |
| return _tmp_dir |
| |
| |
| def _read_log_tail(path: str, max_bytes: int = 65536) -> str: |
| """Return up to *max_bytes* of UTF-8 text from the end of *path*.""" |
| try: |
| with open(path, "rb") as f: |
| f.seek(0, os.SEEK_END) |
| size = f.tell() |
| f.seek(max(0, size - max_bytes)) |
| return f.read().decode("utf-8", errors="replace") |
| except OSError as e: |
| return f"(could not read log: {e})" |
| |
| |
| def get_image_server() -> Dict[str, Any]: |
| """Return the singleton image-server info dict, starting it if needed.""" |
| global _server_proc, _server_info, _server_log_fp, _server_log_path, _atexit_registered |
| |
| if _server_info is not None: |
| return _server_info |
| |
| tmp = get_tmp_dir() |
| port = _free_port() |
| ctrl_sock = os.path.join(tmp, "ctrl.sock") |
| log_path = os.path.join(tmp, "imageserver.log") |
| _server_log_path = log_path |
| |
| imageserver_pkg = str(Path(__file__).resolve().parent.parent) |
| parent_dir = str(Path(imageserver_pkg).parent) |
| |
| env = os.environ.copy() |
| env["PYTHONPATH"] = parent_dir + os.pathsep + env.get("PYTHONPATH", "") |
| |
| _server_log_fp = open( |
| log_path, "a", encoding="utf-8", buffering=1, errors="replace" |
| ) |
| try: |
| _server_log_fp.write( |
| "\n========== imageserver test subprocess log ==========\n" |
| ) |
| _server_log_fp.flush() |
| except OSError: |
| pass |
| |
| proc = subprocess.Popen( |
| [ |
| sys.executable, "-m", "imageserver", |
| "--listen", "127.0.0.1", |
| "--port", str(port), |
| "--control-socket", ctrl_sock, |
| ], |
| cwd=parent_dir, |
| env=env, |
| stdout=_server_log_fp, |
| stderr=_server_log_fp, |
| ) |
| _server_proc = proc |
| |
| try: |
| _wait_for_control_socket(ctrl_sock) |
| except RuntimeError: |
| proc.kill() |
| try: |
| proc.wait(timeout=5) |
| except subprocess.TimeoutExpired: |
| proc.kill() |
| proc.wait(timeout=5) |
| try: |
| _server_log_fp.flush() |
| except OSError: |
| pass |
| tail = _read_log_tail(log_path) |
| try: |
| _server_log_fp.close() |
| except OSError: |
| pass |
| _server_log_fp = None |
| _server_proc = None |
| raise RuntimeError( |
| "Image server failed to start.\n" |
| f"Log file: {log_path}\n" |
| f"--- log tail ---\n{tail}" |
| ) |
| |
| def send(msg: dict) -> dict: |
| return control_socket_send(ctrl_sock, msg) |
| |
| _server_info = { |
| "base_url": f"http://127.0.0.1:{port}", |
| "port": port, |
| "ctrl_sock": ctrl_sock, |
| "send": send, |
| "imageserver_log": log_path, |
| } |
| if not _atexit_registered: |
| atexit.register(shutdown_image_server) |
| _atexit_registered = True |
| sys.stdout.write( |
| "\n[IMAGESERVER_TEST] child image server log file: %s\n\n" % log_path |
| ) |
| sys.stdout.flush() |
| return _server_info |
| |
| |
| def shutdown_image_server() -> None: |
| global _server_proc, _server_info, _tmp_dir, _server_log_fp, _server_log_path |
| ctrl_sock: Optional[str] = None |
| if _server_info is not None: |
| ctrl_sock = _server_info.get("ctrl_sock") |
| |
| if _server_proc is not None: |
| _server_proc.terminate() |
| try: |
| _server_proc.wait(timeout=5) |
| except subprocess.TimeoutExpired: |
| _server_proc.kill() |
| _server_proc.wait(timeout=5) |
| _server_proc = None |
| if _server_log_fp is not None: |
| try: |
| _server_log_fp.flush() |
| _server_log_fp.close() |
| except OSError: |
| pass |
| _server_log_fp = None |
| _server_info = None |
| _server_log_path = None |
| |
| if ctrl_sock: |
| try: |
| os.unlink(ctrl_sock) |
| except FileNotFoundError: |
| pass |
| |
| # Leave temp dir and imageserver.log on disk for debugging; clear pointer only. |
| _tmp_dir = None |
| |
| |
| # ── qemu-nbd server ──────────────────────────────────────────────────── |
| |
| class QemuNbdServer: |
| """Manages a qemu-nbd process exporting a disk image over a Unix socket.""" |
| |
| def __init__( |
| self, |
| image_path: str, |
| socket_path: str, |
| image_size: int = IMAGE_SIZE, |
| image_format: str = "raw", |
| ): |
| self.image_path = image_path |
| self.socket_path = socket_path |
| self.image_size = image_size |
| self.image_format = image_format |
| self._proc: Optional[subprocess.Popen] = None |
| |
| def start(self) -> None: |
| if not os.path.exists(self.image_path): |
| if self.image_format == "raw": |
| with open(self.image_path, "wb") as f: |
| f.truncate(self.image_size) |
| else: |
| raise FileNotFoundError( |
| f"disk image not found for format {self.image_format!r}: {self.image_path}" |
| ) |
| |
| self._proc = subprocess.Popen( |
| [ |
| "qemu-nbd", |
| "--socket", self.socket_path, |
| "--format", self.image_format, |
| "--persistent", |
| "--shared=0", |
| "--cache=none", |
| self.image_path, |
| ], |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| ) |
| _wait_for_nbd_socket(self.socket_path) |
| |
| def stop(self) -> None: |
| if self._proc is not None: |
| for pipe in (self._proc.stdout, self._proc.stderr): |
| if pipe: |
| try: |
| pipe.close() |
| except Exception: |
| pass |
| self._proc.terminate() |
| try: |
| self._proc.wait(timeout=5) |
| except subprocess.TimeoutExpired: |
| self._proc.kill() |
| self._proc.wait(timeout=5) |
| self._proc = None |
| |
| |
| # ── Factory helpers ───────────────────────────────────────────────────── |
| |
| def make_tmp_image(data=None, image_size=IMAGE_SIZE) -> str: |
| """Create a temp raw image file in the shared tmp dir; return path.""" |
| tmp = get_tmp_dir() |
| path = os.path.join(tmp, f"img_{uuid.uuid4().hex[:8]}.raw") |
| if data is not None: |
| with open(path, "wb") as f: |
| f.write(data) |
| else: |
| with open(path, "wb") as f: |
| f.write(randbytes(42, image_size)) |
| return path |
| |
| |
| def make_file_transfer(data=None, image_size=IMAGE_SIZE, idle_timeout_seconds=None): |
| """ |
| Create a temp file + register a file-backend transfer. |
| Returns (transfer_id, url, file_path, cleanup_callable). |
| |
| If *idle_timeout_seconds* is set, it is sent in the transfer config (for idle expiry tests). |
| """ |
| srv = get_image_server() |
| path = make_tmp_image(data=data, image_size=image_size) |
| transfer_id = f"file-{uuid.uuid4().hex[:8]}" |
| cfg = {"backend": "file", "file": path} |
| if idle_timeout_seconds is not None: |
| cfg["idle_timeout_seconds"] = idle_timeout_seconds |
| resp = srv["send"]({ |
| "action": "register", |
| "transfer_id": transfer_id, |
| "config": cfg, |
| }) |
| assert resp["status"] == "ok", f"register failed: {resp}" |
| url = f"{srv['base_url']}/images/{transfer_id}" |
| |
| def cleanup(): |
| srv["send"]({"action": "unregister", "transfer_id": transfer_id}) |
| try: |
| os.unlink(path) |
| except FileNotFoundError: |
| pass |
| |
| return transfer_id, url, path, cleanup |
| |
| |
| def make_nbd_transfer(image_size=IMAGE_SIZE): |
| """ |
| Create a qemu-nbd server + register an NBD-backend transfer. |
| Returns (transfer_id, url, QemuNbdServer, cleanup_callable). |
| """ |
| srv = get_image_server() |
| tmp = get_tmp_dir() |
| img_path = os.path.join(tmp, f"nbd_{uuid.uuid4().hex[:8]}.raw") |
| sock_path = os.path.join(tmp, f"nbd_{uuid.uuid4().hex[:8]}.sock") |
| |
| server = QemuNbdServer(img_path, sock_path, image_size=image_size) |
| server.start() |
| |
| transfer_id = f"nbd-{uuid.uuid4().hex[:8]}" |
| resp = srv["send"]({ |
| "action": "register", |
| "transfer_id": transfer_id, |
| "config": {"backend": "nbd", "socket": sock_path}, |
| }) |
| assert resp["status"] == "ok", f"register failed: {resp}" |
| url = f"{srv['base_url']}/images/{transfer_id}" |
| |
| def cleanup(): |
| srv["send"]({"action": "unregister", "transfer_id": transfer_id}) |
| server.stop() |
| for p in (img_path, sock_path): |
| try: |
| os.unlink(p) |
| except FileNotFoundError: |
| pass |
| |
| return transfer_id, url, server, cleanup |
| |
| |
| def make_nbd_transfer_existing_disk(image_path: str, image_format: str = "qcow2"): |
| """ |
| Start qemu-nbd for an existing on-disk image (e.g. qcow2) and register a transfer. |
| |
| Does not delete *image_path* on cleanup (only the Unix socket under tmp). |
| |
| Returns (transfer_id, url, QemuNbdServer, cleanup_callable). |
| """ |
| srv = get_image_server() |
| tmp = get_tmp_dir() |
| sock_path = os.path.join(tmp, f"nbd_{uuid.uuid4().hex[:8]}.sock") |
| |
| server = QemuNbdServer( |
| image_path, sock_path, image_format=image_format |
| ) |
| server.start() |
| |
| transfer_id = f"nbd-{uuid.uuid4().hex[:8]}" |
| resp = srv["send"]({ |
| "action": "register", |
| "transfer_id": transfer_id, |
| "config": {"backend": "nbd", "socket": sock_path}, |
| }) |
| assert resp["status"] == "ok", f"register failed: {resp}" |
| url = f"{srv['base_url']}/images/{transfer_id}" |
| |
| def cleanup(): |
| srv["send"]({"action": "unregister", "transfer_id": transfer_id}) |
| server.stop() |
| try: |
| os.unlink(sock_path) |
| except FileNotFoundError: |
| pass |
| |
| return transfer_id, url, server, cleanup |
| |
| |
| # ── HTTP helpers ──────────────────────────────────────────────────────── |
| |
| import urllib.request |
| import urllib.error |
| |
| |
| def http_get(url, headers=None, timeout=HTTP_TIMEOUT): |
| req = urllib.request.Request(url, headers=headers or {}) |
| return urllib.request.urlopen(req, timeout=timeout) |
| |
| |
| def http_put(url, data, headers=None, timeout=HTTP_TIMEOUT): |
| hdrs = {"Content-Length": str(len(data))} |
| if headers: |
| hdrs.update(headers) |
| req = urllib.request.Request(url, data=data, headers=hdrs, method="PUT") |
| return urllib.request.urlopen(req, timeout=timeout) |
| |
| |
| def http_post(url, data=b"", headers=None, timeout=HTTP_TIMEOUT): |
| hdrs = {} |
| if headers: |
| hdrs.update(headers) |
| req = urllib.request.Request(url, data=data, headers=hdrs, method="POST") |
| return urllib.request.urlopen(req, timeout=timeout) |
| |
| |
| def http_options(url, timeout=HTTP_TIMEOUT): |
| req = urllib.request.Request(url, method="OPTIONS") |
| return urllib.request.urlopen(req, timeout=timeout) |
| |
| |
| def http_patch(url, data, headers=None, timeout=HTTP_TIMEOUT): |
| hdrs = {} |
| if headers: |
| hdrs.update(headers) |
| req = urllib.request.Request(url, data=data, headers=hdrs, method="PATCH") |
| return urllib.request.urlopen(req, timeout=timeout) |
| |
| |
| # ── Base TestCase with shared setUp/tearDown ──────────────────────────── |
| |
| class ImageServerTestCase(unittest.TestCase): |
| """ |
| Base class for image-server tests. |
| |
| Ensures the image server is running before any test method. |
| Subclasses that need a file or NBD transfer should set them up |
| in setUp() and tear down in tearDown(). |
| """ |
| |
| @classmethod |
| def setUpClass(cls): |
| cls.server = get_image_server() |
| cls.base_url = cls.server["base_url"] |
| |
| def ctrl(self, msg): |
| """Send a control-socket message; wraps server['send'] to avoid descriptor issues.""" |
| return self.server["send"](msg) |
| |
| def _make_tmp_image(self, data=None): |
| return make_tmp_image(data=data) |
| |
| def _register_file_transfer(self, data=None): |
| return make_file_transfer(data=data) |
| |
| def _register_nbd_transfer(self): |
| return make_nbd_transfer() |
| |
| @staticmethod |
| def dump_server_logs(max_bytes: int = 256 * 1024): |
| """Print a tail of the image-server log file (shared by all tests in the run).""" |
| path = _server_log_path |
| if not path or not os.path.isfile(path): |
| return |
| try: |
| if _server_log_fp is not None: |
| _server_log_fp.flush() |
| except OSError: |
| pass |
| try: |
| data = _read_log_tail(path, max_bytes=max_bytes) |
| if data.strip(): |
| sys.stderr.write("\n=== IMAGE SERVER LOG (tail) ===\n") |
| sys.stderr.write(data) |
| if not data.endswith("\n"): |
| sys.stderr.write("\n") |
| sys.stderr.write("=== END IMAGE SERVER LOG ===\n") |
| except Exception: |
| pass |
