cloudflare/cloudflared
Publicmirrored from https://github.com/cloudflare/cloudflaredAvailable
component-tests/util.py
84lines · modecode
| 1 | from contextlib import contextmanager |
| 2 | import logging |
| 3 | import requests |
| 4 | from retrying import retry |
| 5 | import subprocess |
| 6 | import yaml |
| 7 | |
| 8 | from time import sleep |
| 9 | |
| 10 | from constants import METRICS_PORT, MAX_RETRIES, BACKOFF_SECS |
| 11 | |
| 12 | LOGGER = logging.getLogger(__name__) |
| 13 | |
| 14 | |
| 15 | def write_config(path, config): |
| 16 | config_path = path / "config.yaml" |
| 17 | with open(config_path, 'w') as outfile: |
| 18 | yaml.dump(config, outfile) |
| 19 | return config_path |
| 20 | |
| 21 | |
| 22 | def start_cloudflared(path, config, cfd_args=["run"], cfd_pre_args=["tunnel"], new_process=False, allow_input=False, capture_output=True): |
| 23 | config_path = write_config(path, config.full_config) |
| 24 | cmd = [config.cloudflared_binary] |
| 25 | cmd += cfd_pre_args |
| 26 | cmd += ["--config", config_path] |
| 27 | cmd += cfd_args |
| 28 | LOGGER.info(f"Run cmd {cmd} with config {config}") |
| 29 | if new_process: |
| 30 | return run_cloudflared_background(cmd, allow_input, capture_output) |
| 31 | # By setting check=True, it will raise an exception if the process exits with non-zero exit code |
| 32 | return subprocess.run(cmd, check=True, capture_output=capture_output) |
| 33 | |
| 34 | |
| 35 | @contextmanager |
| 36 | def run_cloudflared_background(cmd, allow_input, capture_output): |
| 37 | output = subprocess.PIPE if capture_output else subprocess.DEVNULL |
| 38 | stdin = subprocess.PIPE if allow_input else None |
| 39 | try: |
| 40 | cfd = subprocess.Popen(cmd, stdin=stdin, stdout=output, stderr=output) |
| 41 | yield cfd |
| 42 | finally: |
| 43 | cfd.terminate() |
| 44 | |
| 45 | |
| 46 | @retry(stop_max_attempt_number=MAX_RETRIES, wait_fixed=BACKOFF_SECS * 1000) |
| 47 | def wait_tunnel_ready(expect_connections=4): |
| 48 | url = f'http://localhost:{METRICS_PORT}/ready' |
| 49 | |
| 50 | with requests.Session() as s: |
| 51 | resp = send_request(s, url, True) |
| 52 | assert resp.json()[ |
| 53 | "readyConnections"] == expect_connections, f"Ready endpoint returned {resp.json()} but we expect {expect_connections} ready connections" |
| 54 | |
| 55 | |
| 56 | @retry(stop_max_attempt_number=MAX_RETRIES, wait_fixed=BACKOFF_SECS * 1000) |
| 57 | def check_tunnel_not_ready(): |
| 58 | url = f'http://localhost:{METRICS_PORT}/ready' |
| 59 | |
| 60 | resp = requests.get(url, timeout=1) |
| 61 | assert resp.status_code == 503, f"Expect {url} returns 503, got {resp.status_code}" |
| 62 | |
| 63 | # In some cases we don't need to check response status, such as when sending batch requests to generate logs |
| 64 | |
| 65 | |
| 66 | def send_requests(url, count, require_ok=True): |
| 67 | errors = 0 |
| 68 | with requests.Session() as s: |
| 69 | for _ in range(count): |
| 70 | resp = send_request(s, url, require_ok) |
| 71 | if resp is None: |
| 72 | errors += 1 |
| 73 | sleep(0.01) |
| 74 | if errors > 0: |
| 75 | LOGGER.warning( |
| 76 | f"{errors} out of {count} requests to {url} return non-200 status") |
| 77 | |
| 78 | |
| 79 | @retry(stop_max_attempt_number=MAX_RETRIES, wait_fixed=BACKOFF_SECS * 1000) |
| 80 | def send_request(session, url, require_ok): |
| 81 | resp = session.get(url, timeout=BACKOFF_SECS) |
| 82 | if require_ok: |
| 83 | assert resp.status_code == 200, f"{url} returned {resp}" |
| 84 | return resp if resp.status_code == 200 else None |
| 85 | |