diff --git a/tests/end_to_end/conftest.py b/tests/end_to_end/conftest.py index 6d1ad45b75..4136144927 100644 --- a/tests/end_to_end/conftest.py +++ b/tests/end_to_end/conftest.py @@ -12,6 +12,7 @@ from tests.end_to_end.utils.logger import configure_logging from tests.end_to_end.utils.logger import logger as log from tests.end_to_end.utils.conftest_helper import parse_arguments +import tests.end_to_end.utils.docker_helper as dh def pytest_addoption(parser): @@ -191,6 +192,11 @@ def pytest_sessionfinish(session, exitstatus): if os.path.exists(cache_dir): shutil.rmtree(cache_dir, ignore_errors=False) log.debug(f"Cleared .pytest_cache directory at {cache_dir}") + + # Cleanup docker containers related to aggregator and collaborators, if any. + dh.cleanup_docker_containers(list_of_containers=["aggregator", "collaborator*"]) + # Cleanup docker network created for openfl, if any. + dh.remove_docker_network(["openfl"]) def pytest_configure(config): diff --git a/tests/end_to_end/utils/docker_helper.py b/tests/end_to_end/utils/docker_helper.py index bfb8c214cb..d8178652c0 100644 --- a/tests/end_to_end/utils/docker_helper.py +++ b/tests/end_to_end/utils/docker_helper.py @@ -12,35 +12,40 @@ log = logging.getLogger(__name__) -def remove_docker_network(): +def remove_docker_network(list_of_networks=[constants.DOCKER_NETWORK_NAME]): """ Remove docker network. + Args: + list_of_networks (list): List of network names to remove. """ client = get_docker_client() - networks = client.networks.list(names=[constants.DOCKER_NETWORK_NAME]) + networks = client.networks.list(names=list_of_networks) if not networks: - log.debug(f"Network {constants.DOCKER_NETWORK_NAME} does not exist") + log.debug(f"Network(s) {list_of_networks} does not exist") return for network in networks: log.debug(f"Removing network: {network.name}") network.remove() - log.debug("Docker network removed successfully") + log.debug(f"Docker network(s) {list_of_networks} removed successfully") -def create_docker_network(): +def create_docker_network(list_of_networks=[constants.DOCKER_NETWORK_NAME]): """ Create docker network. + Args: + list_of_networks (list): List of network names to create. """ client = get_docker_client() - networks = client.networks.list(names=[constants.DOCKER_NETWORK_NAME]) + networks = client.networks.list(names=list_of_networks) if networks: - log.info(f"Network {constants.DOCKER_NETWORK_NAME} already exists") + log.info(f"Network(s) {list_of_networks} already exists") return - log.debug(f"Creating network: {constants.DOCKER_NETWORK_NAME}") - network = client.networks.create(constants.DOCKER_NETWORK_NAME) - log.info(f"Network {network.name} created successfully") + for network_name in list_of_networks: + log.debug(f"Creating network: {network_name}") + _ = client.networks.create(network_name) + log.info(f"Docker network(s) {list_of_networks} created successfully") def check_docker_image(): @@ -143,24 +148,24 @@ def get_docker_client(): return client -def cleanup_docker_containers(): +def cleanup_docker_containers(list_of_containers=["aggregator", "collaborator*"]): """ Cleanup the docker containers meant for openfl. + Args: + list_of_containers: List of container names to cleanup. """ log.debug("Cleaning up docker containers") client = get_docker_client() - # List all containers related to openfl - agg_containers = client.containers.list(all=True, filters={"name": "aggregator"}) - col_containers = client.containers.list(all=True, filters={"name": "collaborator*"}) - containers = agg_containers + col_containers - container_names = [] - # Stop and remove all containers - for container in containers: - container.stop() - container.remove() - container_names.append(container.name) - - if containers: - log.info(f"Docker containers {container_names} cleaned up successfully") + for container_name in list_of_containers: + containers = client.containers.list(all=True, filters={"name": container_name}) + container_names = [] + # Stop and remove all containers + for container in containers: + container.stop() + container.remove() + container_names.append(container.name) + + if containers: + log.info(f"Docker containers {container_names} cleaned up successfully") diff --git a/tests/end_to_end/utils/federation_helper.py b/tests/end_to_end/utils/federation_helper.py index 9d3066c514..58fc9ba64f 100644 --- a/tests/end_to_end/utils/federation_helper.py +++ b/tests/end_to_end/utils/federation_helper.py @@ -689,7 +689,7 @@ def setup_collaborator_data(collaborators, model_name, local_bind_path): model_name (str): Model name local_bind_path (str): Local bind path """ - if not pre_existing_data(collaborators, model_name, local_bind_path): + if not pre_existing_data(collaborators): download_data(collaborators, model_name, local_bind_path) log.info("Data setup is complete for all the collaborators") @@ -709,19 +709,12 @@ def download_data(collaborators, model_name, local_bind_path): raise ex.DataSetupException(f"Failed to copy data setup file: {e}") log.info("Downloading the data for the model. This will take some time to complete based on the data size ..") - error_msg = f"Failed to download data for {model_name}" try: - return_code, _, error = run_command( - f"python -v {constants.DATA_SETUP_FILE} {len(collaborators)}", - workspace_path=local_bind_path, - error_msg=error_msg, - return_error=True, - ) - if return_code !=0 or error: - raise ex.DataSetupException(f"{error_msg}: {error}") - + # Data setup file has logic to show progress bar + # Using os.system to make sure it is displayed + os.system(f"cd {local_bind_path}; python {constants.DATA_SETUP_FILE} {len(collaborators)}") except Exception: - raise ex.DataSetupException(f"Failed to download data for {model_name}") # Do not print error, it maybe too long + raise ex.DataSetupException(f"Failed to download data for {model_name}") try: # Move the data to the respective workspace based on the index