@@ -114,29 +114,42 @@ def test_natural_language_query_works_without_dbt_project(self):
114114
115115 This should initially FAIL, then PASS after we fix the spoofing.
116116 """
117- config = SQLBotConfig (profile = "TestProfile" )
118-
119- # Simulate what happens in natural language query processing
120- # This typically involves:
121- # 1. Connection validation (runs dbt debug - THIS FAILS)
122- # 2. Schema introspection
123- # 3. LLM generates SQL
124- # 4. SQL execution (this part works)
125-
126- dbt_service = get_dbt_service (config )
127-
128- # Step 1: Connection validation - this is where it currently fails
129- debug_result = dbt_service .debug ()
130- assert debug_result ['success' ], f"dbt debug failed: { debug_result ['error' ]}
131- assert debug_result ['connection_ok' ], "Connection should be OK"
132-
133- # Step 2: If debug passes, the rest should work
134- # Test a query that would come from natural language processing
135- result = dbt_service .execute_query ("SELECT COUNT(*) FROM test_table WHERE category = 'Fruit'" )
136-
137- assert result .success , f"Natural language query execution failed: { result .error }
138- assert len (result .data ) == 1
139- assert result .data [0 ]['COUNT(*)' ] == '3' # Apple, Banana, Mango
117+ # Set the profiles directory to our test location
118+ import os
119+ old_profiles_dir = os .environ .get ('DBT_PROFILES_DIR' )
120+ os .environ ['DBT_PROFILES_DIR' ] = str (Path (self .test_dir ) / ".dbt" )
121+
122+ try :
123+ config = SQLBotConfig (profile = "TestProfile" )
124+
125+ # Simulate what happens in natural language query processing
126+ # This typically involves:
127+ # 1. Connection validation (runs dbt debug - THIS FAILS)
128+ # 2. Schema introspection
129+ # 3. LLM generates SQL
130+ # 4. SQL execution (this part works)
131+
132+ dbt_service = get_dbt_service (config )
133+
134+ # Step 1: Connection validation - this is where it currently fails
135+ debug_result = dbt_service .debug ()
136+ assert debug_result ['success' ], f"dbt debug failed: { debug_result ['error' ]}
137+ assert debug_result ['connection_ok' ], "Connection should be OK"
138+
139+ # Step 2: If debug passes, the rest should work
140+ # Test a query that would come from natural language processing
141+ result = dbt_service .execute_query ("SELECT COUNT(*) FROM test_table WHERE category = 'Fruit'" )
142+
143+ assert result .success , f"Natural language query execution failed: { result .error }
144+ assert len (result .data ) == 1
145+ assert result .data [0 ]['COUNT(*)' ] == '3' # Apple, Banana, Mango
146+
147+ finally :
148+ # Restore original environment
149+ if old_profiles_dir is not None :
150+ os .environ ['DBT_PROFILES_DIR' ] = old_profiles_dir
151+ elif 'DBT_PROFILES_DIR' in os .environ :
152+ del os .environ ['DBT_PROFILES_DIR' ]
140153
141154 def test_generalized_dbt_operations_work_without_physical_project_file (self ):
142155 """
@@ -145,28 +158,41 @@ def test_generalized_dbt_operations_work_without_physical_project_file(self):
145158 This tests the broader issue: ALL dbt operations should work with virtual spoofing,
146159 not just the direct SQL execution path.
147160 """
148- config = SQLBotConfig (profile = "TestProfile" )
149- dbt_service = get_dbt_service (config )
150-
151- # Test various dbt operations that natural language queries might use
152- operations_to_test = [
153- ("debug" , lambda : dbt_service .debug ()),
154- ("list_models" , lambda : dbt_service .list_models ()),
155- # Add more operations as needed
156- ]
157-
158- for operation_name , operation_func in operations_to_test :
159- try :
160- result = operation_func ()
161- # Each operation should succeed or gracefully handle the virtual environment
162- if isinstance (result , dict ):
163- assert not (result .get ('success' ) is False and 'dbt_project.yml' in str (result .get ('error' , '' ))), \
164- f"Operation { operation_name } { result }
165- # For operations that return other types, just ensure they don't crash
166- except Exception as e :
167- if 'dbt_project.yml' in str (e ):
168- pytest .fail (f"Operation { operation_name } { e } )
169- # Other exceptions might be expected (e.g., no models to list)
161+ # Set the profiles directory to our test location
162+ import os
163+ old_profiles_dir = os .environ .get ('DBT_PROFILES_DIR' )
164+ os .environ ['DBT_PROFILES_DIR' ] = str (Path (self .test_dir ) / ".dbt" )
165+
166+ try :
167+ config = SQLBotConfig (profile = "TestProfile" )
168+ dbt_service = get_dbt_service (config )
169+
170+ # Test various dbt operations that natural language queries might use
171+ operations_to_test = [
172+ ("debug" , lambda : dbt_service .debug ()),
173+ ("list_models" , lambda : dbt_service .list_models ()),
174+ # Add more operations as needed
175+ ]
176+
177+ for operation_name , operation_func in operations_to_test :
178+ try :
179+ result = operation_func ()
180+ # Each operation should succeed or gracefully handle the virtual environment
181+ if isinstance (result , dict ):
182+ assert not (result .get ('success' ) is False and 'dbt_project.yml' in str (result .get ('error' , '' ))), \
183+ f"Operation { operation_name } { result }
184+ # For operations that return other types, just ensure they don't crash
185+ except Exception as e :
186+ if 'dbt_project.yml' in str (e ):
187+ pytest .fail (f"Operation { operation_name } { e } )
188+ # Other exceptions might be expected (e.g., no models to list)
189+
190+ finally :
191+ # Restore original environment
192+ if old_profiles_dir is not None :
193+ os .environ ['DBT_PROFILES_DIR' ] = old_profiles_dir
194+ elif 'DBT_PROFILES_DIR' in os .environ :
195+ del os .environ ['DBT_PROFILES_DIR' ]
170196
171197
172198if __name__ == "__main__" :
0 commit comments