knowledge-base-qa/test_interaction.py at main · codewriterYang/knowledge-base-qa · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
"""
[FILE]: test_interaction.py
[ROLE]: 交互流程测试脚本
[DEPS]: unittest, unittest.mock | @interaction_manager.py, @kb_manager.py
[VIBE_LOGIC]:
- 核心逻辑：测试交互流程的各个环节
- 关键约束：覆盖所有输入验证和错误处理路径
"""

import unittest
from unittest.mock import patch, MagicMock, call
import sys
import io


class TestInteractionManager(unittest.TestCase):
    """交互管理器测试类"""

    def setUp(self):
        """测试前准备"""
        self.mock_kb = MagicMock()
        self.mock_kb.get_stats.return_value = {'doc_count': 5, 'log_count': 0}
        self.mock_kb.load_index.return_value = True
        self.mock_kb.clear_kb.return_value = None
        self.mock_kb.add_document.return_value = None
        self.mock_kb.ask_stream.return_value = ["测试答案"]

    @patch('builtins.input', side_effect=['n'])
    def test_post_demo_exit(self, mock_input):
        """测试演示结束后选择退出"""
        from modules.interaction_manager import InteractionManager

        manager = InteractionManager(self.mock_kb)

        result = manager.start_post_demo_flow()

        self.assertFalse(result)
        mock_input.assert_called_once_with("演示已结束，是否进入交互模式？(y/n): ")

    @patch('builtins.input', side_effect=['invalid', 'Y'])
    def test_post_demo_invalid_then_valid(self, mock_input):
        """测试无效输入后输入有效选项"""
        from modules.interaction_manager import InteractionManager

        manager = InteractionManager(self.mock_kb)
        manager._handle_index_choice = MagicMock(return_value=True)
        manager._chat_interaction = MagicMock()

        result = manager.start_post_demo_flow()

        self.assertTrue(result)
        self.assertEqual(mock_input.call_count, 2)

    @patch('builtins.input', side_effect=['y', 'a'])
    def test_index_choice_a(self, mock_input):
        """测试选择使用已有索引"""
        from modules.interaction_manager import InteractionManager

        manager = InteractionManager(self.mock_kb)
        manager._chat_interaction = MagicMock()

        result = manager.start_post_demo_flow()

        self.assertTrue(result)
        self.mock_kb.load_index.assert_called_once()

    @patch('builtins.input', side_effect=['y', 'b'])
    def test_index_choice_b(self, mock_input):
        """测试选择重新构建索引"""
        from modules.interaction_manager import InteractionManager

        manager = InteractionManager(self.mock_kb)
        manager._chat_interaction = MagicMock()

        with patch('data.simulated_docs.SIMULATED_DOCS', [
            {'id': 'p1', 'title': 'Test', 'content': 'Test content'}
        ]):
            result = manager.start_post_demo_flow()

        self.assertTrue(result)
        self.mock_kb.clear_kb.assert_called_once()
        self.mock_kb.add_document.assert_called_once()

    @patch('builtins.input', side_effect=['y', 'invalid', 'A'])
    def test_index_choice_invalid(self, mock_input):
        """测试索引选择时的无效输入"""
        from modules.interaction_manager import InteractionManager

        manager = InteractionManager(self.mock_kb)
        manager._chat_interaction = MagicMock()

        result = manager.start_post_demo_flow()

        self.assertTrue(result)
        self.assertEqual(mock_input.call_count, 3)

    def test_validate_input(self):
        """测试输入验证函数"""
        from modules.interaction_manager import InteractionManager

        manager = InteractionManager(self.mock_kb)

        with patch('builtins.input', side_effect=['invalid', 'y']):
            result = manager._validate_input(
                "prompt: ",
                ['y', 'n'],
                "error"
            )

        self.assertEqual(result, 'y')

    def test_progress_bar_output(self):
        """测试进度条输出"""
        from modules.interaction_manager import InteractionManager

        manager = InteractionManager(self.mock_kb)

        captured_output = io.StringIO()
        sys.stdout = captured_output

        manager._show_progress_bar(5, 10, prefix="Test:", suffix="(5/10)")

        sys.stdout = sys.__stdout__

        output = captured_output.getvalue()
        self.assertIn("[=====     ]", output)
        self.assertIn("50.0%", output)


if __name__ == '__main__':
    unittest.main(verbosity=2)