Enable Citus behave tests for Postgres v16 (#2914)

and reduce flakiness

.github/workflows/install_deps.py

@@ -45,8 +45,8 @@ def install_packages(what):
     packages['exhibitor'] = packages['zookeeper']
     packages = packages.get(what, [])
     ver = versions.get(what)
-    if float(ver) == 15:
+    if float(ver) >= 15:
-        packages += ['postgresql-{0}-citus-12.0'.format(ver)]
+        packages += ['postgresql-{0}-citus-12.1'.format(ver)]
     subprocess.call(['sudo', 'apt-get', 'update', '-y'])
     return subprocess.call(['sudo', 'apt-get', 'install', '-y', 'postgresql-' + ver, 'expect-dev'] + packages)
 

features/citus.feature

@@ -68,6 +68,6 @@ Feature: citus
     And I receive a response output "+ttl: 20"
     Then postgres4 is registered in the postgres2 as the primary in group 2 after 5 seconds
     When I shut down postgres4
-    Then There is a transaction in progress on postgres0 changing pg_dist_node
+    Then there is a transaction in progress on postgres0 changing pg_dist_node after 5 seconds
     When I run patronictl.py restart batman postgres2 --group 1 --force
     Then a transaction finishes in 20 seconds

features/steps/citus.py

@@ -115,12 +115,18 @@ def count_rows(context, name):
     assert rows == context.insert_counter, "Distributed table doesn't have expected amount of rows"
 
 
-@step("There is a transaction in progress on {name:w} changing pg_dist_node")
+@step("there is a transaction in progress on {name:w} changing pg_dist_node after {time_limit:d} seconds")
-def check_transaction(context, name):
+def check_transaction(context, name, time_limit):
-    cur = context.pctl.query(name, "SELECT xact_start FROM pg_stat_activity WHERE pid <> pg_backend_pid()"
+    time_limit *= context.timeout_multiplier
-                                   " AND state = 'idle in transaction' AND query ~ 'citus_update_node'")
+    max_time = time.time() + int(time_limit)
-    assert cur.rowcount == 1, "There is no idle in transaction updating pg_dist_node"
+    while time.time() < max_time:
-    context.xact_start = cur.fetchone()[0]
+        cur = context.pctl.query(name, "SELECT xact_start FROM pg_stat_activity WHERE pid <> pg_backend_pid()"
+                                       " AND state = 'idle in transaction' AND query ~ 'citus_update_node'")
+        if cur.rowcount == 1:
+            context.xact_start = cur.fetchone()[0]
+            return
+        time.sleep(1)
+    assert False, f"There is no idle in transaction on {name} updating pg_dist_node after {time_limit} seconds"
 
 
 @step("a transaction finishes in {timeout:d} seconds")