| # Copyright 2020 The Fuchsia Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| import attr |
| import datetime |
| import enum |
| |
| from recipe_engine import recipe_api |
| |
| # The format used for the "date" field in HTTP responses from requests to the |
| # tree status page. |
| TREE_STATUS_DATE_FORMAT = "%Y-%m-%d %H:%M:%S.%f" |
| |
| |
| class State(enum.Enum): |
| OPEN = "open" |
| CLOSED = "closed" |
| THROTTLED = "throttled" |
| MAINTENANCE = "maintenance" |
| |
| |
| @attr.s |
| class TreeStatus(object): |
| state = attr.ib(type=str) |
| date = attr.ib(type=datetime.datetime) |
| username = attr.ib(type=str) |
| # Unique ID of this status. |
| key = attr.ib(type=int) |
| |
| @classmethod |
| def from_dict(cls, general_state, date, username, key, **_): |
| """Create a TreeStatus from a dict, ignoring any unneeded values.""" |
| return cls( |
| state=State(general_state), |
| date=datetime.datetime.strptime(date, TREE_STATUS_DATE_FORMAT), |
| username=username, |
| key=key, |
| ) |
| |
| @property |
| def open(self): |
| return self.state in (State.OPEN, State.THROTTLED) |
| |
| |
| class TreeStatusApi(recipe_api.RecipeApi): |
| @property |
| def _tree_status_script(self): |
| return self.resource("tree_status.py") |
| |
| def get(self, hostname, step_name="get current tree status"): |
| """Get the current tree status. |
| |
| Args: |
| hostname (str): Hostname of the tree status page. |
| |
| Returns: |
| A `TreeStatus` corresponding to the current state of the tree. |
| """ |
| step = self.m.step( |
| step_name, |
| [self._tree_status_script, hostname, "get"], |
| stdout=self.m.json.output(), |
| step_test_data=self.test_api.status_step_data, |
| ) |
| status = TreeStatus.from_dict(**step.stdout) |
| step.presentation.step_text = status.state.value |
| step.presentation.links[hostname] = "https://%s" % hostname |
| return status |
| |
| def update( |
| self, |
| hostname, |
| password, |
| message, |
| username=None, |
| last_status=None, |
| step_name="update tree status", |
| ): |
| """Change the tree status text. |
| |
| Args: |
| hostname (str): Hostname of the tree status page. |
| password (str): Password for the service account to access the |
| tree status page. |
| message (str): Message to set as the tree status. |
| username (str): Name or email to use as the author of the tree |
| status. |
| last_status (TreeStatus or None): Only update the tree status if it |
| has not changed since the specified status change. |
| """ |
| # Collision checking *should* be handled by the tree status app itself, |
| # but it doesn't do collision checking for status changes that use the |
| # bot-specific endpoint (as opposed to the web UI form): |
| # https://chromium.googlesource.com/infra/infra/+/09d53b324dd786b96d69c6cbb8ee6c389b22fd3f/appengine/chromium_status/appengine_module/chromium_status/status.py#426 |
| if last_status: |
| current_status = self.get( |
| hostname, step_name="check for tree status collision" |
| ) |
| if current_status.key != last_status.key: |
| raise self.m.step.StepFailure( |
| "mid-air collision detected between tree status changes" |
| ) |
| return self.m.step( |
| step_name, |
| [ |
| self._tree_status_script, |
| hostname, |
| "set", |
| message, |
| "--username", |
| username or self.m.buildbucket.builder_name, |
| "--password", |
| password, |
| ], |
| ) |