tests/test_text_summarization.py

import unittest
from unittest.mock import patch
from transformers.pipelines.text2text_generation import SummarizationPipeline
from level_1_text_summarization import get_pipeline, summarize_text

class TestTextSummarization(unittest.TestCase):

    def test_get_pipeline(self):
        """
        Tests if the get_pipeline function returns a valid summarization pipeline.
        """
        with patch('level_1_text_summarization.pipeline') as mock_pipeline:
            # Mock the pipeline call to avoid actual model loading during tests
            mock_pipeline.return_value = "MockedPipeline"
            summarizer = get_pipeline()
            # We check if the function returns what the mocked pipeline returned.
            # In a more complex scenario, we would assert its type.
            self.assertIsNotNone(summarizer)
            self.assertIsInstance(summarizer, str) # In this mocked case, it's a string.

    @patch('level_1_text_summarization.get_pipeline')
    def test_summarize_text(self, mock_get_pipeline):
        """
        Tests the summarize_text function for basic functionality.
        """
        # Mock the pipeline object itself
        class MockSummarizer:
            def __call__(self, text, max_length, min_length, do_sample):
                return [{'summary_text': 'This is a test summary.'}]

        mock_summarizer_instance = MockSummarizer()
        
        # The text to be summarized
        sample_text = "This is a long piece of text that needs to be summarized for a unit test."
        
        # Call the function with the mocked summarizer
        summary = summarize_text(mock_summarizer_instance, sample_text)
        
        # Assert that the summary is a non-empty string
        self.assertIsInstance(summary, str)
        self.assertTrue(len(summary) > 0)
        self.assertEqual(summary, 'This is a test summary.')

    def test_summarize_text_with_none_pipeline(self):
        """
        Tests that summarize_text returns None if the pipeline is not available.
        """
        summary = summarize_text(None, "Some text")
        self.assertIsNone(summary)

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