reverst the code back

janhq · Jul 18, 2024 · db599fe · db599fe
1 parent d4535a5
commit db599fe
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/test_flows.py b/tests/test_flows.py
@@ -86,6 +86,7 @@ def setUpClass(cls):
         S3Helper()
 
     @test_name("Tokenizer Loading Test")
+    @patch('s3helper.S3HelperAutoTokenizer.from_pretrained')
     def test_tokenizer_loading(self, mock_from_pretrained):
         mock_tokenizer = MagicMock()
         mock_from_pretrained.return_value = mock_tokenizer
@@ -106,6 +107,7 @@ def test_dataset_loading(self):
         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):
         mock_config = MagicMock()
         mock_from_pretrained.return_value = mock_config
@@ -117,6 +119,7 @@ def test_config_loading(self, mock_from_pretrained):
         self.assertEqual(config, mock_config)
 
     @test_name("Model Loading Test")
+    @patch('s3helper.S3HelperAutoModelForCausalLM.from_pretrained')
     def test_model_loading(self, mock_from_pretrained):
         mock_model = MagicMock()
         mock_from_pretrained.return_value = mock_model