@@ -60,6 +60,7 @@ void getDatabaseObjects_returnsTables() throws SQLException {
6060 .thenReturn (false ) // No functions
6161 .thenReturn (false ); // No procedures
6262
63+ when (resultSet .getString ("schema_name" )).thenReturn ("public" );
6364 when (resultSet .getString ("name" )).thenReturn ("users" );
6465 when (resultSet .getString ("type" )).thenReturn ("table" );
6566 when (resultSet .getObject ("row_count" )).thenReturn (100L );
@@ -70,6 +71,7 @@ void getDatabaseObjects_returnsTables() throws SQLException {
7071 assertEquals (1 , objects .size ());
7172 assertEquals ("users" , objects .get (0 ).getName ());
7273 assertEquals ("table" , objects .get (0 ).getType ());
74+ assertEquals ("public" , objects .get (0 ).getSchema ());
7375 }
7476
7577 @ Test
@@ -216,6 +218,7 @@ void scanSchema_returnsSchemaMetadata() throws SQLException {
216218 .thenReturn (true )
217219 .thenReturn (false );
218220
221+ when (resultSet .getString ("schemaname" )).thenReturn ("public" );
219222 when (resultSet .getString ("tablename" )).thenReturn ("users" );
220223 when (resultSet .getString ("type" )).thenReturn ("table" );
221224 when (resultSet .getObject ("row_count" )).thenReturn (100L );
@@ -228,6 +231,7 @@ void scanSchema_returnsSchemaMetadata() throws SQLException {
228231 assertNotNull (schema .getTables ());
229232 assertEquals (1 , schema .getTables ().size ());
230233 assertEquals ("users" , schema .getTables ().get (0 ).getName ());
234+ assertEquals ("public" , schema .getTables ().get (0 ).getSchema ());
231235 }
232236
233237 @ Test
@@ -240,8 +244,10 @@ void getForeignKeys_returnsRelationships() throws SQLException {
240244 .thenReturn (false );
241245
242246 when (resultSet .getString ("constraint_name" )).thenReturn ("fk_orders_user" );
247+ when (resultSet .getString ("source_schema" )).thenReturn ("public" );
243248 when (resultSet .getString ("source_table" )).thenReturn ("orders" );
244249 when (resultSet .getString ("source_column" )).thenReturn ("user_id" );
250+ when (resultSet .getString ("target_schema" )).thenReturn ("public" );
245251 when (resultSet .getString ("target_table" )).thenReturn ("users" );
246252 when (resultSet .getString ("target_column" )).thenReturn ("id" );
247253
@@ -250,6 +256,7 @@ void getForeignKeys_returnsRelationships() throws SQLException {
250256 assertNotNull (relationships );
251257 assertEquals (1 , relationships .size ());
252258 assertEquals ("fk_orders_user" , relationships .get (0 ).getConstraintName ());
259+ // qualifyForConsumers returns just the table name for "public" schema
253260 assertEquals ("orders" , relationships .get (0 ).getFromTable ());
254261 assertEquals ("user_id" , relationships .get (0 ).getFromColumn ());
255262 assertEquals ("users" , relationships .get (0 ).getToTable ());
@@ -294,6 +301,7 @@ void scanSchema_fallsBackToExactCountWhenEstimateMissing() throws SQLException {
294301 when (foreignKeysStatement .executeQuery (anyString ())).thenReturn (foreignKeysResultSet );
295302
296303 when (resultSet .next ()).thenReturn (true ).thenReturn (false );
304+ when (resultSet .getString ("schemaname" )).thenReturn ("public" );
297305 when (resultSet .getString ("tablename" )).thenReturn ("dim_route" );
298306 when (resultSet .getString ("type" )).thenReturn ("table" );
299307 when (resultSet .getObject ("row_count" )).thenReturn (null );
@@ -370,7 +378,7 @@ void nonSystemSchemaPredicate_excludesPgCatalog() {
370378
371379 @ Test
372380 void getTablesAndViews_queryExcludesExtensionViews () throws SQLException {
373- // Capture the SQL query and verify it contains the extension view exclusion
381+ // Capture all SQL queries and find the tables/views one
374382 ArgumentCaptor <String > sqlCaptor = ArgumentCaptor .forClass (String .class );
375383 when (connection .createStatement ()).thenReturn (statement );
376384 when (statement .executeQuery (sqlCaptor .capture ())).thenReturn (resultSet );
@@ -379,20 +387,25 @@ void getTablesAndViews_queryExcludesExtensionViews() throws SQLException {
379387 // getDatabaseObjects calls getTablesAndViews internally
380388 List <DatabaseObject > objects = provider .getDatabaseObjects (connection , "test_db" );
381389
382- // Verify the query was captured
383- String executedQuery = sqlCaptor .getValue ();
384- assertNotNull (executedQuery );
390+ // Find the tables/views query among all executed queries
391+ String tablesViewsQuery = sqlCaptor .getAllValues ().stream ()
392+ .filter (q -> q .contains ("pg_tables" ) || q .contains ("pg_views" ))
393+ .findFirst ()
394+ .orElse (null );
395+
396+ assertNotNull (tablesViewsQuery , "Should have executed a query accessing pg_tables or pg_views" );
385397
386398 // The query for views should contain the extension exclusion
387- // (tables + views are in one UNION query, so check the whole thing)
388- assertTrue (executedQuery .contains ("pg_views" ) || executedQuery .contains ("pg_tables" ),
389- "Query should access PostgreSQL system catalogs" );
399+ assertTrue (tablesViewsQuery .contains ("pg_views" ),
400+ "Query should access pg_views" );
401+ assertTrue (tablesViewsQuery .contains ("pg_tables" ),
402+ "Query should access pg_tables" );
390403
391404 // Verify schema exclusions are present
392- assertTrue (executedQuery .contains ("NOT IN" ),
405+ assertTrue (tablesViewsQuery .contains ("NOT IN" ),
393406 "Query should have NOT IN clause for exclusions" );
394- assertTrue (executedQuery . contains ( "pg_catalog" ) || executedQuery .contains ("pg_stat_statements" ),
395- "Query should exclude system schemas or extension views" );
407+ assertTrue (tablesViewsQuery .contains ("pg_stat_statements" ),
408+ "Query should exclude pg_stat_statements extension views" );
396409 }
397410
398411 @ Test
0 commit comments