rework test_wizard_postgres

NREL · Sep 13, 2024 · a690afe · a690afe
1 parent feecfb1
commit a690afe
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/examples/research_hub/retrieve_docs.py b/examples/research_hub/retrieve_docs.py
@@ -17,7 +17,7 @@
 
 
 logger = logging.getLogger(__name__)
-init_logger(__name__, log_level='DEBUG')
+init_logger(__name__, log_level='INFO')
 init_logger('elm', log_level='INFO')
 
 

diff --git a/tests/test_wizard_postgres.py b/tests/test_wizard_postgres.py
@@ -63,14 +63,21 @@ def invoke_model(self, **kwargs):  # pylint: disable=unused-argument
         return dummy_response
 
 
-def test_postgres():
+def test_postgres(mocker):
     """Test to ensure correct response vector db."""
 
     os.environ["EWIZ_DB_USER"] = "user"
 
+    mock_conn_cm = mocker.MagicMock()
+    mock_conn = mock_conn_cm.__enter__.return_value
+    mock_conn.cursor.return_value = Cursor()
+
+    mock_connect = mocker.patch('psycopg2.connect')
+    mock_connect.return_value = mock_conn_cm
+
     wizard = EnergyWizardPostgres(db_host='Dummy', db_port='Dummy',
                                   db_name='Dummy', db_schema='Dummy',
-                                  db_table='Dummy', cursor=Cursor(),
+                                  db_table='Dummy',
                                   boto_client=BotoClient())
 
     question = 'Is this a dummy question?'