async def test_pool_spawn(client):
""" Test context manager. """
pool = AIOConnectionPool(client, minconn=1, maxconn=1)
assert pool.idle_connection == 0
async with pool.spawn() as conn:
assert pool.shared_connection == 1
_ = await conn.whoami()
assert pool.idle_connection == 1
assert pool.shared_connection == 0
async def test_pool_close(client):
""" Test closing the pool. """
pool = AIOConnectionPool(client, minconn=1, maxconn=1)
await pool.open()
assert pool.closed == False
assert pool.idle_connection == 1
await pool.close()
assert pool.closed == True
assert pool.idle_connection == 0
async def test_pool_get_put(client):
""" Test getting and putting back connection from pool. """
delay = 2
pool = AIOConnectionPool(client, minconn=1, maxconn=1)
with pytest.raises(ClosedPool):
_ = await pool.get()
await pool.open()
assert pool.closed == False
assert pool.idle_connection == 1
task1 = asyncio.ensure_future(keep(pool, delay))
task2 = asyncio.ensure_future(keep(pool, delay))
start = time.time()
await task1
await task2
assert time.time() - start >= delay * 2
# With enough connection in the pool for both tasks.
pool.max_connection = 2
task1 = asyncio.ensure_future(keep(pool, delay))
task2 = asyncio.ensure_future(keep(pool, delay))
start = time.time()
await task1
await task2
assert time.time() - start >= delay
async def test_pool_pass_param(client):
""" Test passing parameter to connect. """
pool = AIOConnectionPool(client, minconn=1, maxconn=1, timeout=0)
with pytest.raises(asyncio.TimeoutError):
await pool.open()
_ = await pool.get()
def main(args: List[str]) -> int:
parser = setup_argparse()
parsed_args = parser.parse_args(args)
event_loop = get_event_loop()
logging.basicConfig(level=logging.INFO)
# Initialize singletons
try:
config = ConfigSchema.from_json_file(
file_handle=parsed_args.config_file)
except KeyError as err:
print("Missing section in config:", err)
return 1
# Set up DB pool
db_conn_pool: Pool = event_loop.run_until_complete(
create_pool(
host=config.database.host,
port=config.database.port,
user=config.database.user,
password=config.database.password,
db=config.database.dbname,
loop=event_loop,
charset="utf8mb4",
use_unicode=True,
))
# Nothing else is needed for migrations - check the command
if parsed_args.command == "migrate":
event_loop.run_until_complete(migrate(db_conn_pool))
return 0
# Async connect broke when TLS is enabled
# See https://github.com/noirello/bonsai/issues/25
if config.ldap.use_tls:
set_connect_async(False)
ldap_client = LDAPClient(url=config.ldap.url, tls=config.ldap.use_tls)
ldap_connection_pool = AIOConnectionPool(client=ldap_client)
slack_client = WebClient(token=config.slack.api_token, run_async=True)
# Setup LDAP pool
ldap_client.set_credentials(
mechanism="SIMPLE",
user=config.ldap.bind_user,
password=config.ldap.bind_password,
)
event_loop.run_until_complete(ldap_connection_pool.open())
# Integrations
google_api = GoogleAPIIntegration(config=config)
ggroups_db = GoogleGroupsDatabaseIntegration(db_conn_pool=db_conn_pool)
requests_db = RequestsDatabaseIntegration(db_conn_pool=db_conn_pool)
schedule_db = ScheduleDatabaseIntegration(db_conn_pool=db_conn_pool)
# Tasks
ldap_load_task = LDAPLoadTask(ldap_conn_pool=ldap_connection_pool,
ldap_config=config.ldap)
# Controllers
ldap_controller = LDAPController(load_task=ldap_load_task)
ggroups_controller = GoogleGroupsController(ggroups_db=ggroups_db,
google_api=google_api)
request_controller = RequestController(
requests_db=requests_db,
ldap=ldap_controller,
config=config,
)
schedule_controller = ScheduleController(schedule_db=schedule_db)
slack_action_controller = SlackActionController(
client=slack_client,
ldap=ldap_controller,
ggroups=ggroups_controller,
request=request_controller,
schedule=schedule_controller,
config=config,
)
slack_event_controller = SlackEventController(client=slack_client,
ggroups=ggroups_controller,
request=request_controller)
verify_controller = SlackVerifyController(slack_config=config.slack,
path="/slack")
# Add singletons to the app context
# This allows views to read them in on initalisation
app = web.Application(middlewares=[verify_controller.verify_request])
app["ConfigSchema"] = config
app["GoogleGroupsController"] = ggroups_controller
app["LDAPController"] = ldap_controller
app["RequestController"] = request_controller
app["SlackActionController"] = slack_action_controller
app["SlackEventController"] = slack_event_controller
app["ScheduleController"] = schedule_controller
# Setup scheduled tasks
scheduler = TaskScheduler()
scheduler.add_task("Refresh LDAP Data", 2, ldap_controller.sync)
scheduler.add_task("Refresh Google Groups Data", 2,
ggroups_controller.sync)
scheduler.add_task("Load schedule from DB", 1, schedule_controller.sync)
# Do one initial run of all scheduled tasks
print("Performing initial run of tasks")
event_loop.run_until_complete(scheduler.run_all())
# Setup routes
setup_routes(app)
# Go go go!
ensure_future(scheduler.run_scheduler(), loop=event_loop)
web.run_app(app, path=config.sockfile, host=config.host, port=config.port)
return 0