tests: add user mute strategy tests

2025-12-24 09:32:46 -08:00
parent 1d70275574
commit 5a4236bc71
1 changed files with 139 additions and 0 deletions
--- a/tests/test_user_mute_strategy.py
+++ 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()))