dvc/tests/func/test_daemon.py

import json
import os
import re
import subprocess
import sys
import time
from collections import defaultdict
from contextlib import contextmanager
from http.server import HTTPServer, SimpleHTTPRequestHandler
from pathlib import Path
from threading import Thread
from typing import Dict, Iterator

import pytest

from dvc import version_tuple
from dvc.daemon import _get_dvc_args
from dvc.env import (
    DVC_ANALYTICS_ENDPOINT,
    DVC_DAEMON_LOGFILE,
    DVC_NO_ANALYTICS,
    DVC_UPDATER_ENDPOINT,
)
from dvc.updater import Updater

version = ".".join(map(str, version_tuple[:3]))
UPDATER_INFO = {
    "version": version,
    "packages": {
        "linux": {
            "deb": f"https://dvc.org/download/linux-deb/dvc-{version}",
            "rpm": f"https://dvc.org/download/linux-rpm/dvc-{version}",
        },
        "windows": {"exe": f"https://dvc.org/download/win/dvc-{version}"},
        "osx": {"pkg": f"https://dvc.org/download/osx/dvc-{version}"},
    },
}

UPDATER_INFO_STR = json.dumps(UPDATER_INFO).encode("utf8")


def make_request_handler():
    class RequestHandler(SimpleHTTPRequestHandler):
        # save requests count for each method
        hits: Dict[str, int] = defaultdict(int)

        def log_message(self, format, *args) -> None:  # noqa: A002
            super().log_message(format, *args)
            if length := self.headers.get("Content-Length"):
                data = self.rfile.read(int(length)).decode("utf8")
                sys.stderr.write(f"{data}\n")

        def do_POST(self):  # noqa: N802
            # analytics endpoint
            self.hits["POST"] += 1
            self.send_response(200)
            super().end_headers()

        def do_GET(self):  # noqa: N802
            # updater endpoint
            self.hits["GET"] += 1
            self.send_response(200)
            self.send_header("Content-type", "application/json")
            self.end_headers()
            self.wfile.write(UPDATER_INFO_STR)

    return RequestHandler


@contextmanager
def make_server(port: int = 0) -> Iterator["HTTPServer"]:
    with HTTPServer(("localhost", port), make_request_handler()) as httpd:
        yield httpd


@pytest.fixture
def server():
    with make_server() as httpd:
        thread = Thread(target=httpd.serve_forever)
        thread.daemon = True
        thread.start()
        yield httpd


def retry_until(pred, timeout):
    timeout_ns = timeout * 1e9
    start = time.perf_counter_ns()
    while time.perf_counter_ns() - start < timeout_ns:
        if pred():
            return
        time.sleep(0.01)
    raise RuntimeError(f"timed out after {timeout}s")


def test_analytics(tmp_path, server):
    addr = server.server_address
    logfile = tmp_path / "logfile"

    env = {
        **os.environ,
        DVC_DAEMON_LOGFILE: str(logfile),
        DVC_ANALYTICS_ENDPOINT: "http://{}:{}".format(*addr),
    }
    env.pop("DVC_TEST", None)
    env.pop("DVC_NO_ANALYTICS", None)

    output = subprocess.check_output(
        [*_get_dvc_args(), "config", "-l", "-vv"],
        env=env,
        text=True,
    )

    pattern = r".*Saving analytics report to (.*)"
    for line in output.splitlines():
        if match := re.search(pattern, line):
            report_file = match.group(1).strip()
            break
    else:
        raise AssertionError("no match for the report file")

    # wait until the file disappears
    retry_until(lambda: not os.path.exists(report_file), 10)
    # wait till the daemon exits
    retry_until(lambda: "exiting with 0" in logfile.read_text(encoding="utf8"), 5)
    assert server.RequestHandlerClass.hits == {"POST": 1}


def test_updater(tmp_dir, dvc, server):
    addr = server.server_address
    logfile = tmp_dir / "logfile"

    env = {
        **os.environ,
        DVC_DAEMON_LOGFILE: str(logfile),
        DVC_UPDATER_ENDPOINT: "http://{}:{}".format(*addr),
        # prevent running analytics daemon
        DVC_NO_ANALYTICS: "true",
    }
    env.pop("DVC_TEST", None)
    env.pop("CI", None)

    subprocess.check_output(
        [*_get_dvc_args(), "version", "-vv"],
        env=env,
        text=True,
    )

    updater_file = Path(dvc.tmp_dir) / Updater.UPDATER_FILE

    # wait until the updater file appears
    retry_until(updater_file.is_file, 10)
    # wait till the daemon exits
    retry_until(lambda: "exiting with 0" in logfile.read_text(encoding="utf8"), 5)
    assert server.RequestHandlerClass.hits == {"GET": 1}
    # check that the file is saved correctly
    assert json.loads(updater_file.read_text(encoding="utf8")) == UPDATER_INFO


if __name__ == "__main__":
    # python -m tests.func.test_daemon [<port>]
    port = int(sys.argv[1]) if len(sys.argv) >= 2 else 0
    with make_server(port) as httpd:
        print(  # noqa:  T201
            "Running server on http://{}:{}".format(*httpd.server_address)
        )
        httpd.serve_forever()