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()