add content hosts to satelite

mass_hosts.py

import asyncio
from socket import timeout
import aiohttp
import json
import uuid
import click
from functools import wraps


def coro(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        return asyncio.run(f(*args, **kwargs))

    return wrapper


SEMAPHORE = asyncio.Semaphore(100)


def gen_name():
    return str(uuid.uuid4()).split("-")[-1]


async def add_host(session, hostname, cv_id=None, sub_id=None, progress=None, pos=None):
    global SEMAPHORE
    auth = aiohttp.BasicAuth("admin", "changeme")
    try:
        async with session.post(
            f"https://{hostname}/api/v2/hosts/subscriptions",
            data=json.dumps(
                {
                    "name": gen_name(),
                    "lifecycle_environment_id": 1,
                    "content_view_id": cv_id,
                }
            ),
            ssl=False,
            headers={"content-type": "application/json"},
            auth=auth,
        ) as reg:
            host_id = (await reg.json()).get("id")
    except asyncio.exceptions.TimeoutError:
        host_id = None
    if sub_id:
        try:
            async with session.put(
                f"https://{hostname}/api/hosts/{host_id}/subscriptions/add_subscriptions",
                data=json.dumps({"subscriptions": [{"id": sub_id}]}),
                ssl=False,
                headers={"content-type": "application/json"},
                auth=auth,
            ) as sub:
                pass
        except asyncio.exceptions.TimeoutError:
            pass
    if progress and (pos is not None):
        progress.update(pos)


@click.command()
@click.option("--hostname", type=str, help="Hostname of the Satellite")
@click.option("--cv-id", type=int, help="Content View id")
@click.option("--sub-id", type=int, help="Subscription id")
@click.option("--count", type=int, default=10, help="Number of hosts to add")
@click.option(
    "--max-connections",
    type=int,
    default=10,
    help="Max number of concurrent connections",
)
@coro
async def add_hosts(hostname, cv_id, sub_id, count, max_connections):
    """Add hosts to a Satellite. If a subscription is provided, add that to the host."""
    with click.progressbar(length=count) as progress:
        async with aiohttp.ClientSession(
            connector=aiohttp.TCPConnector(limit=max_connections, ssl=False),
            timeout=aiohttp.ClientTimeout(total=180)
        ) as session:
            print(f"Starting registration of {count} hosts")
            # await asyncio.gather(
            #     *[
            #         add_host(session, hostname, cv_id, sub_id, progress, pos)
            #         for pos in range(count)
            #     ]
            # )
            tasks = []
            for pos in range(count):
                task = asyncio.ensure_future(add_host(session, hostname, cv_id, sub_id, progress, pos))
                tasks.append(task)
            responses = asyncio.gather(*tasks)
            await responses
        # for c in range(count):
        #     add_host(hostname, cv_id, sub_id)
        #     progress.update(c)


if __name__ == "__main__":
    add_hosts()

# if __name__ == "__main__":
#     hostname = "dhcp-3-231.vms.sat.rdu2.redhat.com"
#     cv_id = 3
#     add_host(hostname, cv_id)
#     add_host(hostname, cv_id, 1)