diff --git a/src/bot.py b/src/bot.py index 9964019..c5c4f98 100644 --- a/src/bot.py +++ b/src/bot.py @@ -1,40 +1,91 @@ import os import sys +import asyncio -from discord.ext import commands +import discord +from discord.ext import commands, tasks from dotenv import load_dotenv +import youtube_dl -from yt import YTDLSource from error import ErrorHandler load_dotenv() +# Suppress noise about console usage from errors +youtube_dl.utils.bug_reports_message = lambda: '' + class Music(commands.Cog): def __init__(self, bot): - self.bot = bot + self._bot = bot + self._queue = asyncio.Queue() + self._queue_lock = asyncio.Lock() + self._ytdl_format_options = { + 'format': 'bestaudio/best', + 'outtmpl': '%(extractor)s-%(id)s-%(title)s.%(ext)s', + 'restrictfilenames': True, + 'noplaylist': True, + 'nocheckcertificate': True, + 'ignoreerrors': False, + 'logtostderr': False, + 'quiet': True, + 'no_warnings': True, + 'default_search': 'auto', + 'source_address': '0.0.0.0' # bind to ipv4 since ipv6 addresses cause issues sometimes + } + self._ffmpeg_options = { + 'options': '-vn' + } + self._ytdl = youtube_dl.YoutubeDL(self._ytdl_format_options) + # pylint: disable=no-member + self._handle_playback.start() + + def cog_unload(self): + # pylint: disable=no-member + self._handle_playback.cancel() + + @tasks.loop() + async def _handle_playback(self): + while True: + await self._queue_lock.acquire() + ctx, url, title = await self._queue.get() + audio = discord.FFmpegPCMAudio(url, **self._ffmpeg_options) + + def after(err): + if err: + print(f"Player error: {err}") + self._queue.task_done() + self._queue_lock.release() + ctx.voice_client.play(audio, after=after) + await ctx.send(f"Now playing: {title}") + + @_handle_playback.before_loop + async def before_handle_playback(self): + await self._bot.wait_until_ready() @commands.command() - async def stream(self, ctx, *, url): + async def play(self, ctx, *, query): async with ctx.typing(): - player = await YTDLSource.from_url(url, loop=self.bot.loop, stream=True) - ctx.voice_client.play(player, after=lambda e: print(f"Player error: {e}") if e else None) - - await ctx.send(f"Now playing: {player.title}") + data = self._ytdl.extract_info(query, download=False) + if 'entries' in data: + data = data['entries'][0] + title = data.get('title') + url = data.get('url') + await self._queue.put((ctx, url, title)) + if ctx.voice_client.is_playing(): + await ctx.send(f"Queued: {title}") @commands.command() async def stop(self, ctx): await ctx.voice_client.disconnect() - @stream.before_invoke + @play.before_invoke async def ensure_voice(self, ctx): if ctx.voice_client is None: if ctx.author.voice: await ctx.author.voice.channel.connect() else: raise commands.CommandError("Author not connected to a voice channel.") - elif ctx.voice_client.is_playing(): - ctx.voice_client.stop() if __name__ == "__main__": diff --git a/src/error.py b/src/error.py index 2b4b457..7f1009e 100644 --- a/src/error.py +++ b/src/error.py @@ -15,6 +15,8 @@ class ErrorHandler(commands.Cog): message = "You are missing the required permissions to run this command!" elif isinstance(error, commands.UserInputError): message = "Something about your input was wrong, please check your input and try again!" + elif isinstance(error, commands.CommandError): + message = str(error) else: message = "Oh no! Something went wrong while running the command!" diff --git a/src/yt.py b/src/yt.py deleted file mode 100644 index 8148bba..0000000 --- a/src/yt.py +++ /dev/null @@ -1,50 +0,0 @@ -import asyncio - -import discord -import youtube_dl - -# Suppress noise about console usage from errors -youtube_dl.utils.bug_reports_message = lambda: '' - - -ytdl_format_options = { - 'format': 'bestaudio/best', - 'outtmpl': '%(extractor)s-%(id)s-%(title)s.%(ext)s', - 'restrictfilenames': True, - 'noplaylist': True, - 'nocheckcertificate': True, - 'ignoreerrors': False, - 'logtostderr': False, - 'quiet': True, - 'no_warnings': True, - 'default_search': 'auto', - 'source_address': '0.0.0.0' # bind to ipv4 since ipv6 addresses cause issues sometimes -} - -ffmpeg_options = { - 'options': '-vn' -} - -ytdl = youtube_dl.YoutubeDL(ytdl_format_options) - - -class YTDLSource(discord.PCMVolumeTransformer): - def __init__(self, source, *, data, volume=0.5): - super().__init__(source, volume) - - self.data = data - - self.title = data.get('title') - self.url = data.get('url') - - @classmethod - async def from_url(cls, url, *, loop=None, stream=False): - loop = loop or asyncio.get_event_loop() - data = await loop.run_in_executor(None, lambda: ytdl.extract_info(url, download=not stream)) - - if 'entries' in data: - # take first item from a playlist - data = data['entries'][0] - - filename = data['url'] if stream else ytdl.prepare_filename(data) - return cls(discord.FFmpegPCMAudio(filename, **ffmpeg_options), data=data) diff --git a/test/conftest.py b/test/conftest.py index 1fa415c..65744e7 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,3 +1,4 @@ +import asyncio import sys from pathlib import Path from unittest.mock import Mock @@ -9,6 +10,8 @@ from pytest_mock import MockerFixture SRC_PATH = str(Path(__file__).parent / '..' / 'src') sys.path.insert(0, SRC_PATH) +from bot import Music + @pytest.fixture def bot(mocker: MockerFixture): @@ -17,6 +20,21 @@ def bot(mocker: MockerFixture): yield bot_mock +@pytest.fixture +def mbot(bot): + mbot_mock = Music(bot) + yield mbot_mock + mbot_mock.cog_unload() + + @pytest.fixture def ctx(mocker: MockerFixture): yield mocker.patch('discord.ext.commands.Context', autospec=True) + + +@pytest.fixture +def event_loop(): + loop = asyncio.get_event_loop() + yield loop + pending = asyncio.all_tasks(loop=loop) + asyncio.ensure_future(asyncio.gather(*pending)) diff --git a/test/test_bot.py b/test/test_bot.py index 1bbe70d..4595652 100644 --- a/test/test_bot.py +++ b/test/test_bot.py @@ -1,54 +1,98 @@ +import asyncio from unittest.mock import Mock, AsyncMock, patch import pytest from discord.ext import commands -from bot import Music - @pytest.mark.asyncio -async def test_bot_ensure_voice(bot, ctx): - mbot = Music(bot) - - # 1. Inside a voice channel - # 1.1 Does not call stop if no sound is playing - ctx.voice_client.is_playing.return_value = False - await mbot.ensure_voice(ctx) - assert ctx.voice_client.stop.call_count == 0 - ctx.reset_mock(return_value=True) - - # 1.2 Does call stop if sound is playing - ctx.voice_client.is_playing.return_value = True - await mbot.ensure_voice(ctx) - assert ctx.voice_client.stop.call_count == 1 - ctx.reset_mock(return_value=True) - - # 2. Not inside a voice channel - # 2.1 Connects to voice channel of author if possible +async def test_bot_ensure_voice(mbot, ctx): + # Connects to voice channel of author if possible ctx.voice_client = None ctx.author.voice = AsyncMock() await mbot.ensure_voice(ctx) - assert ctx.author.voice.channel.connect.call_count == 1 + assert ctx.author.voice.channel.connect.call_count == 1, "Did not connect to voice channel of author" ctx.reset_mock(return_value=True) - # 2.2 Error if author not inside a channel + # Error if author not inside a channel ctx.voice_client = None ctx.author.voice = None with pytest.raises(commands.CommandError): await mbot.ensure_voice(ctx) -@pytest.mark.asyncio -async def test_bot_stream(bot, ctx): - mbot = Music(bot) +def mock_ytdl_extract_info(ytdl, url, title): + ytdl.extract_info.return_value = {"entries": [{"url": url, "title": title}]} - with patch('bot.YTDLSource', new_callable=AsyncMock) as ytdl_source: - player = Mock() - ytdl_source.from_url.return_value = player - url = 'A Day To Remember - All I Want' - # pylint: disable=too-many-function-args - await mbot.stream(mbot, ctx, url=url) - assert ytdl_source.from_url.await_args.args == (url,) - assert ytdl_source.from_url.await_args.kwargs == {"loop": bot.loop, "stream": True} - assert ctx.voice_client.play.call_args.args == (player,) - assert ctx.send.call_count == 1 + +def mock_ffmpeg_pcm_audio(ffmpeg_pcm_audio): + deliver_us_audio = Mock() + ffmpeg_pcm_audio.return_value = deliver_us_audio + return deliver_us_audio + + +@pytest.mark.asyncio +async def test_bot_playback(mbot, ctx): + with patch.object(mbot, '_ytdl') as ytdl: + with patch('discord.FFmpegPCMAudio') as ffmpeg_pcm_audio: + ctx.voice_client.is_playing.return_value = False + url = "https://www.youtube.com/watch?v=Wr9LZ1hAFpQ" + title = "In Flames - Deliver Us (Official Video)" + mock_ytdl_extract_info(ytdl, url, title) + deliver_us_audio = mock_ffmpeg_pcm_audio(ffmpeg_pcm_audio) + query = 'in flames deliver us' + # pylint: disable=too-many-function-args + await mbot.play(mbot, ctx, query=query) + assert \ + ytdl.extract_info.call_args.args == (query,) and ytdl.extract_info.call_args.kwargs == {"download": False}, \ + f"ytdl.extract_info was not called with {query}, {{ download: False }}" + assert \ + ffmpeg_pcm_audio.call_args is None, \ + f"FFmpegPCMAudio was immediately called with {url} instead of being queued" + assert \ + ctx.voice_client.play.call_args is None, \ + "Did immediately playback audio instead of being queued" + assert \ + ctx.send.call_args is None, \ + "Did immediately send 'Now playing:' message of being queued" + await asyncio.sleep(0) + assert \ + ffmpeg_pcm_audio.call_args.args == (url,), \ + f"FFmpegPCMAudio was not called with {url}" + assert \ + ctx.voice_client.play.call_args.args == (deliver_us_audio,), \ + "Did not playback correct audio" + assert \ + ctx.send.call_args.args == (f"Now playing: {title}",), \ + "Did not send 'Now playing:' message" + + ctx.voice_client.is_playing.return_value = True + url = "https://www.youtube.com/watch?v=pMDcYX2wRSg" + title = "Three Days Grace - Time of Dying (lyrics)" + mock_ytdl_extract_info(ytdl, url, title) + time_of_dying_audio = mock_ffmpeg_pcm_audio(ffmpeg_pcm_audio) + # pylint: disable=too-many-function-args + query = "three days grace time of dying" + await mbot.play(mbot, ctx, query=query) + assert \ + ytdl.extract_info.call_args.args == (query,) and ytdl.extract_info.call_args.kwargs == {"download": False}, \ + f"ytdl.extract_info was not called with {query}, {{ download: False }}" + assert \ + not ffmpeg_pcm_audio.call_args.args == (url,), \ + f"FFmpegPCMAudio was immediately called with {url} instead of being queued" + assert \ + not ctx.voice_client.play.call_args.args == (time_of_dying_audio,), \ + "Did immediately playback audio instead of being queued" + assert \ + ctx.send.call_args.args == (f"Queued: {title}",), \ + "Did not send 'Queued:' message" + await asyncio.sleep(0) + # Still no playback because previous song not finished + assert \ + not ffmpeg_pcm_audio.call_args.args == (url,), \ + f"FFmpegPCMAudio was called with {url} before previous song finished" + ctx.voice_client.play.call_args.kwargs["after"](None) # Execute callback for song finish event + await asyncio.sleep(0) + assert \ + ffmpeg_pcm_audio.call_args.args == (url,), \ + f"FFmpegPCMAudio was not called with {url}"