Source code for pyarchops_helpers.helpers
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Helpers for `pyarchops tests` packages."""
import os
import re
import time
import socket
from contextlib import contextmanager
from typing import Iterator
from fabric.api import local
from fabric.context_managers import settings, quiet
import paramiko
from paramiko.ssh_exception import (
AuthenticationException,
BadHostKeyException,
SSHException
)
import netifaces
[docs]def wait_for_ssh(
host: dict,
initial_wait: int = 0,
interval: int = 0,
retries: int = 1
) -> bool:
""" waits for ssh to become available """
ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
time.sleep(initial_wait)
for _ in range(retries):
try:
ssh.connect(host['ip'], username=host['user'], port=host['port'])
return True
except (AuthenticationException, BadHostKeyException):
return True
except (SSHException, socket.error) as error:
print(error)
time.sleep(interval)
return False
[docs]@contextmanager
def ephemeral_docker_container(**kwargs: dict) -> Iterator[dict]:
"""
prepares a docker container, yelding a dict
with its configuration before deleting the container
"""
try:
image = str(kwargs['image'])
container_id = docker_up(image)
yield dict({
'user': 'root',
'ip': dockerd_ip(),
'port': docker_container_port(container_id),
'pkey': 'tests/fixtures/id_rsa'
})
docker_down(container_id)
docker_rm(container_id)
except Exception as error:
docker_down(container_id)
docker_rm(container_id)
raise error
[docs]def docker_up(image: str, privileged: bool = False) -> str:
"""
runs a docker container
params:
string image: name of the docker image
bool privileged: use docker --privileged flag
returns:
string: stdout of the docker run
"""
flags = ''
if privileged:
flags = '--privileged'
container_id = local(
"docker run --privileged -ti "
" --cpu-period=100000 --cpu-quota=10000 "
" -v /sys/fs/cgroup:/sys/fs/cgroup:ro "
" -d -P {} {}".format(flags, image),
capture=True
)
host = {
'ip': dockerd_ip(),
'user': 'root',
'port': docker_container_port(container_id)
}
wait_for_ssh(
host=host,
initial_wait=1,
retries=10,
interval=1
)
return str(container_id)
[docs]def docker_down(container_id: str) -> None:
"""
kills the docker container
params:
string container: docker id of the container to stop
"""
with settings(quiet()):
local('docker kill %s' % container_id)
[docs]def docker_rm(container_id: str) -> None:
"""
removes a docker container
params:
string container: docker id of the container to remove
"""
with settings(quiet()):
local('docker rm --force %s' % container_id)
[docs]def docker_container_port(container_id: str) -> int:
"""
returns the ssh port number for a docker instance
params:
string container: docker container id
returns:
string: port number
"""
with settings(quiet()):
output = local(
'docker port %s 22' % container_id,
capture=True
)
return int(output.split(':')[1])
[docs]def dockerd_ip() -> str:
"""
returns the ip address of the docker daemon
params:
string docker_host_string: URL of the docker daemon
returns:
string: ip address of the docker host
"""
if 'DOCKER_HOST' in os.environ:
docker_host_env = os.environ.get('DOCKER_HOST')
ipaddr = re.search( # type: ignore
r'[tcp|http]:\/\/(.*):\d.*', docker_host_env).group(1) # noqa
else:
ipaddr = netifaces.ifaddresses(
'docker0')[netifaces.AF_INET][0]['addr']
return str(ipaddr)