diff --git a/src/test/regress/expected/citus_local_tables_mx.out b/src/test/regress/expected/citus_local_tables_mx.out
index 3ed66c36d..17edd2132 100644
--- a/src/test/regress/expected/citus_local_tables_mx.out
+++ b/src/test/regress/expected/citus_local_tables_mx.out
@@ -628,7 +628,7 @@ ERROR:  cannot build foreign key between reference table and a partition
 -- this should work
 alter table citus_local_parent_1 add constraint fkey_to_drop_test foreign key(a) references ref(a);
 -- this should undistribute the table, and the entries should be gone from pg_dist_partition
-select logicalrelid from pg_dist_partition where logicalrelid::text like 'citus_local_parent%';
+select logicalrelid from pg_dist_partition where logicalrelid::text like 'citus_local_parent%' order by logicalrelid;
          logicalrelid
 ---------------------------------------------------------------------
  citus_local_parent_1
@@ -640,6 +640,12 @@ select logicalrelid from pg_dist_partition where logicalrelid::text like 'citus_
 
 set client_min_messages to error;
 alter table citus_local_parent_1 drop constraint fkey_to_drop_test;
+SELECT undistribute_table('citus_local_parent_1', true);
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
 select logicalrelid from pg_dist_partition where logicalrelid::text like 'citus_local_parent%';
  logicalrelid
 ---------------------------------------------------------------------
diff --git a/src/test/regress/sql/citus_local_tables_mx.sql b/src/test/regress/sql/citus_local_tables_mx.sql
index d74e18da6..73cd0ee9e 100644
--- a/src/test/regress/sql/citus_local_tables_mx.sql
+++ b/src/test/regress/sql/citus_local_tables_mx.sql
@@ -337,9 +337,10 @@ alter table citus_local_parent_1_child_1 add foreign key(a) references ref(a);
 -- this should work
 alter table citus_local_parent_1 add constraint fkey_to_drop_test foreign key(a) references ref(a);
 -- this should undistribute the table, and the entries should be gone from pg_dist_partition
-select logicalrelid from pg_dist_partition where logicalrelid::text like 'citus_local_parent%';
+select logicalrelid from pg_dist_partition where logicalrelid::text like 'citus_local_parent%' order by logicalrelid;
 set client_min_messages to error;
 alter table citus_local_parent_1 drop constraint fkey_to_drop_test;
+SELECT undistribute_table('citus_local_parent_1', true);
 select logicalrelid from pg_dist_partition where logicalrelid::text like 'citus_local_parent%';
 
 -- verify attaching partition with a foreign key errors out