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