From 11d6460be4b1444880c12e2e6a794a95fc0de86b Mon Sep 17 00:00:00 2001 From: bachvudinh Date: Thu, 18 Jul 2024 16:40:22 +0700 Subject: [PATCH] finish debugging and ready for merging --- tests/test_flows.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/tests/test_flows.py b/tests/test_flows.py index 7c3af9c..d22825e 100644 --- a/tests/test_flows.py +++ b/tests/test_flows.py @@ -99,16 +99,13 @@ def test_tokenizer_loading(self, mock_from_pretrained): @test_name("Dataset Loading Test") @patch('s3helper.s3_load_dataset') - def test_dataset_loading(self, mock_s3_load_dataset): - mock_dataset = MagicMock() - mock_s3_load_dataset.return_value = mock_dataset - - dataset = s3_load_dataset(self.dataset_name, file_format='parquet', split='train') - print(dataset) - # mock_s3_load_dataset.assert_called_once_with(self.dataset_name, file_format='parquet', split='train') - self.assertIsNotNone(dataset) - # self.assertEqual(dataset, mock_dataset) - + def test_dataset_loading(self): + try: + dataset = s3_load_dataset(self.dataset_name, file_format='parquet', split='train') + print(dataset) + self.assertIsNotNone(dataset) + except Exception as e: + self.fail(f"s3_load_dataset raised an exception: {e}") @test_name("Config Loading Test") @patch('s3helper.S3HelperAutoConfig.from_pretrained') def test_config_loading(self, mock_from_pretrained):