File size: 4,535 Bytes
7ab7df1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1521cd2
7ab7df1
 
 
 
 
 
 
 
 
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
"""Tests for ui/app.py module."""

from unittest.mock import MagicMock, Mock, patch

import pytest

from ui.app import create_app, initialize_chain


@patch("ui.app.create_qa_chain")
@patch("ui.app.load_or_create_vectorstore")
@patch("ui.app.create_embeddings")
def test_initialize_chain(mock_create_embeddings, mock_load_vectorstore, mock_create_qa_chain):
    """Test initialize_chain function."""
    # Setup mocks
    mock_embeddings = MagicMock()
    mock_create_embeddings.return_value = mock_embeddings

    mock_vectorstore = MagicMock()
    mock_vectorstore.get.return_value = {
        "metadatas": [
            {"source": "pdf/test1.pdf"},
            {"source": "pdf/test2.pdf"},
        ]
    }
    mock_load_vectorstore.return_value = mock_vectorstore

    mock_qa_chain = MagicMock()
    mock_create_qa_chain.return_value = mock_qa_chain

    # Test
    chain, sources = initialize_chain()

    assert chain == mock_qa_chain
    assert len(sources) == 2
    assert "test1.pdf" in sources
    assert "test2.pdf" in sources


@patch("ui.app.create_qa_chain")
@patch("ui.app.load_or_create_vectorstore")
@patch("ui.app.create_embeddings")
def test_initialize_chain_empty_metadatas(
    mock_create_embeddings, mock_load_vectorstore, mock_create_qa_chain
):
    """Test initialize_chain with empty metadatas."""
    # Setup mocks
    mock_embeddings = MagicMock()
    mock_create_embeddings.return_value = mock_embeddings

    mock_vectorstore = MagicMock()
    mock_vectorstore.get.return_value = {"metadatas": []}
    mock_load_vectorstore.return_value = mock_vectorstore

    mock_qa_chain = MagicMock()
    mock_create_qa_chain.return_value = mock_qa_chain

    # Test
    chain, sources = initialize_chain()

    assert chain == mock_qa_chain
    assert sources == []


@patch("ui.app.create_qa_chain")
@patch("ui.app.load_or_create_vectorstore")
@patch("ui.app.create_embeddings")
def test_initialize_chain_no_collection(
    mock_create_embeddings, mock_load_vectorstore, mock_create_qa_chain
):
    """Test initialize_chain with no collection."""
    # Setup mocks
    mock_embeddings = MagicMock()
    mock_create_embeddings.return_value = mock_embeddings

    mock_vectorstore = MagicMock()
    mock_vectorstore.get.return_value = None
    mock_load_vectorstore.return_value = mock_vectorstore

    mock_qa_chain = MagicMock()
    mock_create_qa_chain.return_value = mock_qa_chain

    # Test
    chain, sources = initialize_chain()

    assert chain == mock_qa_chain
    assert sources == []


@patch("ui.app.update_hybrid_alpha_visibility")
@patch("ui.app.update_rag_controls")
@patch("ui.app.create_respond_handler")
@patch("ui.app.create_stream_chat_response")
@patch("ui.app.create_ui_components")
@patch("ui.app.initialize_chain")
def test_create_app(
    mock_initialize,
    mock_create_components,
    mock_create_stream,
    mock_create_respond,
    mock_update_rag,
    mock_update_hybrid,
):
    """Test create_app function."""
    # Setup mocks
    mock_qa_chain = MagicMock()
    mock_initialize.return_value = (mock_qa_chain, ["test1.pdf", "test2.pdf"])

    mock_components = {
        "demo": MagicMock(),
        "msg": MagicMock(),
        "chatbot": MagicMock(),
        "rag_enabled": MagicMock(),
        "search_type": MagicMock(),
        "doc_filter": MagicMock(),
        "query_rewriting": MagicMock(),
        "reranking": MagicMock(),
        "hybrid_alpha": MagicMock(),
        "context_box": MagicMock(),
        "search_col": MagicMock(),
        "filter_col": MagicMock(),
        "context_section": MagicMock(),
        "advanced_options": MagicMock(),
        "submit": MagicMock(),
        "clear": MagicMock(),
    }
    mock_create_components.return_value = mock_components

    mock_stream_fn = MagicMock()
    mock_create_stream.return_value = mock_stream_fn

    mock_respond_fn = MagicMock()
    mock_create_respond.return_value = mock_respond_fn

    # Test
    app = create_app()

    assert app == mock_components["demo"]
    mock_initialize.assert_called_once()
    mock_create_components.assert_called_once()
    mock_create_stream.assert_called_once_with(mock_qa_chain, ["test1.pdf", "test2.pdf"])
    mock_create_respond.assert_called_once_with(mock_stream_fn)

    # Check event handlers were attached
    assert mock_components["msg"].submit.called
    assert mock_components["submit"].click.called
    assert mock_components["clear"].click.called
    assert mock_components["rag_enabled"].change.called
    assert mock_components["search_type"].change.called