aki_prj23_transparenzregister/tests/ai/sentiment_pipeline_test.py

"""Unit test for sentiment pipeline."""

from unittest.mock import Mock, patch

import pytest

from aki_prj23_transparenzregister.ai.sentiment_pipeline import (
    SentimentPipeline,
)
from aki_prj23_transparenzregister.config.config_template import MongoConnection


@pytest.fixture()
def mock_mongo_connection() -> MongoConnection:
    """Mock MongoConnector class.

    Args:
        mocker (any): Library mocker

    Returns:
        Mock: Mocked MongoConnector
    """
    return MongoConnection("", "", None, "" "", "")


@pytest.fixture()
def mock_mongo_connector(mocker: Mock) -> Mock:
    """Mock MongoConnector class.

    Args:
        mocker (any): Library mocker

    Returns:
        Mock: Mocked MongoConnector
    """
    mock = Mock()
    mocker.patch(
        "aki_prj23_transparenzregister.utils.mongo.connector.MongoConnector",
        return_value=mock,
    )
    mock.database = {"news": Mock()}
    return mock


@pytest.fixture()
def mock_spacy(mocker: Mock) -> Mock:
    """Mock MongoConnector class.

    Args:
        mocker (any): Library mocker

    Returns:
        Mock: Mocked MongoConnector
    """
    mock = Mock()
    mocker.patch(
        "aki_prj23_transparenzregister.ai.sentiment_service.SentimentAnalysisService.init_spacy",
        return_value=mock,
    )
    return mock


@patch(
    "aki_prj23_transparenzregister.ai.sentiment_service.SentimentAnalysisService.sentiment_spacy"
)
def test_sentiment_pipeline_existing_sentiment(
    mock_sentiment_spacy: Mock,
    mock_mongo_connector: Mock,
    mock_mongo_connection: MongoConnection,
    mock_spacy: Mock,
) -> None:
    # Configure the mock to return a specific sentiment result
    mock_sentiment_spacy.return_value = ("positive", 0.8)

    # Create an instance of the SentimentPipeline
    sentiment_pipeline = SentimentPipeline(mock_mongo_connection)

    # Mock the news collection and documents for testing
    mock_collection = Mock()
    mock_documents = [
        {
            "_id": "document1",
            "text": "This is a positive text.",
            "sentiment": {"label": "neutral", "score": 0.5},
        }
    ]

    # Set the collection to the mock_collection
    sentiment_pipeline.news_obj.collection = mock_collection

    # Mock the find method of the collection to return the mock documents
    mock_collection.find.return_value = mock_documents

    # Call the process_documents method
    sentiment_pipeline.process_documents("text", "use_spacy")

    # Ensure that sentiment_spacy was called with the correct text
    mock_sentiment_spacy.assert_called_once_with("This is a positive text.")

    # Ensure that the document in the collection was not updated with sentiment
    # mock_collection.update_one.assert_not_called()


@patch(
    "aki_prj23_transparenzregister.ai.sentiment_service.SentimentAnalysisService.sentiment_spacy"
)
def test_sentiment_pipeline_no_documents(
    mock_sentiment_spacy: Mock,
    mock_mongo_connector: Mock,
    mock_mongo_connection: MongoConnection,
    mock_spacy: Mock,
) -> None:
    # Configure the mock to return a specific sentiment result
    mock_sentiment_spacy.return_value = ("positive", 0.8)

    # Create an instance of the SentimentPipeline
    sentiment_pipeline = SentimentPipeline(mock_mongo_connection)

    # Mock the news collection to return an empty result
    mock_collection = Mock()
    mock_collection.find.return_value = []

    # Set the collection to the mock_collection
    sentiment_pipeline.news_obj.collection = mock_collection

    # Call the process_documents method
    sentiment_pipeline.process_documents("text", "use_spacy")

    # Ensure that sentiment_spacy was not called
    mock_sentiment_spacy.assert_not_called()

    # Ensure that the document in the collection was not updated with sentiment
    mock_collection.update_one.assert_not_called()


@patch(
    "aki_prj23_transparenzregister.ai.sentiment_service.SentimentAnalysisService.sentiment_spacy"
)
def test_sentiment_pipeline_with_spacy(
    mock_sentiment_spacy: Mock,
    mock_mongo_connector: Mock,
    mock_mongo_connection: MongoConnection,
    mock_spacy: Mock,
) -> None:
    # Configure the mock to return a specific sentiment result
    mock_sentiment_spacy.return_value = ("positive", 0.8)

    # Create an instance of the SentimentPipeline
    sentiment_pipeline = SentimentPipeline(mock_mongo_connection)

    # Mock the news collection and documents for testing
    mock_collection = Mock()
    mock_documents = [{"_id": "document1", "text": "This is a positive text."}]

    # Set the collection to the mock_collection
    sentiment_pipeline.news_obj.collection = mock_collection

    # Mock the find method of the collection to return the mock documents
    mock_collection.find.return_value = mock_documents

    # Call the process_documents method
    sentiment_pipeline.process_documents("text", "use_spacy")

    # Ensure that sentiment_spacy was called with the correct text
    mock_sentiment_spacy.assert_called_once_with("This is a positive text.")

    # Ensure that the document in the collection was updated with the sentiment result
    mock_collection.update_one.assert_called_once_with(
        {"_id": "document1"},
        {"$set": {"sentiment": {"label": "positive", "score": 0.8}}},
    )


# Mocking the SentimentAnalysisService methods
@patch(
    "aki_prj23_transparenzregister.ai.sentiment_service.SentimentAnalysisService.sentiment_transformer"
)
def test_sentiment_pipeline_with_transformer(
    mock_sentiment_transformer: Mock,
    mock_mongo_connector: Mock,
    mock_mongo_connection: MongoConnection,
    mock_spacy: Mock,
) -> None:
    # Configure the mock to return a specific sentiment result
    mock_sentiment_transformer.return_value = ("negative", 0.6)

    # Create an instance of the SentimentPipeline
    sentiment_pipeline = SentimentPipeline(mock_mongo_connection)

    # Mock the news collection and documents for testing
    mock_collection = Mock()
    mock_documents = [{"_id": "document2", "text": "This is a negative text."}]

    # Set the collection to the mock_collection
    sentiment_pipeline.news_obj.collection = mock_collection

    # Mock the find method of the collection to return the mock documents
    mock_collection.find.return_value = mock_documents

    # Call the process_documents method
    sentiment_pipeline.process_documents("text", "use_transformer")

    # Ensure that sentiment_transformer was called with the correct text
    mock_sentiment_transformer.assert_called_once_with("This is a negative text.")

    # Ensure that the document in the collection was updated with the sentiment result
    mock_collection.update_one.assert_called_once_with(
        {"_id": "document2"},
        {"$set": {"sentiment": {"label": "negative", "score": 0.6}}},
    )