From 349ea74353b93d714cdee392400374c03ea6ce0d Mon Sep 17 00:00:00 2001 From: Tim Ryan Date: Mon, 17 Jul 2023 14:27:55 -0400 Subject: [PATCH] Revert "RM-96 fix vertica driver test" This reverts commit 3581ff60c2b28ce9dc60db6f48c44d444f83cf27. --- tests/unit/db/vertica/test_vertica_db_driver.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/unit/db/vertica/test_vertica_db_driver.py b/tests/unit/db/vertica/test_vertica_db_driver.py index 21f1434d0..ef6415a33 100644 --- a/tests/unit/db/vertica/test_vertica_db_driver.py +++ b/tests/unit/db/vertica/test_vertica_db_driver.py @@ -1,5 +1,5 @@ from .base_test_vertica_db_driver import BaseTestVerticaDBDriver -from mock import Mock +from mock import Mock, MagicMock from ...records.format_hints import (vertica_format_hints) import sqlalchemy @@ -8,7 +8,10 @@ class TestVerticaDBDriver(BaseTestVerticaDBDriver): def test_unload(self): mock_result = Mock(name='result') mock_result.rows = 579 - self.mock_db_engine.execute.return_value.fetchall.return_value = [mock_result] + mock_connection = MagicMock(name='connection') + self.mock_db_engine.connect.return_value \ + .__enter__.return_value = mock_connection + mock_connection.execute.return_value.fetchall.return_value = [mock_result] self.mock_records_unload_plan.processing_instructions.fail_if_dont_understand = True self.mock_records_unload_plan.processing_instructions.fail_if_cant_handle_hint = True @@ -25,7 +28,10 @@ def test_unload(self): def test_unload_to_non_s3(self): mock_result = Mock(name='result') mock_result.rows = 579 - self.mock_db_engine.execute.return_value.fetchall.return_value = [mock_result] + mock_connection = MagicMock(name='connection') + self.mock_db_engine.connect.return_value \ + .__enter__.return_value = mock_connection + mock_connection.execute.return_value.fetchall.return_value = [mock_result] self.mock_records_unload_plan.processing_instructions.fail_if_dont_understand = True self.mock_records_unload_plan.processing_instructions.fail_if_cant_handle_hint = True