graceful shutdown for aiohttp session

bot.py

@@ -6,16 +6,40 @@ from discord.ext import commands
 import heavynode
 
 
+class CustomBot(commands.Bot):
+    def __init__(self, *args, **kwargs):
+        self.cleanup_sync = []
+        self.cleanup_async = []
+        super().__init__(*args, **kwargs)
+
+    def add_cleanup_sync(self, func):
+        self.cleanup_sync.append(func)
+
+    def add_cleanup_async(self, coro):
+        self.cleanup_async.append(coro)
+
+    async def close(self):
+        for func in self.cleanup_sync:
+            func()
+        for coro in self.cleanup_async:
+            await coro()
+        await super().close()
+
+
 logging.basicConfig()
+
+server_id = 'd030737c'
 hn = heavynode.Client(os.environ['heavynode_token'])
-bot = commands.Bot(command_prefix='!')
+
+bot = CustomBot(command_prefix='!')
+bot.add_cleanup_async(hn.shutdown)
 
 
 @bot.command()
 @commands.has_role('Server admin')
 async def add(ctx, player):
     """Add a player to the server whitelist.Must use exact Minecraft name."""
-    await hn.send_command('d030737c', f'whitelist add {player}')
+    await hn.send_command(server_id, f'whitelist add {player}')
     await ctx.send(f'"{player}" added to whitelist.')
 
 
@@ -23,7 +47,7 @@ async def add(ctx, player):
 @commands.has_role('Server admin')
 async def remove(ctx, player):
     """Remove a player from the server whitelist. Must use exact Minecraft name."""
-    await hn.send_command('d030737c', f'whitelist remove {player}')
+    await hn.send_command(server_id, f'whitelist remove {player}')
     await ctx.send(f'"{player}" removed from whitelist.')
 
 
@@ -36,4 +60,4 @@ async def whitelist_error(ctx, error):
         raise error
 
 
-bot.run(os.environ['discord_token'], reconnect=True)
+bot.run(os.environ['discord_token'], reconnect=True)

heavynode.py

@@ -34,3 +34,6 @@ class Client:
     async def send_command(self, serverid, cmd):
         payload = {'command': cmd}
         return await self.make_request('POST', f'/client/servers/{serverid}/command', json=payload)
+
+    async def shutdown(self):
+        await self.session.close()