diff --git a/src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/11.1-1.sql b/src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/11.1-1.sql
index 134ae1b82..b35373dc9 100644
--- a/src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/11.1-1.sql
+++ b/src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/11.1-1.sql
@@ -2,7 +2,7 @@ CREATE FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint)
     RETURNS VOID
     LANGUAGE C STRICT
     AS 'MODULE_PATHNAME',$$citus_wait_for_rebalance_job$$;
-COMMENT ON FUNCTION pg_catalog.citus_wait_for_rebalance_job()
+COMMENT ON FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint)
     IS 'blocks till the job identified by jobid is at a terminal state, errors if no such job exists';
 
-GRANT EXECUTE ON FUNCTION pg_catalog.citus_wait_for_rebalance_job() TO PUBLIC;
+GRANT EXECUTE ON FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint) TO PUBLIC;
diff --git a/src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/latest.sql b/src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/latest.sql
index 134ae1b82..b35373dc9 100644
--- a/src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/latest.sql
+++ b/src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/latest.sql
@@ -2,7 +2,7 @@ CREATE FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint)
     RETURNS VOID
     LANGUAGE C STRICT
     AS 'MODULE_PATHNAME',$$citus_wait_for_rebalance_job$$;
-COMMENT ON FUNCTION pg_catalog.citus_wait_for_rebalance_job()
+COMMENT ON FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint)
     IS 'blocks till the job identified by jobid is at a terminal state, errors if no such job exists';
 
-GRANT EXECUTE ON FUNCTION pg_catalog.citus_wait_for_rebalance_job() TO PUBLIC;
+GRANT EXECUTE ON FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint) TO PUBLIC;
diff --git a/src/test/regress/expected/rebalance_job.out b/src/test/regress/expected/rebalance_job.out
new file mode 100644
index 000000000..b68d03a1c
--- /dev/null
+++ b/src/test/regress/expected/rebalance_job.out
@@ -0,0 +1,27 @@
+CREATE SCHEMA rebalance_job;
+SET search_path TO rebalance_job;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+SET citus.next_shard_id TO 3536400;
+CREATE TABLE results (a int);
+-- simple job that inserts 1 into results to show that query runs
+SELECT a FROM results WHERE a = 1; -- verify result is not in there
+ a
+---------------------------------------------------------------------
+(0 rows)
+
+INSERT INTO pg_dist_rebalance_jobs (command) VALUES ($job$ INSERT INTO rebalance_job.results VALUES ( 1 ); $job$) RETURNING jobid \gset
+SELECT citus_wait_for_rebalance_job(:jobid); -- wait for the job to be finished
+ citus_wait_for_rebalance_job
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT a FROM results WHERE a = 1; -- verify result is there
+ a
+---------------------------------------------------------------------
+ 1
+(1 row)
+
+SET client_min_messages TO WARNING;
+DROP SCHEMA rebalance_job CASCADE;
diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule
index ab8d092ab..a7cdc9159 100644
--- a/src/test/regress/multi_schedule
+++ b/src/test/regress/multi_schedule
@@ -97,6 +97,7 @@ test: issue_5248 issue_5099
 test: object_propagation_debug
 test: undistribute_table
 test: run_command_on_all_nodes
+test: rebalance_job
 
 
 # ---------
diff --git a/src/test/regress/sql/rebalance_job.sql b/src/test/regress/sql/rebalance_job.sql
new file mode 100644
index 000000000..a55b6eabc
--- /dev/null
+++ b/src/test/regress/sql/rebalance_job.sql
@@ -0,0 +1,16 @@
+CREATE SCHEMA rebalance_job;
+SET search_path TO rebalance_job;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+SET citus.next_shard_id TO 3536400;
+
+CREATE TABLE results (a int);
+
+-- simple job that inserts 1 into results to show that query runs
+SELECT a FROM results WHERE a = 1; -- verify result is not in there
+INSERT INTO pg_dist_rebalance_jobs (command) VALUES ($job$ INSERT INTO rebalance_job.results VALUES ( 1 ); $job$) RETURNING jobid \gset
+SELECT citus_wait_for_rebalance_job(:jobid); -- wait for the job to be finished
+SELECT a FROM results WHERE a = 1; -- verify result is there
+
+SET client_min_messages TO WARNING;
+DROP SCHEMA rebalance_job CASCADE;