New postgres test

NREL · May 23, 2024 · 6f0fa77 · 6f0fa77
1 parent 36696bf
commit 6f0fa77
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/elm/wizard.py b/elm/wizard.py
@@ -61,8 +61,6 @@ def query_vector_db(self, query, limit=100):
             ranked strings/scores outputs.
         """
 
-    def test(self, question):
-        return question
     def engineer_query(self, query, token_budget=None, new_info_threshold=0.7,
                        convo=False):
         """Engineer a query for GPT using the corpus of information

diff --git a/tests/test_postgres.py b/tests/test_postgres.py
@@ -35,15 +35,16 @@ def query_call(*args, **kwargs):  # pylint: disable=unused-argument
 
 def test_ref_list(mocker):
     """Test to ensure correct response vector db."""
-    wizard_mock = mocker.patch('elm.wizard.EnergyWizardPostgres', autospec=True)
+    wizard_mock = mocker.patch('elm.wizard.EnergyWizardPostgres',
+                               autospec=True)
 
     wizard = wizard_mock.return_value
     wizard.messages = []
     wizard.MODEL_INSTRUCTION = "Model instruction dummy."
     wizard.token_budget = 500
     wizard.model = "dummy-model-name"
     wizard.count_tokens = mocker.Mock(return_value=50)
-    wizard.make_ref_list.side_effect= MockClass.ref_call
+    wizard.make_ref_list.side_effect = MockClass.ref_call
     wizard.query_vector_db.side_effect = MockClass.query_call
 
     question = "What is a dummy question?"