Fix citus tests. Metadata sync could be slow after coordinator switch

features/citus.feature

@@ -10,20 +10,20 @@ Feature: citus
     And I start postgres3 in citus group 1
     Then replication works from postgres0 to postgres1 after 15 seconds
     Then replication works from postgres2 to postgres3 after 15 seconds
-    And postgres0 is registered in the postgres0 as the worker in group 0
-    And postgres2 is registered in the postgres0 as the worker in group 1
+    And postgres0 is registered in the postgres0 as the worker in group 0 after 5 seconds
+    And postgres2 is registered in the postgres0 as the worker in group 1 after 5 seconds
 
   Scenario: coordinator failover updates pg_dist_node
     Given I run patronictl.py failover batman --group 0 --candidate postgres1 --force
     Then postgres1 role is the primary after 10 seconds
     And replication works from postgres1 to postgres0 after 15 seconds
+    And postgres1 is registered in the postgres2 as the worker in group 0 after 5 seconds
     And "sync" key in a group 0 in DCS has sync_standby=postgres0 after 15 seconds
-    And postgres1 is registered in the postgres2 as the worker in group 0
     When I run patronictl.py failover batman --group 0 --candidate postgres0 --force
     Then postgres0 role is the primary after 10 seconds
     And replication works from postgres0 to postgres1 after 15 seconds
+    And postgres0 is registered in the postgres2 as the worker in group 0 after 5 seconds
     And "sync" key in a group 0 in DCS has sync_standby=postgres1 after 15 seconds
-    And postgres0 is registered in the postgres2 as the worker in group 0
 
   Scenario: worker switchover doesn't break client queries on the coordinator
     Given I create a distributed table on postgres0
@@ -32,15 +32,15 @@ Feature: citus
     Then I receive a response returncode 0
     And postgres3 role is the primary after 10 seconds
     And replication works from postgres3 to postgres2 after 15 seconds
+    And postgres3 is registered in the postgres0 as the worker in group 1 after 5 seconds
     And "sync" key in a group 1 in DCS has sync_standby=postgres2 after 15 seconds
-    And postgres3 is registered in the postgres0 as the worker in group 1
     And a thread is still alive
     When I run patronictl.py switchover batman --group 1 --force
     Then I receive a response returncode 0
     And postgres2 role is the primary after 10 seconds
     And replication works from postgres2 to postgres3 after 15 seconds
+    And postgres2 is registered in the postgres0 as the worker in group 1 after 5 seconds
     And "sync" key in a group 1 in DCS has sync_standby=postgres3 after 15 seconds
-    And postgres2 is registered in the postgres0 as the worker in group 1
     And a thread is still alive
     When I stop a thread
     Then a distributed table on postgres0 has expected rows
@@ -52,7 +52,7 @@ Feature: citus
     Then I receive a response returncode 0
     And postgres2 role is the primary after 10 seconds
     And replication works from postgres2 to postgres3 after 15 seconds
-    And postgres2 is registered in the postgres0 as the worker in group 1
+    And postgres2 is registered in the postgres0 as the worker in group 1 after 5 seconds
     And a thread is still alive
     When I stop a thread
     Then a distributed table on postgres0 has expected rows
@@ -65,7 +65,7 @@ Feature: citus
     Then I receive a response returncode 0
     And I receive a response output "+ttl: 20"
     When I sleep for 2 seconds
-    Then postgres4 is registered in the postgres2 as the worker in group 2
+    Then postgres4 is registered in the postgres2 as the worker in group 2 after 5 seconds
     When I shut down postgres4
     Then There is a transaction in progress on postgres0 changing pg_dist_node
     When I run patronictl.py restart batman postgres2 --group 1 --force

features/steps/citus.py

@@ -44,12 +44,24 @@ def start_citus(context, name, group):
     return context.pctl.start(name, custom_config={"citus": {"database": "postgres", "group": int(group)}})
 
 
-@step('{name1:w} is registered in the {name2:w} as the worker in group {group:d}')
-def check_registration(context, name1, name2, group):
+@step('{name1:w} is registered in the {name2:w} as the worker in group {group:d} after {time_limit:d} seconds')
+def check_registration(context, name1, name2, group, time_limit):
+    time_limit *= context.timeout_multiplier
+    max_time = time.time() + int(time_limit)
+    pg_dist_value = None
     worker_port = int(context.pctl.query(name1, "SHOW port").fetchone()[0])
-    r = context.pctl.query(name2, "SELECT nodeport FROM pg_catalog.pg_dist_node WHERE groupid = {0}".format(group))
-    assert worker_port == r.fetchone()[0],\
-        "Worker {0} is not registered in pg_dist_node on the coordinator {1}".format(name1, name2)
+    while time.time() < max_time:
+        try:
+            r = context.pctl.query(name2,
+                                   "SELECT nodeport FROM pg_catalog.pg_dist_node WHERE groupid = {0}".format(group))
+            pg_dist_value = r.fetchone()[0]
+            if pg_dist_value == worker_port:
+                return
+        except Exception:
+            pass
+        time.sleep(1)
+    assert False, "Worker {0} is not registered in pg_dist_node on the coordinator {1} after {2} seconds"\
+        .format(name1, name2, time_limit)
 
 
 @step('I create a distributed table on {name:w}')