From 1887f64270f8bab89bc02d4f2351efd4475b43cb Mon Sep 17 00:00:00 2001 From: Pavan Kumar Date: Fri, 29 Nov 2024 23:01:58 +0530 Subject: [PATCH 1/8] Unit Testcase Chathistory --- code/tests/test_chat_history.py | 605 ++++++++++++++++++++++++++++++++ 1 file changed, 605 insertions(+) create mode 100644 code/tests/test_chat_history.py diff --git a/code/tests/test_chat_history.py b/code/tests/test_chat_history.py new file mode 100644 index 000000000..4ac0cf567 --- /dev/null +++ b/code/tests/test_chat_history.py @@ -0,0 +1,605 @@ +""" +This module tests the entry point for the application. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from create_app import create_app + + +@pytest.fixture +def client(): + """Create a test client for the app.""" + app = create_app() + app.testing = True + return app.test_client() + + +@pytest.fixture +def mock_conversation_client(): + """Mock the database client.""" + with patch( + "backend.batch.utilities.chat_history.database_factory.DatabaseFactory.get_conversation_client" + ) as mock: + mock_conversation_client = AsyncMock() + mock.return_value = mock_conversation_client + yield mock_conversation_client + + +@pytest.fixture +def mock_config_helper(): + """Mock the ConfigHelper to control the config behavior.""" + with patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) as mock: + mock_config = MagicMock() + mock_config.enable_chat_history = ( + True # Ensure chat history is enabled for the test + ) + mock.return_value = mock_config + yield mock_config + + +class TestListConversations: + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_list_conversations_success( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test that the list_conversations endpoint works when everything is set up correctly.""" + # Given + get_active_config_or_default_mock.return_value.prompts.conversational_flow = ( + "custom" + ) + get_active_config_or_default_mock.enable_chat_history = True + mock_conversation_client.get_conversations = AsyncMock( + return_value=[{"conversation_id": "1", "content": "Hello, world!"}] + ) + + # When + response = client.get("/api/history/list?offset=0") + + # Then + assert response.status_code == 200 + assert response.json == [{"conversation_id": "1", "content": "Hello, world!"}] + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_list_conversations_no_history( + self, get_active_config_or_default_mock, client + ): + """Test that the list_conversations endpoint returns an error if chat history is not enabled.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = False + + # When + response = client.get("/api/history/list?offset=0") + + # Then + assert response.status_code == 400 + assert response.json == {"error": "Chat history is not available"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_list_conversations_db_error( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test that the list_conversations endpoint returns an error if the database is not available.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversations = AsyncMock( + side_effect=Exception("Database error") + ) + + # When + response = client.get("/api/history/list?offset=0") + + # Then + assert response.status_code == 500 + assert response.json == { + "error": "Error while listing historical conversations" + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_list_conversations_no_conversations( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test that the list_conversations endpoint returns an error if no conversations are found.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversations = AsyncMock( + return_value="invalid response" + ) + + # When + response = client.get("/api/history/list?offset=0") + + # Then + assert response.status_code == 404 + assert response.json == { + "error": "No conversations for 00000000-0000-0000-0000-000000000000 were found" + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_rename_conversation_success( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test that the rename_conversation endpoint works correctly.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversations = AsyncMock( + return_value={"conversation_id": "1", "title": "Old Title"} + ) + mock_conversation_client.upsert_conversation = AsyncMock( + return_value={"conversation_id": "1", "title": "New Title"} + ) + + request_json = {"conversation_id": "1", "title": "New Title"} + + # When + response = client.post("/api/history/rename", json=request_json) + + # Then + assert response.status_code == 200 + assert response.json == {"conversation_id": "1", "title": "New Title"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_rename_conversation_no_history( + self, get_active_config_or_default_mock, client + ): + """Test that the rename_conversation endpoint returns an error if chat history is not enabled.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = False + + request_json = {"conversation_id": "1", "title": "New Title"} + + # When + response = client.post("/api/history/rename", json=request_json) + + # Then + assert response.status_code == 400 + assert response.json == {"error": "Chat history is not available"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_rename_conversation_missing_conversation_id( + self, get_active_config_or_default_mock, client + ): + """Test that the rename_conversation endpoint returns an error if conversation_id is missing.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + + request_json = {"title": "New Title"} + + # When + response = client.post("/api/history/rename", json=request_json) + + # Then + assert response.status_code == 400 + assert response.json == {"error": "conversation_id is required"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_rename_conversation_empty_title( + self, get_active_config_or_default_mock, client + ): + """Test that the rename_conversation endpoint returns an error if the title is empty.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + + request_json = {"conversation_id": "1", "title": ""} + + # When + response = client.post("/api/history/rename", json=request_json) + + # Then + assert response.status_code == 400 + assert response.json == {"error": "A non-empty title is required"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + @patch( + "backend.batch.utilities.chat_history.database_factory.DatabaseFactory.get_conversation_client" + ) + def test_rename_conversation_db_error( + self, mock_conversation_client, get_active_config_or_default_mock, client + ): + """Test that the rename_conversation endpoint returns an error if the database is not available.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.return_value.get_conversation = AsyncMock( + side_effect=Exception("Database error") + ) + + request_json = {"conversation_id": "1", "title": "New Title"} + + # When + response = client.post("/api/history/rename", json=request_json) + + # Then + assert response.status_code == 500 + assert response.json == {"error": "Error while renaming conversation"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_rename_conversation_not_found( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test that the rename_conversation endpoint returns an error if the conversation is not found.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversation = AsyncMock(return_value=None) + + request_json = {"conversation_id": "1", "title": "New Title"} + + # When + response = client.post("/api/history/rename", json=request_json) + + # Then + assert response.status_code == 400 + assert response.json == { + "error": "Conversation 1 was not found. It either does not exist or the logged in user does not have access to it." + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_get_conversation_success( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test that the get_conversation endpoint works correctly.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversation = AsyncMock( + return_value={"conversation_id": "1", "title": "Sample Conversation"} + ) + mock_conversation_client.get_messages = AsyncMock( + return_value=[ + { + "id": "1", + "role": "user", + "content": "Hello, world!", + "createdAt": "2024-11-29T12:00:00Z", + } + ] + ) + + request_json = {"conversation_id": "1"} + + # When + response = client.post("/api/history/read", json=request_json) + + # Then + assert response.status_code == 200 + assert response.json == { + "conversation_id": "1", + "messages": [ + { + "id": "1", + "role": "user", + "content": "Hello, world!", + "createdAt": "2024-11-29T12:00:00Z", + "feedback": None, + } + ], + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_get_conversation_no_history( + self, get_active_config_or_default_mock, client + ): + """Test that the get_conversation endpoint returns an error if chat history is not enabled.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = False + + request_json = {"conversation_id": "1"} + + # When + response = client.post("/api/history/read", json=request_json) + + # Then + assert response.status_code == 400 + assert response.json == {"error": "Chat history is not available"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_get_conversation_missing_conversation_id( + self, get_active_config_or_default_mock, client + ): + """Test that the get_conversation endpoint returns an error if conversation_id is missing.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + + request_json = {} + + # When + response = client.post("/api/history/read", json=request_json) + + # Then + assert response.status_code == 400 + assert response.json == {"error": "conversation_id is required"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + @patch( + "backend.batch.utilities.chat_history.database_factory.DatabaseFactory.get_conversation_client" + ) + def test_get_conversation_db_error( + self, mock_conversation_client, get_active_config_or_default_mock, client + ): + """Test that the get_conversation endpoint returns an error if the database is not available.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.return_value.get_conversation = AsyncMock( + side_effect=Exception("Database error") + ) + + request_json = {"conversation_id": "1"} + + # When + response = client.post("/api/history/read", json=request_json) + + # Then + assert response.status_code == 500 + assert response.json == {"error": "Error while fetching conversation history"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_get_conversation_not_found( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test that the get_conversation endpoint returns an error if the conversation is not found.""" + # Given + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversation = AsyncMock(return_value=None) + + request_json = {"conversation_id": "1"} + + # When + response = client.post("/api/history/read", json=request_json) + + # Then + assert response.status_code == 400 + assert response.json == { + "error": "Conversation 1 was not found. It either does not exist or the logged in user does not have access to it." + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_delete_conversation_success( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test that the delete_conversation endpoint works correctly.""" + + # Setup mocks + get_active_config_or_default_mock.return_value.enable_chat_history = True + + # Mock the database client + mock_conversation_client.delete_messages = AsyncMock(return_value=None) + mock_conversation_client.delete_conversation = AsyncMock(return_value=None) + + # Define request data + request_json = {"conversation_id": "conv123"} + + # Make DELETE request to delete the conversation + response = client.delete("/api/history/delete", json=request_json) + + # Assert the response status and data + assert response.status_code == 200 + assert response.json == { + "message": "Successfully deleted conversation and messages", + "conversation_id": "conv123", + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_delete_conversation_no_chat_history( + self, get_active_config_or_default_mock, client + ): + """Test when chat history is not enabled in the configuration.""" + + # Setup mocks + get_active_config_or_default_mock.return_value.enable_chat_history = False + + # Define request data + request_json = {"conversation_id": "conv123"} + + # Make DELETE request to delete the conversation + response = client.delete("/api/history/delete", json=request_json) + + # Assert the response status and error message + assert response.status_code == 400 + assert response.json == {"error": "Chat history is not available"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_delete_conversation_missing_conversation_id( + self, get_active_config_or_default_mock, client + ): + """Test when the conversation_id is missing in the request.""" + + # Setup mocks + get_active_config_or_default_mock.return_value.enable_chat_history = True + + # Define request data (missing conversation_id) + request_json = {} + + # Make DELETE request to delete the conversation + response = client.delete("/api/history/delete", json=request_json) + + # Assert the response status and error message + assert response.status_code == 400 + assert response.json == { + "error": "Conversation None was not found. It either does not exist or the logged in user does not have access to it." + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_delete_conversation_database_error( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test when the database client connection fails.""" + + # Setup mocks + get_active_config_or_default_mock.return_value.enable_chat_history = True + + # Mock a failure in the database client connection + mock_conversation_client.connect.side_effect = Exception( + "Database not available" + ) + + # Define request data + request_json = {"conversation_id": "conv123"} + + # Make DELETE request to delete the conversation + response = client.delete("/api/history/delete", json=request_json) + + # Assert the response status and error message + assert response.status_code == 500 + assert response.json == {"error": "Error while deleting conversation history"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_delete_conversation_internal_error( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + """Test when an unexpected internal error occurs during conversation deletion.""" + + # Setup mocks + get_active_config_or_default_mock.return_value.enable_chat_history = True + + # Mock an unexpected error in the database client deletion + mock_conversation_client.delete_messages.side_effect = Exception( + "Unexpected error" + ) + + # Define request data + request_json = {"conversation_id": "conv123"} + + # Make DELETE request to delete the conversation + response = client.delete("/api/history/delete", json=request_json) + + # Assert the response status and error message + assert response.status_code == 500 + assert response.json == {"error": "Error while deleting conversation history"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_delete_all_conversations_success( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversation = AsyncMock( + return_value=[{"id": "conv1"}, {"id": "conv2"}] + ) + + response = client.delete("/api/history/delete_all") + assert response.status_code == 200 + assert response.json == { + "message": "Successfully deleted all conversations and messages for user 00000000-0000-0000-0000-000000000000" + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_delete_all_conversations_no_chat_history( + self, get_active_config_or_default_mock, client + ): + get_active_config_or_default_mock.return_value.enable_chat_history = False + response = client.delete("/api/history/delete_all") + assert response.status_code == 400 + assert response.json == {"error": "Chat history is not available"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_update_conversation_success( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversation.return_value = { + "title": "Test Title", + "updatedAt": "2024-12-01", + "id": "conv1", + } + mock_conversation_client.create_message.return_value = "success" + request_json = { + "conversation_id": "conv1", + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + ], + } + + # When + response = client.post("/api/history/update", json=request_json) + + assert response.status_code == 200 + assert response.json == { + "data": { + "conversation_id": "conv1", + "date": "2024-12-01", + "title": "Test Title", + }, + "success": True, + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_update_conversation_no_chat_history( + self, get_active_config_or_default_mock, client + ): + get_active_config_or_default_mock.return_value.enable_chat_history = False + response = client.post( + "/api/history/update", json={}, headers={"Content-Type": "application/json"} + ) + assert response.status_code == 400 + assert response.json == {"error": "Chat history is not available"} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_get_frontend_settings_success( + self, get_active_config_or_default_mock, client + ): + get_active_config_or_default_mock.return_value.enable_chat_history = True + response = client.get("/api/history/frontend_settings") + assert response.status_code == 200 + assert response.json == {"CHAT_HISTORY_ENABLED": True} + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_get_frontend_settings_error( + self, get_active_config_or_default_mock, client + ): + get_active_config_or_default_mock.side_effect = Exception("Test Error") + response = client.get("/api/history/frontend_settings") + assert response.status_code == 500 + assert response.json == {"error": "Error while getting frontend settings"} From 335476256836e0ffc639c0e87f7262150b718bbf Mon Sep 17 00:00:00 2001 From: Pavan Kumar Date: Mon, 2 Dec 2024 20:17:30 +0530 Subject: [PATCH 2/8] Unit Testing for Chat History Functionality - Utilities --- .../test_postgres_search_handler.py | 169 ++++ code/tests/test_chat_history.py | 16 +- .../helpers/test_azure_postgres_helper.py | 741 ++++++++++++++++++ .../helpers/test_postgress_embedder.py | 211 +++++ 4 files changed, 1122 insertions(+), 15 deletions(-) create mode 100644 code/tests/search_utilities/test_postgres_search_handler.py create mode 100644 code/tests/utilities/helpers/test_azure_postgres_helper.py create mode 100644 code/tests/utilities/helpers/test_postgress_embedder.py diff --git a/code/tests/search_utilities/test_postgres_search_handler.py b/code/tests/search_utilities/test_postgres_search_handler.py new file mode 100644 index 000000000..eead10dd3 --- /dev/null +++ b/code/tests/search_utilities/test_postgres_search_handler.py @@ -0,0 +1,169 @@ +import pytest +from unittest.mock import MagicMock, patch +from backend.batch.utilities.common.source_document import SourceDocument +from backend.batch.utilities.search.postgres_search_handler import AzurePostgresHandler + + +@pytest.fixture(autouse=True) +def env_helper_mock(): + mock = MagicMock() + mock.POSTGRESQL_USER = "test_user" + mock.POSTGRESQL_PASSWORD = "test_password" + mock.POSTGRESQL_HOST = "test_host" + mock.POSTGRESQL_DB = "test_db" + return mock + + +@pytest.fixture(autouse=True) +def mock_search_client(): + with patch( + "backend.batch.utilities.search.postgres_search_handler.AzurePostgresHelper" + ) as mock: + search_client = mock.return_value.get_search_client.return_value + yield search_client + + +@pytest.fixture +def handler(env_helper_mock, mock_search_client): + with patch( + "backend.batch.utilities.search.postgres_search_handler", + return_value=mock_search_client, + ): + return AzurePostgresHandler(env_helper_mock) + + +def test_query_search(handler, mock_search_client): + mock_llm_helper = MagicMock() + mock_search_client.llm_helper = mock_llm_helper + + mock_llm_helper.generate_embeddings.return_value = [1, 2, 3] + + mock_search_client.get_search_indexes.return_value = [ + { + "id": "1", + "title": "Title1", + "chunk": "Chunk1", + "offset": 0, + "page_number": 1, + "content": "Content1", + "source": "Source1", + }, + { + "id": "2", + "title": "Title2", + "chunk": "Chunk2", + "offset": 1, + "page_number": 2, + "content": "Content2", + "source": "Source2", + }, + ] + + mock_search_client.get_search_client.return_value = mock_search_client + handler.azure_postgres_helper = mock_search_client + + result = handler.query_search("Sample question") + + mock_llm_helper.generate_embeddings.assert_called_once_with("Sample question") + mock_search_client.get_search_indexes.assert_called_once() + assert len(result) == 2 + assert isinstance(result[0], SourceDocument) + assert result[0].id == "1" + assert result[0].title == "Title1" + assert result[1].content == "Content2" + + +def test_convert_to_source_documents(handler): + search_results = [ + { + "id": "1", + "title": "Title1", + "chunk": "Chunk1", + "offset": 0, + "page_number": 1, + "content": "Content1", + "source": "Source1", + }, + { + "id": "2", + "title": "Title2", + "chunk": "Chunk2", + "offset": 1, + "page_number": 2, + "content": "Content2", + "source": "Source2", + }, + ] + + result = handler._convert_to_source_documents(search_results) + + assert len(result) == 2 + assert result[0].id == "1" + assert result[0].content == "Content1" + assert result[1].page_number == 2 + + +def test_create_search_client(handler, mock_search_client): + handler.azure_postgres_helper.get_search_client = MagicMock( + return_value=mock_search_client + ) + + result = handler.create_search_client() + + assert result == mock_search_client + + +def test_get_files(handler): + mock_get_files = MagicMock(return_value=["test1.txt", "test2.txt"]) + handler.azure_postgres_helper.get_files = mock_get_files + + result = handler.get_files() + + assert len(result) == 2 + assert result[0] == "test1.txt" + assert result[1] == "test2.txt" + + +def test_delete_files(handler): + files_to_delete = {"test1.txt": [1, 2], "test2.txt": [3]} + mock_delete_documents = MagicMock() + handler.azure_postgres_helper.delete_documents = mock_delete_documents + + result = handler.delete_files(files_to_delete) + + mock_delete_documents.assert_called_once_with([{"id": 1}, {"id": 2}, {"id": 3}]) + assert "test1.txt" in result + + +# Test case for delete_from_index method +def test_delete_from_index(handler): + blob_url = "https://example.com/blob" + + # Mocking methods + mock_search_by_blob_url = MagicMock(return_value=[{"id": "1", "title": "Title1"}]) + mock_output_results = MagicMock(return_value={"test1.txt": ["1"]}) + mock_delete_files = MagicMock(return_value="test1.txt") + + handler.search_by_blob_url = mock_search_by_blob_url + handler.output_results = mock_output_results + handler.delete_files = mock_delete_files + + handler.delete_from_index(blob_url) + + mock_search_by_blob_url.assert_called_once_with(blob_url) + mock_output_results.assert_called_once() + mock_delete_files.assert_called_once_with({"test1.txt": ["1"]}) + + +# Test case for get_unique_files method +def test_get_unique_files(handler): + mock_get_unique_files = MagicMock( + return_value=[{"title": "test1.txt"}, {"title": "test2.txt"}] + ) + handler.azure_postgres_helper.get_unique_files = mock_get_unique_files + + result = handler.get_unique_files() + + assert len(result) == 2 + assert result[0] == "test1.txt" + assert result[1] == "test2.txt" diff --git a/code/tests/test_chat_history.py b/code/tests/test_chat_history.py index 4ac0cf567..f1b8bdcb1 100644 --- a/code/tests/test_chat_history.py +++ b/code/tests/test_chat_history.py @@ -2,7 +2,7 @@ This module tests the entry point for the application. """ -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest from create_app import create_app @@ -27,20 +27,6 @@ def mock_conversation_client(): yield mock_conversation_client -@pytest.fixture -def mock_config_helper(): - """Mock the ConfigHelper to control the config behavior.""" - with patch( - "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" - ) as mock: - mock_config = MagicMock() - mock_config.enable_chat_history = ( - True # Ensure chat history is enabled for the test - ) - mock.return_value = mock_config - yield mock_config - - class TestListConversations: @patch( "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" diff --git a/code/tests/utilities/helpers/test_azure_postgres_helper.py b/code/tests/utilities/helpers/test_azure_postgres_helper.py new file mode 100644 index 000000000..907d0bdfd --- /dev/null +++ b/code/tests/utilities/helpers/test_azure_postgres_helper.py @@ -0,0 +1,741 @@ +import unittest +from unittest.mock import MagicMock, patch +import psycopg2 +from backend.batch.utilities.helpers.azure_postgres_helper import AzurePostgresHelper + + +class TestAzurePostgresHelper(unittest.TestCase): + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + def test_create_search_client_success(self, mock_connect, mock_credential): + # Arrange + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + + mock_connection = MagicMock() + mock_connect.return_value = mock_connection + + helper = AzurePostgresHelper() + helper.env_helper.POSTGRESQL_USER = "mock_user" + helper.env_helper.POSTGRESQL_HOST = "mock_host" + helper.env_helper.POSTGRESQL_DATABASE = "mock_database" + + # Act + connection = helper._create_search_client() + + # Assert + self.assertEqual(connection, mock_connection) + mock_credential.return_value.get_token.assert_called_once_with( + "https://ossrdbms-aad.database.windows.net/.default" + ) + mock_connect.assert_called_once_with( + "host=mock_host user=mock_user dbname=mock_database password=mock-access-token" + ) + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + def test_get_search_client_reuses_connection(self, mock_connect): + # Arrange + mock_connection = MagicMock() + mock_connection.closed = 0 # Simulate an open connection + mock_connect.return_value = mock_connection + + helper = AzurePostgresHelper() + helper.conn = mock_connection + + # Act + connection = helper.get_search_client() + + # Assert + self.assertEqual(connection, mock_connection) + mock_connect.assert_not_called() # Ensure no new connection is created + + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.RealDictCursor") + def test_get_search_indexes_success( + self, mock_cursor, mock_connect, mock_credential + ): + # Arrange + # Mock the EnvHelper and set required attributes + mock_env_helper = MagicMock() + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + + # Mock the database connection and cursor + mock_connection = MagicMock() + mock_connect.return_value = mock_connection + mock_cursor_instance = MagicMock() + mock_cursor.return_value = mock_cursor_instance + + # Mock the behavior of the context manager for the cursor + mock_cursor_context = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor_context + mock_results = [{"id": 1, "title": "Test"}] + mock_cursor_context.fetchall.return_value = mock_results + + # Replace EnvHelper in AzurePostgresHelper with the mocked version + helper = AzurePostgresHelper() + helper.env_helper = mock_env_helper + + # Embedding vector for the test + embedding_vector = [1, 2, 3] + + # Act + results = helper.get_search_indexes(embedding_vector) + + # Assert + self.assertEqual(results, mock_results) + mock_connect.assert_called_once_with( + "host=mock_host user=mock_user dbname=mock_database password=mock-access-token" + ) + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + def test_get_search_indexes_query_error(self, mock_connect): + # Arrange + # Mock the EnvHelper and set required attributes + mock_env_helper = MagicMock() + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + mock_connection = MagicMock() + mock_connect.return_value = mock_connection + + def raise_exception(*args, **kwargs): + raise Exception("Query execution error") + + mock_cursor_instance = MagicMock() + mock_cursor_instance.execute.side_effect = raise_exception + + mock_connection.cursor.return_value.__enter__.return_value = ( + mock_cursor_instance + ) + + helper = AzurePostgresHelper() + helper.env_helper = mock_env_helper + embedding_vector = [1, 2, 3] + + # Act & Assert + with self.assertRaises(Exception) as context: + helper.get_search_indexes(embedding_vector) + + self.assertEqual(str(context.exception), "Query execution error") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + def test_create_search_client_connection_error(self, mock_connect): + # Arrange + # Mock the EnvHelper and set required attributes + mock_env_helper = MagicMock() + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + def raise_exception(*args, **kwargs): + raise Exception("Connection error") + + mock_connect.side_effect = raise_exception + + helper = AzurePostgresHelper() + helper.env_helper = mock_env_helper + + # Act & Assert + with self.assertRaises(Exception) as context: + helper._create_search_client() + + self.assertEqual(str(context.exception), "Connection error") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_get_files_success(self, mock_env_helper, mock_connect): + # Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Arrange: Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the result of the cursor's fetchall() method + mock_cursor.fetchall.return_value = [ + {"id": 1, "title": "Title 1"}, + {"id": 2, "title": "Title 2"}, + ] + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act: Call the method under test + result = helper.get_files() + + # Assert: Check that the correct results are returned + self.assertEqual( + result, [{"id": 1, "title": "Title 1"}, {"id": 2, "title": "Title 2"}] + ) + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_get_files_no_results(self, mock_env_helper, mock_connect): + # Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Arrange: Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the result of the cursor's fetchall() method to return an empty list + mock_cursor.fetchall.return_value = [] + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act: Call the method under test + result = helper.get_files() + + # Assert: Check that the result is None + self.assertIsNone(result) + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + def test_get_files_db_error(self, mock_logger, mock_env_helper, mock_connect): + # Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Arrange: Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Simulate a database error when executing the query + mock_cursor.fetchall.side_effect = psycopg2.Error("Database error") + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act & Assert: Ensure that the exception is raised and the error is logged + with self.assertRaises(psycopg2.Error): + helper.get_files() + + mock_logger.error.assert_called_with( + "Database error while fetching titles: Database error" + ) + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + def test_get_files_unexpected_error( + self, mock_logger, mock_env_helper, mock_connect + ): + # Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Arrange: Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Simulate an unexpected error + mock_cursor.fetchall.side_effect = Exception("Unexpected error") + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act & Assert: Ensure that the exception is raised and the error is logged + with self.assertRaises(Exception): + helper.get_files() + + mock_logger.error.assert_called_with( + "Unexpected error while fetching titles: Unexpected error" + ) + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_delete_documents_success(self, mock_env_helper, mock_logger, mock_connect): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the behavior of cursor.rowcount and execute + mock_cursor.rowcount = 3 # Simulate 3 rows deleted + mock_cursor.execute.return_value = None + + ids_to_delete = [{"id": 1}, {"id": 2}, {"id": 3}] + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act: Call the method under test + result = helper.delete_documents(ids_to_delete) + + # Assert: Check that the correct number of rows were deleted + self.assertEqual(result, 3) + mock_connection.commit.assert_called_once() + mock_connection.close.assert_called_once() + mock_logger.info.assert_called_with("Deleted 3 documents.") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_delete_documents_no_ids(self, mock_env_helper, mock_logger, mock_connect): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # No IDs to delete + ids_to_delete = [] + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act: Call the method under test + result = helper.delete_documents(ids_to_delete) + + # Assert: Check that no rows were deleted and a warning was logged + self.assertEqual(result, 0) + mock_logger.warning.assert_called_with("No IDs provided for deletion.") + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_delete_documents_db_error( + self, mock_env_helper, mock_logger, mock_connect + ): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Simulate a database error during execution + mock_cursor.execute.side_effect = psycopg2.Error("Database error") + + ids_to_delete = [{"id": 1}, {"id": 2}] + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act & Assert: Ensure that the exception is raised and the error is logged + with self.assertRaises(psycopg2.Error): + helper.delete_documents(ids_to_delete) + + mock_logger.error.assert_called_with( + "Database error while deleting documents: Database error" + ) + mock_connection.rollback.assert_called_once() + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_delete_documents_unexpected_error( + self, mock_env_helper, mock_logger, mock_connect + ): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Simulate an unexpected error + mock_cursor.execute.side_effect = Exception("Unexpected error") + + ids_to_delete = [{"id": 1}, {"id": 2}] + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act & Assert: Ensure that the exception is raised and the error is logged + with self.assertRaises(Exception): + helper.delete_documents(ids_to_delete) + + mock_logger.error.assert_called_with( + "Unexpected error while deleting documents: Unexpected error" + ) + mock_connection.rollback.assert_called_once() + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_perform_search_success(self, mock_env_helper, mock_logger, mock_connect): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the behavior of cursor's execute and fetchall + mock_cursor.fetchall.return_value = [ + { + "title": "Test Title", + "content": "Test Content", + "metadata": "Test Metadata", + } + ] + + title_to_search = "Test Title" + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act: Call the method under test + result = helper.perform_search(title_to_search) + + # Assert: Check that the results match the expected data + self.assertEqual(len(result), 1) # One result returned + self.assertEqual(result[0]["title"], "Test Title") + self.assertEqual(result[0]["content"], "Test Content") + self.assertEqual(result[0]["metadata"], "Test Metadata") + + # Ensure the connection was closed + mock_connection.close.assert_called_once() + mock_logger.info.assert_called_with("Retrieved 1 search result(s).") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_perform_search_no_results( + self, mock_env_helper, mock_logger, mock_connect + ): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the behavior of cursor's execute and fetchall to return no results + mock_cursor.fetchall.return_value = [] + + title_to_search = "Nonexistent Title" + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act: Call the method under test + result = helper.perform_search(title_to_search) + + # Assert: Check that no results were returned + self.assertEqual(result, []) # Empty list returned for no results + + # Ensure the connection was closed + mock_connection.close.assert_called_once() + mock_logger.info.assert_called_with("Retrieved 0 search result(s).") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_perform_search_error(self, mock_env_helper, mock_logger, mock_connect): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Simulate an error during the execution of the query + mock_cursor.execute.side_effect = Exception("Database error") + + title_to_search = "Test Title" + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act & Assert: Ensure that the exception is raised and the error is logged + with self.assertRaises(Exception): + helper.perform_search(title_to_search) + + mock_logger.error.assert_called_with( + "Error executing search query: Database error" + ) + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_get_unique_files_success(self, mock_env_helper, mock_logger, mock_connect): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the behavior of cursor's execute and fetchall + mock_cursor.fetchall.return_value = [ + {"title": "Unique Title 1"}, + {"title": "Unique Title 2"}, + ] + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act: Call the method under test + result = helper.get_unique_files() + + # Assert: Check that the results match the expected data + self.assertEqual(len(result), 2) # Two unique titles returned + self.assertEqual(result[0]["title"], "Unique Title 1") + self.assertEqual(result[1]["title"], "Unique Title 2") + + # Ensure the connection was closed + mock_connection.close.assert_called_once() + mock_logger.info.assert_called_with("Retrieved 2 unique title(s).") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_get_unique_files_no_results( + self, mock_env_helper, mock_logger, mock_connect + ): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the behavior of cursor's execute and fetchall to return no results + mock_cursor.fetchall.return_value = [] + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act: Call the method under test + result = helper.get_unique_files() + + # Assert: Check that no results were returned + self.assertEqual(result, []) # Empty list returned for no results + + # Ensure the connection was closed + mock_connection.close.assert_called_once() + mock_logger.info.assert_called_with("Retrieved 0 unique title(s).") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_get_unique_files_error(self, mock_env_helper, mock_logger, mock_connect): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Simulate an error during the execution of the query + mock_cursor.execute.side_effect = Exception("Database error") + + # Create an instance of the helper + helper = AzurePostgresHelper() + + # Act & Assert: Ensure that the exception is raised and the error is logged + with self.assertRaises(Exception): + helper.get_unique_files() + + mock_logger.error.assert_called_with( + "Error executing search query: Database error" + ) + mock_connection.close.assert_called_once() + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_search_by_blob_url_success( + self, mock_env_helper, mock_logger, mock_connect + ): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the behavior of cursor's execute and fetchall + mock_cursor.fetchall.return_value = [ + {"id": 1, "title": "Title 1"}, + {"id": 2, "title": "Title 2"}, + ] + + # Create an instance of the helper + helper = AzurePostgresHelper() + blob_url = "mock_blob_url" + + # Act: Call the method under test + result = helper.search_by_blob_url(blob_url) + + # Assert: Check that the results match the expected data + self.assertEqual(len(result), 2) # Two titles returned + self.assertEqual(result[0]["title"], "Title 1") + self.assertEqual(result[1]["title"], "Title 2") + + # Ensure the connection was closed + mock_connection.close.assert_called_once() + mock_logger.info.assert_called_with("Retrieved 2 unique title(s).") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_search_by_blob_url_no_results( + self, mock_env_helper, mock_logger, mock_connect + ): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Mock the behavior of cursor's execute and fetchall to return no results + mock_cursor.fetchall.return_value = [] + + # Create an instance of the helper + helper = AzurePostgresHelper() + blob_url = "mock_blob_url" + + # Act: Call the method under test + result = helper.search_by_blob_url(blob_url) + + # Assert: Check that no results were returned + self.assertEqual(result, []) # Empty list returned for no results + + # Ensure the connection was closed + mock_connection.close.assert_called_once() + mock_logger.info.assert_called_with("Retrieved 0 unique title(s).") + + @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") + @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") + def test_search_by_blob_url_error(self, mock_env_helper, mock_logger, mock_connect): + # Arrange: Mock the EnvHelper attributes + mock_env_helper.POSTGRESQL_USER = "mock_user" + mock_env_helper.POSTGRESQL_HOST = "mock_host" + mock_env_helper.POSTGRESQL_DATABASE = "mock_database" + mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + + # Mock the connection and cursor + mock_connection = MagicMock() + mock_cursor = MagicMock() + mock_connection.cursor.return_value.__enter__.return_value = mock_cursor + mock_connect.return_value = mock_connection + + # Simulate an error during the execution of the query + mock_cursor.execute.side_effect = Exception("Database error") + + # Create an instance of the helper + helper = AzurePostgresHelper() + blob_url = "mock_blob_url" + + # Act & Assert: Ensure that the exception is raised and the error is logged + with self.assertRaises(Exception): + helper.search_by_blob_url(blob_url) + + mock_logger.error.assert_called_with( + "Error executing search query: Database error" + ) + mock_connection.close.assert_called_once() diff --git a/code/tests/utilities/helpers/test_postgress_embedder.py b/code/tests/utilities/helpers/test_postgress_embedder.py new file mode 100644 index 000000000..1f76bae3d --- /dev/null +++ b/code/tests/utilities/helpers/test_postgress_embedder.py @@ -0,0 +1,211 @@ +from unittest.mock import MagicMock, patch, call + +import pytest +from backend.batch.utilities.helpers.embedders.postgres_embedder import PostgresEmbedder +from backend.batch.utilities.common.source_document import SourceDocument +from backend.batch.utilities.helpers.config.embedding_config import EmbeddingConfig +from backend.batch.utilities.document_loading.strategies import LoadingStrategy +from backend.batch.utilities.document_loading import LoadingSettings +from backend.batch.utilities.document_chunking.chunking_strategy import ChunkingSettings + +CHUNKING_SETTINGS = ChunkingSettings({"strategy": "layout", "size": 1, "overlap": 0}) +LOADING_SETTINGS = LoadingSettings({"strategy": LoadingStrategy.LAYOUT}) + + +@pytest.fixture(autouse=True) +def llm_helper_mock(): + with patch( + "backend.batch.utilities.helpers.embedders.postgres_embedder.LLMHelper" + ) as mock: + llm_helper = mock.return_value + llm_helper.get_embedding_model.return_value.embed_query.return_value = [ + 0 + ] * 1536 + mock_completion = llm_helper.get_chat_completion.return_value + choice = MagicMock() + choice.message.content = "This is a caption for an image" + mock_completion.choices = [choice] + llm_helper.generate_embeddings.return_value = [123] + yield llm_helper + + +@pytest.fixture(autouse=True) +def env_helper_mock(): + with patch( + "backend.batch.utilities.helpers.embedders.push_embedder.EnvHelper" + ) as mock: + env_helper = mock.return_value + yield env_helper + + +@pytest.fixture(autouse=True) +def azure_postgres_helper_mock(): + with patch( + "backend.batch.utilities.helpers.embedders.postgres_embedder.AzurePostgresHelper" + ) as mock: + yield mock + + +@pytest.fixture(autouse=True) +def mock_config_helper(): + with patch( + "backend.batch.utilities.helpers.embedders.postgres_embedder.ConfigHelper" + ) as mock: + config_helper = mock.get_active_config_or_default.return_value + config_helper.document_processors = [ + EmbeddingConfig( + "jpg", + CHUNKING_SETTINGS, + LOADING_SETTINGS, + use_advanced_image_processing=True, + ), + EmbeddingConfig( + "pdf", + CHUNKING_SETTINGS, + LOADING_SETTINGS, + use_advanced_image_processing=False, + ), + ] + config_helper.get_advanced_image_processing_image_types.return_value = { + "jpeg", + "jpg", + "png", + } + yield config_helper + + +@pytest.fixture(autouse=True) +def document_loading_mock(): + with patch( + "backend.batch.utilities.helpers.embedders.postgres_embedder.DocumentLoading" + ) as mock: + expected_documents = [ + SourceDocument(content="some content", source="some source") + ] + mock.return_value.load.return_value = expected_documents + yield mock + + +@pytest.fixture(autouse=True) +def document_chunking_mock(): + with patch( + "backend.batch.utilities.helpers.embedders.postgres_embedder.DocumentChunking" + ) as mock: + expected_chunked_documents = [ + SourceDocument( + content="some content", + source="some source", + id="some id", + title="some-title", + offset=1, + chunk=1, + page_number=1, + chunk_id="some chunk id", + ), + SourceDocument( + content="some other content", + source="some other source", + id="some other id", + title="some other-title", + offset=2, + chunk=2, + page_number=2, + chunk_id="some other chunk id", + ), + ] + mock.return_value.chunk.return_value = expected_chunked_documents + yield mock + + +def test_embed_file( + document_chunking_mock, + document_loading_mock, + llm_helper_mock, + azure_postgres_helper_mock, +): + postgres_embedder = PostgresEmbedder(MagicMock(), MagicMock()) + # Setup test data + source_url = "https://example.com/document.pdf" + file_name = "document.pdf" + file_extension = "pdf" + embedding_config = MagicMock() + postgres_embedder.embedding_configs[file_extension] = ( + embedding_config # This needs to be adapted if `self.embedder` isn't set. + ) + + # Mock methods + llm_helper_mock.generate_embeddings.return_value = [0.1, 0.2, 0.3] + azure_postgres_helper_mock.create_search_indexes.return_value = True + + # Execute + postgres_embedder.embed_file(source_url, file_name) + + # Assert method calls + document_loading_mock.return_value.load.assert_called_once_with( + source_url, embedding_config.loading + ) + document_chunking_mock.return_value.chunk.assert_called_once_with( + document_loading_mock.return_value.load.return_value, embedding_config.chunking + ) + llm_helper_mock.generate_embeddings.assert_has_calls( + [call("some content"), call("some other content")] + ) + + +def test_advanced_image_processing_not_implemented(): + postgres_embedder = PostgresEmbedder(MagicMock(), MagicMock()) + # Test for unsupported advanced image processing + file_extension = "jpg" + embedding_config = MagicMock() + embedding_config.use_advanced_image_processing = True + postgres_embedder.embedding_configs[file_extension] = embedding_config + + # Mock config method + postgres_embedder.config.get_advanced_image_processing_image_types = MagicMock( + return_value=["jpg", "png"] + ) + + # Use pytest.raises to check the exception + with pytest.raises(NotImplementedError) as context: + postgres_embedder.embed_file("https://example.com/image.jpg", "image.jpg") + + # Assert that the exception message matches the expected one + assert ( + str(context.value) + == "Advanced image processing is not supported in PostgresEmbedder." + ) + + +def test_postgres_embed_file_loads_documents(document_loading_mock, env_helper_mock): + # given + push_embedder = PostgresEmbedder(MagicMock(), env_helper_mock) + source_url = "some-url" + + # when + push_embedder.embed_file( + source_url, + "some-file-name.pdf", + ) + + # then + document_loading_mock.return_value.load.assert_called_once_with( + source_url, LOADING_SETTINGS + ) + + +def test_postgres_embed_file_chunks_documents( + document_loading_mock, document_chunking_mock, env_helper_mock +): + # given + push_embedder = PostgresEmbedder(MagicMock(), env_helper_mock) + + # when + push_embedder.embed_file( + "some-url", + "some-file-name.pdf", + ) + + # then + document_chunking_mock.return_value.chunk.assert_called_once_with( + document_loading_mock.return_value.load.return_value, CHUNKING_SETTINGS + ) From 270386c0b7dbea2aaa748f4aed617fc03dc15fff Mon Sep 17 00:00:00 2001 From: Pavan Kumar Date: Mon, 2 Dec 2024 21:41:25 +0530 Subject: [PATCH 3/8] Mock default credentials --- .../helpers/test_azure_postgres_helper.py | 174 ++++++++++++++++-- 1 file changed, 158 insertions(+), 16 deletions(-) diff --git a/code/tests/utilities/helpers/test_azure_postgres_helper.py b/code/tests/utilities/helpers/test_azure_postgres_helper.py index 907d0bdfd..404e6f6df 100644 --- a/code/tests/utilities/helpers/test_azure_postgres_helper.py +++ b/code/tests/utilities/helpers/test_azure_postgres_helper.py @@ -134,8 +134,11 @@ def raise_exception(*args, **kwargs): self.assertEqual(str(context.exception), "Query execution error") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") - def test_create_search_client_connection_error(self, mock_connect): + def test_create_search_client_connection_error(self, mock_connect, mock_credential): # Arrange # Mock the EnvHelper and set required attributes mock_env_helper = MagicMock() @@ -144,6 +147,11 @@ def test_create_search_client_connection_error(self, mock_connect): mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + def raise_exception(*args, **kwargs): raise Exception("Connection error") @@ -158,15 +166,23 @@ def raise_exception(*args, **kwargs): self.assertEqual(str(context.exception), "Connection error") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_get_files_success(self, mock_env_helper, mock_connect): + def test_get_files_success(self, mock_env_helper, mock_connect, mock_credential): # Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Arrange: Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -191,15 +207,23 @@ def test_get_files_success(self, mock_env_helper, mock_connect): ) mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_get_files_no_results(self, mock_env_helper, mock_connect): + def test_get_files_no_results(self, mock_env_helper, mock_connect, mock_credential): # Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Arrange: Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -219,16 +243,26 @@ def test_get_files_no_results(self, mock_env_helper, mock_connect): self.assertIsNone(result) mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") - def test_get_files_db_error(self, mock_logger, mock_env_helper, mock_connect): + def test_get_files_db_error( + self, mock_logger, mock_env_helper, mock_connect, mock_credential + ): # Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Arrange: Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -250,11 +284,14 @@ def test_get_files_db_error(self, mock_logger, mock_env_helper, mock_connect): ) mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") def test_get_files_unexpected_error( - self, mock_logger, mock_env_helper, mock_connect + self, mock_logger, mock_env_helper, mock_connect, mock_credential ): # Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" @@ -262,6 +299,11 @@ def test_get_files_unexpected_error( mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Arrange: Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -283,16 +325,26 @@ def test_get_files_unexpected_error( ) mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_delete_documents_success(self, mock_env_helper, mock_logger, mock_connect): + def test_delete_documents_success( + self, mock_env_helper, mock_logger, mock_connect, mock_credential + ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -317,16 +369,26 @@ def test_delete_documents_success(self, mock_env_helper, mock_logger, mock_conne mock_connection.close.assert_called_once() mock_logger.info.assert_called_with("Deleted 3 documents.") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_delete_documents_no_ids(self, mock_env_helper, mock_logger, mock_connect): + def test_delete_documents_no_ids( + self, mock_env_helper, mock_logger, mock_connect, mock_credential + ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -347,11 +409,14 @@ def test_delete_documents_no_ids(self, mock_env_helper, mock_logger, mock_connec mock_logger.warning.assert_called_with("No IDs provided for deletion.") mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") def test_delete_documents_db_error( - self, mock_env_helper, mock_logger, mock_connect + self, mock_env_helper, mock_logger, mock_connect, mock_credential ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" @@ -359,6 +424,11 @@ def test_delete_documents_db_error( mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -383,11 +453,14 @@ def test_delete_documents_db_error( mock_connection.rollback.assert_called_once() mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") def test_delete_documents_unexpected_error( - self, mock_env_helper, mock_logger, mock_connect + self, mock_env_helper, mock_logger, mock_connect, mock_credential ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" @@ -395,6 +468,11 @@ def test_delete_documents_unexpected_error( mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -419,16 +497,26 @@ def test_delete_documents_unexpected_error( mock_connection.rollback.assert_called_once() mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_perform_search_success(self, mock_env_helper, mock_logger, mock_connect): + def test_perform_search_success( + self, mock_env_helper, mock_logger, mock_connect, mock_credential + ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -462,11 +550,14 @@ def test_perform_search_success(self, mock_env_helper, mock_logger, mock_connect mock_connection.close.assert_called_once() mock_logger.info.assert_called_with("Retrieved 1 search result(s).") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") def test_perform_search_no_results( - self, mock_env_helper, mock_logger, mock_connect + self, mock_env_helper, mock_logger, mock_connect, mock_credential ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" @@ -474,6 +565,11 @@ def test_perform_search_no_results( mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -498,16 +594,26 @@ def test_perform_search_no_results( mock_connection.close.assert_called_once() mock_logger.info.assert_called_with("Retrieved 0 search result(s).") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_perform_search_error(self, mock_env_helper, mock_logger, mock_connect): + def test_perform_search_error( + self, mock_env_helper, mock_logger, mock_connect, mock_credential + ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -531,16 +637,26 @@ def test_perform_search_error(self, mock_env_helper, mock_logger, mock_connect): ) mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_get_unique_files_success(self, mock_env_helper, mock_logger, mock_connect): + def test_get_unique_files_success( + self, mock_env_helper, mock_logger, mock_connect, mock_credential + ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -568,11 +684,14 @@ def test_get_unique_files_success(self, mock_env_helper, mock_logger, mock_conne mock_connection.close.assert_called_once() mock_logger.info.assert_called_with("Retrieved 2 unique title(s).") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") def test_get_unique_files_no_results( - self, mock_env_helper, mock_logger, mock_connect + self, mock_env_helper, mock_logger, mock_connect, mock_credential ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" @@ -580,6 +699,11 @@ def test_get_unique_files_no_results( mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -602,16 +726,26 @@ def test_get_unique_files_no_results( mock_connection.close.assert_called_once() mock_logger.info.assert_called_with("Retrieved 0 unique title(s).") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_get_unique_files_error(self, mock_env_helper, mock_logger, mock_connect): + def test_get_unique_files_error( + self, mock_env_helper, mock_logger, mock_connect, mock_credential + ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -633,11 +767,14 @@ def test_get_unique_files_error(self, mock_env_helper, mock_logger, mock_connect ) mock_connection.close.assert_called_once() + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") def test_search_by_blob_url_success( - self, mock_env_helper, mock_logger, mock_connect + self, mock_env_helper, mock_logger, mock_connect, mock_credential ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" @@ -645,6 +782,11 @@ def test_search_by_blob_url_success( mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() From 6198de76bd7f008cb2a781c259f91ec76a0590dd Mon Sep 17 00:00:00 2001 From: Pavan Kumar Date: Mon, 2 Dec 2024 22:10:15 +0530 Subject: [PATCH 4/8] mock default credentials --- .../helpers/test_azure_postgres_helper.py | 32 +++++++++++++++++-- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/code/tests/utilities/helpers/test_azure_postgres_helper.py b/code/tests/utilities/helpers/test_azure_postgres_helper.py index 404e6f6df..e2cb9c438 100644 --- a/code/tests/utilities/helpers/test_azure_postgres_helper.py +++ b/code/tests/utilities/helpers/test_azure_postgres_helper.py @@ -101,8 +101,11 @@ def test_get_search_indexes_success( "host=mock_host user=mock_user dbname=mock_database password=mock-access-token" ) + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") - def test_get_search_indexes_query_error(self, mock_connect): + def test_get_search_indexes_query_error(self, mock_connect, mock_credential): # Arrange # Mock the EnvHelper and set required attributes mock_env_helper = MagicMock() @@ -111,6 +114,11 @@ def test_get_search_indexes_query_error(self, mock_connect): mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + mock_connection = MagicMock() mock_connect.return_value = mock_connection @@ -815,11 +823,14 @@ def test_search_by_blob_url_success( mock_connection.close.assert_called_once() mock_logger.info.assert_called_with("Retrieved 2 unique title(s).") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") def test_search_by_blob_url_no_results( - self, mock_env_helper, mock_logger, mock_connect + self, mock_env_helper, mock_logger, mock_connect, mock_credential ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" @@ -827,6 +838,11 @@ def test_search_by_blob_url_no_results( mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() @@ -850,16 +866,26 @@ def test_search_by_blob_url_no_results( mock_connection.close.assert_called_once() mock_logger.info.assert_called_with("Retrieved 0 unique title(s).") + @patch( + "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential" + ) @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect") @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger") @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper") - def test_search_by_blob_url_error(self, mock_env_helper, mock_logger, mock_connect): + def test_search_by_blob_url_error( + self, mock_env_helper, mock_logger, mock_connect, mock_credential + ): # Arrange: Mock the EnvHelper attributes mock_env_helper.POSTGRESQL_USER = "mock_user" mock_env_helper.POSTGRESQL_HOST = "mock_host" mock_env_helper.POSTGRESQL_DATABASE = "mock_database" mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5 + # Mock access token retrieval + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + # Mock the connection and cursor mock_connection = MagicMock() mock_cursor = MagicMock() From 11d78d09b99c0a4818aa491fbade986e45f1e68a Mon Sep 17 00:00:00 2001 From: Pavan Kumar Date: Tue, 3 Dec 2024 19:56:14 +0530 Subject: [PATCH 5/8] Set PostgreSQL as the Default Database for CWYD Deployment & unittestcases --- .../batch/utilities/helpers/env_helper.py | 4 +- .../test_postgres_search_handler.py | 49 ++++++++ code/tests/test_chat_history.py | 119 +++++++++++++++++- .../helpers/test_database_factory.py | 89 +++++++++++++ infra/main.bicep | 14 ++- infra/main.json | 4 +- 6 files changed, 270 insertions(+), 9 deletions(-) create mode 100644 code/tests/utilities/helpers/test_database_factory.py diff --git a/code/backend/batch/utilities/helpers/env_helper.py b/code/backend/batch/utilities/helpers/env_helper.py index e1d40a162..db2f74ebb 100644 --- a/code/backend/batch/utilities/helpers/env_helper.py +++ b/code/backend/batch/utilities/helpers/env_helper.py @@ -91,7 +91,9 @@ def __load_config(self, **kwargs) -> None: # Chat History DB Integration Settings # Set default values based on DATABASE_TYPE - self.DATABASE_TYPE = os.getenv("DATABASE_TYPE", "").strip() or "CosmosDB" + self.DATABASE_TYPE = ( + os.getenv("DATABASE_TYPE", "").strip() or DatabaseType.POSTGRESQL.value + ) # Cosmos DB configuration if self.DATABASE_TYPE == DatabaseType.COSMOSDB.value: azure_cosmosdb_info = self.get_info_from_env("AZURE_COSMOSDB_INFO", "") diff --git a/code/tests/search_utilities/test_postgres_search_handler.py b/code/tests/search_utilities/test_postgres_search_handler.py index eead10dd3..1c8117791 100644 --- a/code/tests/search_utilities/test_postgres_search_handler.py +++ b/code/tests/search_utilities/test_postgres_search_handler.py @@ -1,3 +1,4 @@ +import json import pytest from unittest.mock import MagicMock, patch from backend.batch.utilities.common.source_document import SourceDocument @@ -124,6 +125,54 @@ def test_get_files(handler): assert result[1] == "test2.txt" +def test_output_results(handler): + results = [ + {"id": "1", "title": "file1.txt"}, + {"id": "2", "title": "file2.txt"}, + {"id": "3", "title": "file1.txt"}, + {"id": "4", "title": "file3.txt"}, + {"id": "5", "title": "file2.txt"}, + ] + + expected_output = { + "file1.txt": ["1", "3"], + "file2.txt": ["2", "5"], + "file3.txt": ["4"], + } + + result = handler.output_results(results) + + assert result == expected_output + assert len(result) == 3 + assert "file1.txt" in result + assert result["file2.txt"] == ["2", "5"] + + +def test_process_results(handler): + results = [ + {"metadata": json.dumps({"chunk": "Chunk1"}), "content": "Content1"}, + {"metadata": json.dumps({"chunk": "Chunk2"}), "content": "Content2"}, + ] + expected_output = [["Chunk1", "Content1"], ["Chunk2", "Content2"]] + result = handler.process_results(results) + assert result == expected_output + + +def test_process_results_none(handler): + result = handler.process_results(None) + assert result == [] + + +def test_process_results_missing_chunk(handler): + results = [ + {"metadata": json.dumps({}), "content": "Content1"}, + {"metadata": json.dumps({"chunk": "Chunk2"}), "content": "Content2"}, + ] + expected_output = [[0, "Content1"], ["Chunk2", "Content2"]] + result = handler.process_results(results) + assert result == expected_output + + def test_delete_files(handler): files_to_delete = {"test1.txt": [1, 2], "test2.txt": [3]} mock_delete_documents = MagicMock() diff --git a/code/tests/test_chat_history.py b/code/tests/test_chat_history.py index f1b8bdcb1..6ef805d50 100644 --- a/code/tests/test_chat_history.py +++ b/code/tests/test_chat_history.py @@ -2,7 +2,7 @@ This module tests the entry point for the application. """ -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from create_app import create_app @@ -555,6 +555,54 @@ def test_update_conversation_success( "success": True, } + @patch("backend.api.chat_history.AsyncAzureOpenAI") + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_update_conversation_new_success( + self, + get_active_config_or_default_mock, + azure_openai_mock: MagicMock, + mock_conversation_client, + client, + ): + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversation.return_value = [] + mock_conversation_client.create_message.return_value = "success" + mock_conversation_client.create_conversation.return_value = { + "title": "Test Title", + "updatedAt": "2024-12-01", + "id": "conv1", + } + request_json = { + "conversation_id": "conv1", + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + ], + } + + openai_client_mock = azure_openai_mock.return_value + + mock_response = MagicMock() + mock_response.choices = [MagicMock(message=MagicMock(content="Test Title"))] + + openai_client_mock.chat.completions.create = AsyncMock( + return_value=mock_response + ) + + response = client.post("/api/history/update", json=request_json) + + assert response.status_code == 200 + assert response.json == { + "data": { + "conversation_id": "conv1", + "date": "2024-12-01", + "title": "Test Title", + }, + "success": True, + } + @patch( "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" ) @@ -568,6 +616,75 @@ def test_update_conversation_no_chat_history( assert response.status_code == 400 assert response.json == {"error": "Chat history is not available"} + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_update_conversation_connect_error( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.get_conversation.return_value = { + "title": "Test Title", + "updatedAt": "2024-12-01", + "id": "conv1", + } + request_json = { + "conversation_id": "conv1", + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + ], + } + mock_conversation_client.connect.side_effect = Exception("Unexpected error") + + # Make the API call + response = client.post( + "/api/history/update", + json=request_json, + headers={"Content-Type": "application/json"}, + ) + + # Assert response + assert response.status_code == 500 + assert response.json == { + "error": "Error while updating the conversation history" + } + + @patch( + "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" + ) + def test_update_conversation_error( + self, get_active_config_or_default_mock, mock_conversation_client, client + ): + get_active_config_or_default_mock.return_value.enable_chat_history = True + mock_conversation_client.create_message.side_effect = Exception( + "Unexpected error" + ) + mock_conversation_client.get_conversation.return_value = { + "title": "Test Title", + "updatedAt": "2024-12-01", + "id": "conv1", + } + request_json = { + "conversation_id": "conv1", + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + ], + } + + response = client.post( + "/api/history/update", + json=request_json, + headers={"Content-Type": "application/json"}, + ) + + # Assert response + assert response.status_code == 500 + assert response.json == { + "error": "Error while updating the conversation history" + } + @patch( "backend.batch.utilities.helpers.config.config_helper.ConfigHelper.get_active_config_or_default" ) diff --git a/code/tests/utilities/helpers/test_database_factory.py b/code/tests/utilities/helpers/test_database_factory.py new file mode 100644 index 000000000..0a1734171 --- /dev/null +++ b/code/tests/utilities/helpers/test_database_factory.py @@ -0,0 +1,89 @@ +import pytest +from unittest.mock import patch, MagicMock +from backend.batch.utilities.helpers.config.database_type import DatabaseType +from backend.batch.utilities.chat_history.cosmosdb import CosmosConversationClient +from backend.batch.utilities.chat_history.database_factory import DatabaseFactory +from backend.batch.utilities.chat_history.postgresdbservice import ( + PostgresConversationClient, +) + + +@patch("backend.batch.utilities.chat_history.database_factory.DefaultAzureCredential") +@patch("backend.batch.utilities.chat_history.database_factory.EnvHelper") +@patch( + "backend.batch.utilities.chat_history.database_factory.CosmosConversationClient", + autospec=True, +) +def test_get_conversation_client_cosmos( + mock_cosmos_client, mock_env_helper, mock_credential +): + # Configure the EnvHelper mock + mock_env_instance = mock_env_helper.return_value + mock_env_instance.DATABASE_TYPE = DatabaseType.COSMOSDB.value + mock_env_instance.AZURE_COSMOSDB_ACCOUNT = "cosmos_account" + mock_env_instance.AZURE_COSMOSDB_DATABASE = "cosmos_database" + mock_env_instance.AZURE_COSMOSDB_CONVERSATIONS_CONTAINER = "conversations_container" + mock_env_instance.AZURE_COSMOSDB_ENABLE_FEEDBACK = False + mock_env_instance.AZURE_COSMOSDB_ACCOUNT_KEY = None + + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + mock_credential_instance = mock_credential.return_value + + # Mock the CosmosConversationClient instance + mock_cosmos_instance = MagicMock(spec=CosmosConversationClient) + mock_cosmos_client.return_value = mock_cosmos_instance + + # Call the method + client = DatabaseFactory.get_conversation_client() + + # Assert the CosmosConversationClient was called with correct arguments + mock_cosmos_client.assert_called_once_with( + cosmosdb_endpoint="https://cosmos_account.documents.azure.com:443/", + credential=mock_credential_instance, + database_name="cosmos_database", + container_name="conversations_container", + enable_message_feedback=False, + ) + assert isinstance(client, CosmosConversationClient) + assert client == mock_cosmos_instance + + +@patch("backend.batch.utilities.chat_history.database_factory.DefaultAzureCredential") +@patch("backend.batch.utilities.chat_history.database_factory.EnvHelper") +@patch( + "backend.batch.utilities.chat_history.database_factory.PostgresConversationClient", + autospec=True, +) +def test_get_conversation_client_postgres( + mock_postgres_client, mock_env_helper, mock_credential +): + mock_env_instance = mock_env_helper.return_value + mock_env_instance.DATABASE_TYPE = DatabaseType.POSTGRESQL.value + mock_env_instance.POSTGRESQL_USER = "postgres_user" + mock_env_instance.POSTGRESQL_HOST = "postgres_host" + mock_env_instance.POSTGRESQL_DATABASE = "postgres_database" + + mock_access_token = MagicMock() + mock_access_token.token = "mock-access-token" + mock_credential.return_value.get_token.return_value = mock_access_token + + mock_postgres_instance = MagicMock(spec=PostgresConversationClient) + mock_postgres_client.return_value = mock_postgres_instance + + client = DatabaseFactory.get_conversation_client() + + mock_postgres_client.assert_called_once_with( + user="postgres_user", host="postgres_host", database="postgres_database" + ) + assert isinstance(client, PostgresConversationClient) + + +@patch("backend.batch.utilities.chat_history.database_factory.EnvHelper") +def test_get_conversation_client_invalid_database_type(mock_env_helper): + mock_env_instance = mock_env_helper.return_value + mock_env_instance.DATABASE_TYPE = "INVALID_DB" + + with pytest.raises(ValueError, match="Unsupported DATABASE_TYPE"): + DatabaseFactory.get_conversation_client() diff --git a/infra/main.bicep b/infra/main.bicep index 2d061e9d9..3a4f426a7 100644 --- a/infra/main.bicep +++ b/infra/main.bicep @@ -306,7 +306,7 @@ param azureMachineLearningName string = 'aml-${resourceToken}' 'CosmosDB' 'PostgreSQL' ]) -param databaseType string = 'CosmosDB' +param databaseType string = 'PostgreSQL' @description('Azure Cosmos DB Account Name') param azureCosmosDBAccountName string = 'cosmos-${resourceToken}' @@ -1258,13 +1258,17 @@ module createIndex './core/database/deploy_create_table_script.bicep' = if (data keyVaultName: keyvault.outputs.name postgresSqlServerName: postgresDBModule.outputs.postgresDbOutput.postgresSQLName webAppPrincipalName: hostingModel == 'code' ? web.outputs.FRONTEND_API_NAME : web_docker.outputs.FRONTEND_API_NAME - adminAppPrincipalName: hostingModel == 'code' ? adminweb.outputs.WEBSITE_ADMIN_NAME : adminweb_docker.outputs.WEBSITE_ADMIN_NAME + adminAppPrincipalName: hostingModel == 'code' + ? adminweb.outputs.WEBSITE_ADMIN_NAME + : adminweb_docker.outputs.WEBSITE_ADMIN_NAME managedIdentityName: managedIdentityModule.outputs.managedIdentityOutput.name } scope: rg - dependsOn: hostingModel == 'code' ? [keyvault, postgresDBModule, storekeys, web, adminweb] : [ - [keyvault, postgresDBModule, storekeys, web_docker, adminweb_docker] - ] + dependsOn: hostingModel == 'code' + ? [keyvault, postgresDBModule, storekeys, web, adminweb] + : [ + [keyvault, postgresDBModule, storekeys, web_docker, adminweb_docker] + ] } output APPLICATIONINSIGHTS_CONNECTION_STRING string = monitoring.outputs.applicationInsightsConnectionString diff --git a/infra/main.json b/infra/main.json index 328f4b9a1..c7bd6a8e7 100644 --- a/infra/main.json +++ b/infra/main.json @@ -5,7 +5,7 @@ "_generator": { "name": "bicep", "version": "0.30.23.60470", - "templateHash": "15176273744623029817" + "templateHash": "10654587243799689217" } }, "parameters": { @@ -614,7 +614,7 @@ }, "databaseType": { "type": "string", - "defaultValue": "CosmosDB", + "defaultValue": "PostgreSQL", "allowedValues": [ "CosmosDB", "PostgreSQL" From dbba5883f43befdeda997e2d4321dea7e1402c0c Mon Sep 17 00:00:00 2001 From: Pavan Kumar Date: Tue, 3 Dec 2024 20:29:26 +0530 Subject: [PATCH 6/8] fix DATABASE_TYPE for testcases --- code/tests/functional/app_config.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/code/tests/functional/app_config.py b/code/tests/functional/app_config.py index 0dd14b2f6..a072d7f92 100644 --- a/code/tests/functional/app_config.py +++ b/code/tests/functional/app_config.py @@ -13,7 +13,9 @@ class AppConfig: config: dict[str, str | None] = { "APPLICATIONINSIGHTS_ENABLED": "False", "AZURE_AUTH_TYPE": "keys", - "AZURE_BLOB_STORAGE_INFO": '{"accountName": "some-blob-account-name", "containerName": "some-blob-container-name", "accountKey": "' + encoded_account_key + '"}', + "AZURE_BLOB_STORAGE_INFO": '{"accountName": "some-blob-account-name", "containerName": "some-blob-container-name", "accountKey": "' + + encoded_account_key + + '"}', "AZURE_COMPUTER_VISION_KEY": "some-computer-vision-key", "AZURE_CONTENT_SAFETY_ENDPOINT": "some-content-safety-endpoint", "AZURE_CONTENT_SAFETY_KEY": "some-content-safety-key", @@ -80,6 +82,7 @@ class AppConfig: "OPENAI_API_TYPE": None, "OPENAI_API_KEY": None, "OPENAI_API_VERSION": None, + "DATABASE_TYPE": "CosmosDB", } def __init__(self, config_overrides: dict[str, str | None] = {}) -> None: From dfb2fe61c6126777b0c5391114e94da110b53661 Mon Sep 17 00:00:00 2001 From: Pavan Kumar Date: Tue, 3 Dec 2024 20:46:01 +0530 Subject: [PATCH 7/8] fix boolean value --- code/backend/batch/utilities/helpers/env_helper.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/code/backend/batch/utilities/helpers/env_helper.py b/code/backend/batch/utilities/helpers/env_helper.py index db2f74ebb..fa82e2796 100644 --- a/code/backend/batch/utilities/helpers/env_helper.py +++ b/code/backend/batch/utilities/helpers/env_helper.py @@ -124,8 +124,8 @@ def __load_config(self, **kwargs) -> None: self.POSTGRESQL_DATABASE = azure_postgresql_info.get("dbname", "") self.POSTGRESQL_HOST = azure_postgresql_info.get("host", "") # Ensure integrated vectorization is disabled for PostgreSQL - self.AZURE_SEARCH_USE_INTEGRATED_VECTORIZATION = "False" - self.USE_ADVANCED_IMAGE_PROCESSING = "False" + self.AZURE_SEARCH_USE_INTEGRATED_VECTORIZATION = False + self.USE_ADVANCED_IMAGE_PROCESSING = False else: raise ValueError( "Unsupported DATABASE_TYPE. Please set DATABASE_TYPE to 'CosmosDB' or 'PostgreSQL'." From 0e11306f83a52135f654284060567ec48f16f787 Mon Sep 17 00:00:00 2001 From: Pavan Kumar Date: Tue, 3 Dec 2024 21:23:42 +0530 Subject: [PATCH 8/8] fix testcase --- code/tests/utilities/helpers/test_env_helper.py | 1 + 1 file changed, 1 insertion(+) diff --git a/code/tests/utilities/helpers/test_env_helper.py b/code/tests/utilities/helpers/test_env_helper.py index 10e1de308..8acd1e497 100644 --- a/code/tests/utilities/helpers/test_env_helper.py +++ b/code/tests/utilities/helpers/test_env_helper.py @@ -133,6 +133,7 @@ def test_azure_speech_recognizer_languages_default(monkeypatch: MonkeyPatch): ) def test_use_advanced_image_processing(monkeypatch: MonkeyPatch, value, expected): # given + monkeypatch.setenv("DATABASE_TYPE", "CosmosDB") if value is not None: monkeypatch.setenv("USE_ADVANCED_IMAGE_PROCESSING", value)