Skip to content

Commit

Permalink
Add generic wait component (#255)
Browse files Browse the repository at this point in the history
This change adds a wait component that can wait based on time, a url
status, or a command result.

Signed-off-by: Alex Schultz <[email protected]>
  • Loading branch information
mwhahaha authored Oct 13, 2021
1 parent cb559d1 commit e7d8015
Show file tree
Hide file tree
Showing 3 changed files with 307 additions and 0 deletions.
202 changes: 202 additions & 0 deletions directord/components/builtin_wait.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,202 @@
# Copyright Alex Schultz <[email protected]>. All Rights Reserved
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

import requests
import time

from directord import components


class Component(components.ComponentBase):
def __init__(self):
"""Initialize the component wait class."""

super().__init__(desc="Wait until a condition is met")
self.requires_lock = False

def args(self):
"""Set default arguments for a component."""

super().args()
condition_group = self.parser.add_mutually_exclusive_group(
required=True
)
condition_group.add_argument(
"--seconds",
type=int,
help="Wait for the provided seconds",
)
condition_group.add_argument(
"--url", type=str, help="Wait for URL to return 2xx or 3xx"
)
condition_group.add_argument(
"--cmd",
action="store_true",
help="Wait for the provided command to returns successful",
)
self.parser.add_argument(
"--retry",
default=0,
type=int,
help="Number of times to retry condition (ignored with --seconds)",
)
self.parser.add_argument(
"--retry-wait",
default=0,
type=int,
help="Time to wait between retries(ignored with --seconds)",
)
self.parser.add_argument(
"--insecure",
action="store_true",
help="Allow insecure server connections when using SSL",
)

def server(self, exec_array, data, arg_vars):
"""Return data from formatted transfer action.
:param exec_array: Input array from action
:type exec_array: List
:param data: Formatted data hash
:type data: Dictionary
:param arg_vars: Pre-Formatted arguments
:type arg_vars: Dictionary
:returns: Dictionary
"""

super().server(exec_array=exec_array, data=data, arg_vars=arg_vars)
if self.known_args.seconds:
data["seconds"] = self.known_args.seconds
elif self.known_args.url:
data["url"] = self.known_args.url
elif self.known_args.cmd:
data["command"] = " ".join(self.unknown_args)
data["retry"] = self.known_args.retry
data["retry_wait"] = self.known_args.retry_wait
data["insecure"] = self.known_args.insecure
return data

def client(self, cache, job):
"""Wait for condition.
Command operations are rendered with cached data from the args dict.
:param cache: Caching object used to template items within a command.
:type cache: Object
:param job: Information containing the original job specification.
:type job: Dictionary
:returns: tuple
"""

self.log.debug("client(): job: %s, cache: %s", job, cache)
seconds = job.get("seconds")
url = job.get("url")
cmd = job.get("command")
retry = job.get("retry", 0)
retry_wait = job.get("retry_wait", 0)
insecure = job.get("insecure", False)

out = b""
err = b""
success = True
msg = None
if seconds is not None:
time.sleep(seconds)
return out, err, success, msg
elif url is not None:
out, err, success = self._fetch_url(
url, not insecure, retry, retry_wait
)
if not success:
msg = f"URL did not return a 2xx or 3xx. Retired {retry} times"
return out, err, success, msg
elif cmd is not None:
out, err, success = self._run_cmd(
cmd, cache.get("envs"), retry, retry_wait
)
if not success:
msg = f"Command was not successful. Retried {retry} times."
return out, err, success, msg
else:
self.log.error("Invalid wait condition provided")
return out, err, False, None

def _fetch_url(
self, url: str, verify: bool, retry: int = 0, retry_wait: int = 0
):
"""Fetch url with retry.
Fetch a url and return True if response code is 2xx or 3xx.
:param url: Url string to fetch
:type url: String
:param verify: Boolean to manage ssl validation
:type verify: Boolean
:param retry: Number of retries on failure
:type retry: Integer
:param: retry_wait: Number of seconds to wait between retry
:type retry_wait: Integer
:returns: tuple
"""
stdout = b""
stderr = b""
outcome = False
count = 0
while not outcome and count < (retry + 1):
count = count + 1
try:
r = requests.get(url, verify=verify)
stdout = f"Response code was {r.response_code}"
if r.response_code >= 200 and r.response_code < 400:
outcome = True
except Exception as e:
stderr = f"Exception occured while fetching url {e}"
self.log.error(e)
if not outcome and retry_wait > 0:
self.log.debug("Url fetch failed, retrying with wait...")
time.sleep(retry_wait)

return stdout, stderr, outcome

def _run_cmd(
self, cmd: str, env: dict, retry: int = 0, retry_wait: int = 0
):
"""Run command with retry.
Run a command and if not successful, retry.
:param cmd: Command string to run
:type cmd: String
:param env: Environment dict to pass when running
:type env: Dictionary
:param retry: Number of retries on failure
:type retry: Integer
:param: retry_wait: Number of seconds to wait between retry
:type retry_wait: Integer
:returns: tuple
"""
job_stdout = []
job_stderr = []
outcome = False
count = 0
while not outcome and count < (retry + 1):
count = count + 1
stdout, stderr, outcome = self.run_command(command=cmd, env=env)
job_stdout.append(stdout)
job_stderr.append(stderr)
if not outcome and retry_wait > 0:
self.log.debug("Command failed, retrying with wait...")
time.sleep(retry_wait)

return b"".join(job_stdout), b"".join(job_stderr), outcome
87 changes: 87 additions & 0 deletions directord/tests/test_components.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
from directord.components import builtin_run
from directord.components import builtin_workdir
from directord.components import builtin_queuesentinel
from directord.components import builtin_wait


class TestComponents(unittest.TestCase):
Expand All @@ -50,6 +51,7 @@ def setUp(self):
self._run = builtin_run.Component()
self._workdir = builtin_workdir.Component()
self._queuesentinal = builtin_queuesentinel.Component()
self._wait = builtin_wait.Component()
for item in [
self._dnf,
self._service,
Expand Down Expand Up @@ -636,3 +638,88 @@ def test_blueprinter_failed(self, mock_log_debug):
blueprinted_content, "Can't compile non template nodes"
)
mock_log_debug.assert_called()

@patch("time.sleep")
def test_wait_seconds(self, mock_sleep):
stdout, stderr, outcome, return_info = self._wait.client(
cache=tests.FakeCache(), job={"seconds": 5}
)
self.assertTrue(outcome)

mock_sleep.assert_called_once_with(5)

@patch("requests.get")
@patch("time.sleep")
def test_wait_url(self, mock_sleep, mock_get):
r_500 = MagicMock()
r_500.response_code = 500
r_400 = MagicMock()
r_400.response_code = 400
r_200 = MagicMock()
r_200.response_code = 200
mock_get.side_effect = [r_500, r_400, r_200]

stdout, stderr, outcome, return_info = self._wait.client(
cache=tests.FakeCache(),
job={"url": "http://localhost", "retry": 5, "retry_wait": 5},
)
self.assertTrue(outcome)

sleep_calls = [call(5), call(5)]
self.assertEqual(mock_sleep.mock_calls, sleep_calls)
get_calls = [
call("http://localhost", verify=True),
call("http://localhost", verify=True),
call("http://localhost", verify=True),
]
self.assertEqual(mock_get.mock_calls, get_calls)

@patch("requests.get")
@patch("time.sleep")
def test_wait_url_fail(self, mock_sleep, mock_get):
r_500 = MagicMock()
r_500.response_code = 500
mock_get.return_value = r_500
stdout, stderr, outcome, return_info = self._wait.client(
cache=tests.FakeCache(),
job={"url": "http://localhost", "insecure": True},
)
self.assertFalse(outcome)
mock_get.assert_called_once_with("http://localhost", verify=False)

@patch("directord.components.ComponentBase.run_command", autospec=True)
@patch("time.sleep")
def test_wait_cmd(self, mock_sleep, mock_run_command):
mock_run_command.side_effect = [
(b"", b"", False),
(b"", b"", True),
]
stdout, stderr, outcome, return_info = self._wait.client(
cache=tests.FakeCache(),
job={
"command": "curl -k http://google.com",
"retry": 5,
"retry_wait": 5,
},
)
self.assertTrue(outcome)

mock_sleep.assert_called_once_with(5)
run_calls = [
call(command="curl -k http://google.com", env=None),
call(command="curl -k http://google.com", env=None),
]
self.assertEqual(mock_run_command.mock_calls, run_calls)

@patch("directord.components.ComponentBase.run_command", autospec=True)
@patch("time.sleep")
def test_wait_cmd_fail(self, mock_sleep, mock_run_command):
mock_run_command.return_value = (b"", b"", False)

stdout, stderr, outcome, return_info = self._wait.client(
cache=tests.FakeCache(), job={"command": "foo"}
)
self.assertFalse(outcome)

mock_sleep.assert_not_called()
mock_run_command.assert_called_once_with(command="foo", env=None)
18 changes: 18 additions & 0 deletions docs/components.md
Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,24 @@ Block client task execution until a specific job has entered a completed state.
> `JOB_WAIT` requires the job SHA to block. This is most useful for component developers
using callback jobs.

##### `WAIT`

Syntax: `[CMD ...]`

Conditional wait based on time, url, or command

* `--seconds` **INTEGER** Wait for the provided seconds.
* `--url` **STRING** URL to fetch and check for a 2xx or 3xx response.
* `--cmd` **STRING** Run the provided command string and check success.
* `--retry` **INTEGER** Number of retries on failure. Default: 0
* `--retry-wait` **INTEGER** Number of seconds to wait before retrying. Default: 0
* `--insecure` **BOOLEAN** Allow insecure service connections when using SSL (works only with --url).

> The options `--seconds`, `--url`, and `--cmd` are mutually exclusive and one of
the options must be provided. When `--cmd` is specified any additional that are
not specifically parameters are assumes to be a command to run similar to the RUN
component.

### User defined Components

User defined components are expected to be in the
Expand Down

0 comments on commit e7d8015

Please sign in to comment.