|
| 1 | +"""pytest-trackflaky plugin implementation.""" |
| 2 | + |
| 3 | +import pytest |
| 4 | +import subprocess |
| 5 | +from urllib import request |
| 6 | +import os |
| 7 | +import json |
| 8 | +from time import time |
| 9 | +import unittest |
| 10 | +import threading |
| 11 | + |
| 12 | + |
| 13 | +# Global state for run tracking |
| 14 | +_run_id = None |
| 15 | +_run_id_lock = threading.Lock() |
| 16 | +server = os.environ.get("CI_SERVER_URL", None) |
| 17 | + |
| 18 | + |
| 19 | +class SnowflakeGenerator: |
| 20 | + """ |
| 21 | + Generates Twitter-style Snowflake IDs. |
| 22 | +
|
| 23 | + Format (64 bits): |
| 24 | + - 41 bits: timestamp in milliseconds since custom epoch |
| 25 | + - 10 bits: worker/machine ID |
| 26 | + - 12 bits: sequence number |
| 27 | + """ |
| 28 | + |
| 29 | + # Custom epoch (2024-01-01 00:00:00 UTC in milliseconds) |
| 30 | + EPOCH = 1704067200000 |
| 31 | + |
| 32 | + # Bit allocation |
| 33 | + TIMESTAMP_BITS = 41 |
| 34 | + WORKER_BITS = 10 |
| 35 | + SEQUENCE_BITS = 12 |
| 36 | + |
| 37 | + # Max values |
| 38 | + MAX_WORKER_ID = (1 << WORKER_BITS) - 1 |
| 39 | + MAX_SEQUENCE = (1 << SEQUENCE_BITS) - 1 |
| 40 | + |
| 41 | + # Bit shifts |
| 42 | + TIMESTAMP_SHIFT = WORKER_BITS + SEQUENCE_BITS |
| 43 | + WORKER_SHIFT = SEQUENCE_BITS |
| 44 | + |
| 45 | + def __init__(self, worker_id=None): |
| 46 | + """Initialize the snowflake generator.""" |
| 47 | + if worker_id is None: |
| 48 | + # Try to get worker ID from environment or use process ID |
| 49 | + worker_id = os.getpid() & self.MAX_WORKER_ID |
| 50 | + |
| 51 | + if worker_id > self.MAX_WORKER_ID or worker_id < 0: |
| 52 | + raise ValueError(f"Worker ID must be between 0 and {self.MAX_WORKER_ID}") |
| 53 | + |
| 54 | + self.worker_id = worker_id |
| 55 | + self.sequence = 0 |
| 56 | + self.last_timestamp = -1 |
| 57 | + self.lock = threading.Lock() |
| 58 | + |
| 59 | + def _current_timestamp(self): |
| 60 | + """Get current timestamp in milliseconds since epoch.""" |
| 61 | + return int(time() * 1000) |
| 62 | + |
| 63 | + def generate(self): |
| 64 | + """Generate a new Snowflake ID.""" |
| 65 | + with self.lock: |
| 66 | + timestamp = self._current_timestamp() - self.EPOCH |
| 67 | + |
| 68 | + if timestamp < self.last_timestamp: |
| 69 | + raise Exception("Clock moved backwards. Refusing to generate ID.") |
| 70 | + |
| 71 | + if timestamp == self.last_timestamp: |
| 72 | + self.sequence = (self.sequence + 1) & self.MAX_SEQUENCE |
| 73 | + if self.sequence == 0: |
| 74 | + # Sequence exhausted, wait for next millisecond |
| 75 | + while timestamp <= self.last_timestamp: |
| 76 | + timestamp = self._current_timestamp() - self.EPOCH |
| 77 | + else: |
| 78 | + self.sequence = 0 |
| 79 | + |
| 80 | + self.last_timestamp = timestamp |
| 81 | + |
| 82 | + # Combine all parts |
| 83 | + snowflake_id = ( |
| 84 | + (timestamp << self.TIMESTAMP_SHIFT) |
| 85 | + | (self.worker_id << self.WORKER_SHIFT) |
| 86 | + | self.sequence |
| 87 | + ) |
| 88 | + |
| 89 | + return snowflake_id |
| 90 | + |
| 91 | + |
| 92 | +# Global snowflake generator |
| 93 | +_snowflake_gen = SnowflakeGenerator() |
| 94 | + |
| 95 | + |
| 96 | +def get_git_sha(): |
| 97 | + """Get the current git commit SHA.""" |
| 98 | + try: |
| 99 | + return ( |
| 100 | + subprocess.check_output(["git", "rev-parse", "HEAD"]) |
| 101 | + .decode("ASCII") |
| 102 | + .strip() |
| 103 | + ) |
| 104 | + except subprocess.CalledProcessError: |
| 105 | + return None |
| 106 | + |
| 107 | + |
| 108 | +def get_git_branch(): |
| 109 | + """Get the current git branch name.""" |
| 110 | + try: |
| 111 | + return ( |
| 112 | + subprocess.check_output(["git", "rev-parse", "--abbrev-ref", "HEAD"]) |
| 113 | + .decode("ASCII") |
| 114 | + .strip() |
| 115 | + ) |
| 116 | + except subprocess.CalledProcessError: |
| 117 | + return None |
| 118 | + |
| 119 | + |
| 120 | +def get_git_repository(): |
| 121 | + """ |
| 122 | + Detect the git repository from remotes. |
| 123 | +
|
| 124 | + Checks refs in order: master@upstream, main@upstream, master@origin, main@origin. |
| 125 | + Returns repository in "owner/repo" format (e.g., "ElementsProject/lightning"). |
| 126 | + """ |
| 127 | + # Check these refs in order |
| 128 | + refs_to_check = [ |
| 129 | + "master@upstream", |
| 130 | + "main@upstream", |
| 131 | + "master@origin", |
| 132 | + "main@origin" |
| 133 | + ] |
| 134 | + |
| 135 | + for ref in refs_to_check: |
| 136 | + try: |
| 137 | + # Try to get the URL for this ref |
| 138 | + remote_url = ( |
| 139 | + subprocess.check_output( |
| 140 | + ["git", "config", "--get", f"remote.{ref.split('@')[1]}.url"], |
| 141 | + stderr=subprocess.DEVNULL |
| 142 | + ) |
| 143 | + .decode("ASCII") |
| 144 | + .strip() |
| 145 | + ) |
| 146 | + |
| 147 | + # Parse the URL to extract owner/repo |
| 148 | + # Handle various formats: |
| 149 | + # - https://github.com/owner/repo.git |
| 150 | + # - [email protected]:owner/repo.git |
| 151 | + # - https://github.com/owner/repo |
| 152 | + |
| 153 | + if remote_url.startswith("git@"): |
| 154 | + # SSH format: [email protected]:owner/repo.git |
| 155 | + path = remote_url.split(":", 1)[1] |
| 156 | + elif "://" in remote_url: |
| 157 | + # HTTPS format: https://github.com/owner/repo.git |
| 158 | + path = remote_url.split("://", 1)[1] |
| 159 | + # Remove the domain part |
| 160 | + if "/" in path: |
| 161 | + path = "/".join(path.split("/")[1:]) |
| 162 | + else: |
| 163 | + # Unknown format, try next ref |
| 164 | + continue |
| 165 | + |
| 166 | + # Remove .git suffix if present |
| 167 | + if path.endswith(".git"): |
| 168 | + path = path[:-4] |
| 169 | + |
| 170 | + # Ensure we have owner/repo format |
| 171 | + parts = path.split("/") |
| 172 | + if len(parts) >= 2: |
| 173 | + return f"{parts[-2]}/{parts[-1]}" |
| 174 | + |
| 175 | + except subprocess.CalledProcessError: |
| 176 | + # This ref doesn't exist, try the next one |
| 177 | + continue |
| 178 | + |
| 179 | + return None |
| 180 | + |
| 181 | + |
| 182 | +def get_run_id(): |
| 183 | + """Get or generate the run ID for this test session.""" |
| 184 | + global _run_id |
| 185 | + with _run_id_lock: |
| 186 | + if _run_id is None: |
| 187 | + _run_id = _snowflake_gen.generate() |
| 188 | + return _run_id |
| 189 | + |
| 190 | + |
| 191 | +def set_run_id(run_id): |
| 192 | + """Set the run ID (used by workers to inherit from main process).""" |
| 193 | + global _run_id |
| 194 | + with _run_id_lock: |
| 195 | + _run_id = run_id |
| 196 | + |
| 197 | + |
| 198 | +def get_base_result(): |
| 199 | + """Collect base result information from environment and git.""" |
| 200 | + github_sha = get_git_sha() |
| 201 | + github_ref_name = get_git_branch() |
| 202 | + github_run_id = os.environ.get("GITHUB_RUN_ID", None) |
| 203 | + run_number = os.environ.get("GITHUB_RUN_NUMBER", None) |
| 204 | + |
| 205 | + # Auto-detect repository from git remotes if not in environment |
| 206 | + github_repository = os.environ.get("GITHUB_REPOSITORY", None) |
| 207 | + if not github_repository: |
| 208 | + github_repository = get_git_repository() |
| 209 | + |
| 210 | + return { |
| 211 | + "run_id": get_run_id(), |
| 212 | + "github_repository": github_repository, |
| 213 | + "github_sha": os.environ.get("GITHUB_SHA", github_sha), |
| 214 | + "github_ref": os.environ.get("GITHUB_REF", None), |
| 215 | + "github_ref_name": github_ref_name, |
| 216 | + "github_run_id": int(github_run_id) if github_run_id else None, |
| 217 | + "github_head_ref": os.environ.get("GITHUB_HEAD_REF", None), |
| 218 | + "github_run_number": int(run_number) if run_number else None, |
| 219 | + "github_base_ref": os.environ.get("GITHUB_BASE_REF", None), |
| 220 | + "github_run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT", None), |
| 221 | + } |
| 222 | + |
| 223 | + |
| 224 | +def pytest_configure(config): |
| 225 | + """Generate a unique run ID when pytest starts.""" |
| 226 | + # Generate the run ID early so it's available for all tests |
| 227 | + get_run_id() |
| 228 | + |
| 229 | + |
| 230 | +def pytest_report_header(config): |
| 231 | + """Add run ID to pytest header.""" |
| 232 | + run_id = get_run_id() |
| 233 | + return f"Run ID: {run_id}, server: {server}" |
| 234 | + |
| 235 | + |
| 236 | +def pytest_configure_node(node): |
| 237 | + """ |
| 238 | + Configure worker nodes to inherit the run ID from the main process. |
| 239 | +
|
| 240 | + This hook is called by pytest-xdist to configure worker nodes. |
| 241 | + """ |
| 242 | + node.workerinput["trackflaky_run_id"] = get_run_id() |
| 243 | + |
| 244 | + |
| 245 | +@pytest.hookimpl(tryfirst=True) |
| 246 | +def pytest_sessionstart(session): |
| 247 | + """ |
| 248 | + Initialize run ID from worker input if this is a worker process. |
| 249 | +
|
| 250 | + This runs on worker nodes to receive the run ID from the main process. |
| 251 | + """ |
| 252 | + if hasattr(session.config, "workerinput"): |
| 253 | + # We're in a worker process |
| 254 | + workerinput = session.config.workerinput |
| 255 | + if "trackflaky_run_id" in workerinput: |
| 256 | + set_run_id(workerinput["trackflaky_run_id"]) |
| 257 | + |
| 258 | + |
| 259 | +@pytest.hookimpl(hookwrapper=True) |
| 260 | +def pytest_pyfunc_call(pyfuncitem): |
| 261 | + """Hook into pytest test execution to track test outcomes.""" |
| 262 | + |
| 263 | + result = get_base_result() |
| 264 | + result["testname"] = pyfuncitem.name |
| 265 | + result["start_time"] = int(time()) |
| 266 | + |
| 267 | + outcome = yield |
| 268 | + |
| 269 | + result["end_time"] = int(time()) |
| 270 | + |
| 271 | + if outcome.excinfo is None: |
| 272 | + result["outcome"] = "success" |
| 273 | + elif outcome.excinfo[0] == unittest.case.SkipTest: |
| 274 | + result["outcome"] = "skip" |
| 275 | + else: |
| 276 | + result["outcome"] = "fail" |
| 277 | + |
| 278 | + print(result) |
| 279 | + |
| 280 | + if not server: |
| 281 | + return |
| 282 | + |
| 283 | + try: |
| 284 | + req = request.Request(f"{server}/hook/test", method="POST") |
| 285 | + req.add_header("Content-Type", "application/json") |
| 286 | + |
| 287 | + request.urlopen( |
| 288 | + req, |
| 289 | + data=json.dumps(result).encode("ASCII"), |
| 290 | + ) |
| 291 | + except ConnectionError as e: |
| 292 | + print(f"Could not report testrun: {e}") |
| 293 | + except Exception as e: |
| 294 | + import warnings |
| 295 | + |
| 296 | + warnings.warn(f"Error reporting testrun: {e}") |
0 commit comments