server.py

# -*- coding: utf-8 -*-
"""WireMock Server Management."""
import atexit
import socket
import time
from subprocess import PIPE, STDOUT, Popen

import requests
from importlib.resources import files

from wiremock.server.exceptions import (
    WireMockServerAlreadyStartedError,
    WireMockServerNotStartedError,
)


class WireMockServer(object):

    DEFAULT_JAVA = "java"  # Assume java in PATH
    DEFAULT_JAR = files("wiremock") / "server" / "wiremock-standalone-2.35.1.jar"

    def __init__(
        self,
        java_path=DEFAULT_JAVA,
        jar_path=DEFAULT_JAR,
        port=None,
        max_attempts=10,
        root_dir=None,
    ):
        self.java_path = java_path
        self.jar_path = jar_path
        self.port = port or self._get_free_port()
        self.__subprocess = None
        self.__running = False
        self.max_attempts = max_attempts
        self.root_dir = root_dir

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

    def __exit__(self, type, value, traceback):
        self.stop()

    @property
    def is_running(self):
        return self.__running

    def start(self):
        if self.is_running:
            raise WireMockServerAlreadyStartedError(
                "WireMockServer already started on port {}".format(self.port)
            )

        cmd = [
            self.java_path,
            "-jar",
            self.jar_path,
            "--port",
            str(self.port),
            "--local-response-templating",
        ]
        if self.root_dir is not None:
            cmd.append("--root-dir")
            cmd.append(str(self.root_dir))
        try:
            self.__subprocess = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=STDOUT)
        except OSError as e:
            raise WireMockServerNotStartedError(str(e))  # Problem with Java

        time.sleep(0.1)
        if self.__subprocess.poll() is not None:
            # Process complete - server not started
            raise WireMockServerNotStartedError(
                "\n".join(
                    [
                        "returncode: {}".format(self.__subprocess.returncode),
                        "stdout:",
                        str(self.__subprocess.stdout.read()),
                    ]
                )
            )

        # Call the /__admin endpoint as a check for running state
        attempts = 0
        success = False
        while attempts < self.max_attempts:
            try:
                attempts += 1
                resp = requests.get("http://localhost:{}/__admin".format(self.port))
                if resp.status_code == 200:
                    success = True
                    break
            except requests.exceptions.ConnectionError:
                pass

            time.sleep(0.25)

        if not success:
            raise WireMockServerNotStartedError(
                "unable to get a successful GET http://localhost:{}/__admin response".format(
                    self.port
                )
            )

        atexit.register(self.stop, raise_on_error=False)
        self.__running = True

    def stop(self, raise_on_error=True):
        try:
            self.__subprocess.kill()
        except AttributeError:
            if raise_on_error:
                raise WireMockServerNotStartedError()

    def _get_free_port(self):
        s = socket.socket(socket.AF_INET, type=socket.SOCK_STREAM)
        s.bind(("localhost", 0))
        address, port = s.getsockname()
        s.close()
        return port