diff --git a/tests/test_user_mute_strategy.py b/tests/test_user_mute_strategy.py new file mode 100644 index 000000000..674a273b5 --- /dev/null +++ b/tests/test_user_mute_strategy.py @@ -0,0 +1,139 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import unittest + +from pipecat.frames.frames import ( + BotStartedSpeakingFrame, + BotStoppedSpeakingFrame, + FunctionCallCancelFrame, + FunctionCallFromLLM, + FunctionCallResultFrame, + FunctionCallsStartedFrame, + InterruptionFrame, +) +from pipecat.turns.mute.always_user_mute_strategy import AlwaysUserMuteStrategy +from pipecat.turns.mute.first_speech_user_mute_strategy import FirstSpeechUserMuteStrategy +from pipecat.turns.mute.function_call_user_mute_strategy import FunctionCallUserMuteStrategy +from pipecat.turns.mute.mute_until_first_bot_complete_user_mute_strategy import ( + MuteUntilFirstBotCompleteUserMuteStrategy, +) + + +class TestAlwaysUserMuteStrategy(unittest.IsolatedAsyncioTestCase): + async def test_user_mute_strategy(self): + strategy = AlwaysUserMuteStrategy() + + self.assertTrue(await strategy.process_frame(BotStartedSpeakingFrame())) + self.assertTrue(await strategy.process_frame(InterruptionFrame())) + self.assertFalse(await strategy.process_frame(BotStoppedSpeakingFrame())) + self.assertFalse(await strategy.process_frame(InterruptionFrame())) + + +class TestFirstSpeechUserMuteStrategy(unittest.IsolatedAsyncioTestCase): + async def test_user_mute_strategy(self): + strategy = FirstSpeechUserMuteStrategy() + + self.assertFalse(await strategy.process_frame(InterruptionFrame())) + self.assertTrue(await strategy.process_frame(BotStartedSpeakingFrame())) + self.assertTrue(await strategy.process_frame(InterruptionFrame())) + self.assertFalse(await strategy.process_frame(BotStoppedSpeakingFrame())) + self.assertFalse(await strategy.process_frame(InterruptionFrame())) + + +class TestMuteUntilFirstBotCompleteUserMuteStrategy(unittest.IsolatedAsyncioTestCase): + async def test_user_mute_strategy(self): + strategy = MuteUntilFirstBotCompleteUserMuteStrategy() + + self.assertTrue(await strategy.process_frame(InterruptionFrame())) + self.assertTrue(await strategy.process_frame(BotStartedSpeakingFrame())) + self.assertTrue(await strategy.process_frame(InterruptionFrame())) + self.assertFalse(await strategy.process_frame(BotStoppedSpeakingFrame())) + self.assertFalse(await strategy.process_frame(InterruptionFrame())) + + +class TestFunctionCallUserMuteStrategy(unittest.IsolatedAsyncioTestCase): + async def test_user_mute_strategy(self): + strategy = FunctionCallUserMuteStrategy() + + self.assertFalse(await strategy.process_frame(InterruptionFrame())) + # First function call (cancelled) + self.assertTrue( + await strategy.process_frame( + FunctionCallsStartedFrame( + function_calls=[ + FunctionCallFromLLM( + function_name="fn_1", tool_call_id="1", arguments={}, context=None + ) + ] + ) + ) + ) + self.assertTrue(await strategy.process_frame(InterruptionFrame())) + self.assertFalse( + await strategy.process_frame( + FunctionCallCancelFrame(function_name="fn_1", tool_call_id="1") + ) + ) + self.assertFalse(await strategy.process_frame(InterruptionFrame())) + + # Second function call (finished) + self.assertTrue( + await strategy.process_frame( + FunctionCallsStartedFrame( + function_calls=[ + FunctionCallFromLLM( + function_name="fn_2", tool_call_id="2", arguments={}, context=None + ) + ] + ) + ) + ) + self.assertTrue(await strategy.process_frame(InterruptionFrame())) + self.assertFalse( + await strategy.process_frame( + FunctionCallResultFrame( + function_name="fn_2", tool_call_id="2", arguments={}, result={} + ) + ) + ) + self.assertFalse(await strategy.process_frame(InterruptionFrame())) + + # Multiple function calls + self.assertTrue( + await strategy.process_frame( + FunctionCallsStartedFrame( + function_calls=[ + FunctionCallFromLLM( + function_name="fn_3", tool_call_id="3", arguments={}, context=None + ), + FunctionCallFromLLM( + function_name="fn_4", tool_call_id="4", arguments={}, context=None + ), + ] + ) + ) + ) + self.assertTrue(await strategy.process_frame(InterruptionFrame())) + # First function call is done, we still should be muted since there's + # another one ongoing. + self.assertTrue( + await strategy.process_frame( + FunctionCallResultFrame( + function_name="fn_3", tool_call_id="3", arguments={}, result={} + ) + ) + ) + self.assertTrue(await strategy.process_frame(InterruptionFrame())) + # Last function call finishes. + self.assertFalse( + await strategy.process_frame( + FunctionCallResultFrame( + function_name="fn_4", tool_call_id="4", arguments={}, result={} + ) + ) + ) + self.assertFalse(await strategy.process_frame(InterruptionFrame()))