From 2e61a9d55230e3f7ea31c34702db30ef27bf031a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Villemain?= <Cedric.Villemain@Data-Bene.io>
Date: Fri, 7 Mar 2025 12:05:48 +0100
Subject: [PATCH] Update test to use citus_is_coordinator()

with a wrapper to make it a stable function and ensure part of the query
is executed on coordinator.

Thanks @onurctirtir for the test/tip.
---
 .../regress/expected/multi_modifications.out  | 41 ++++++++-----------
 src/test/regress/sql/multi_modifications.sql  | 34 +++++++--------
 2 files changed, 31 insertions(+), 44 deletions(-)

diff --git a/src/test/regress/expected/multi_modifications.out b/src/test/regress/expected/multi_modifications.out
index 15b84a1de..93f6c8c45 100644
--- a/src/test/regress/expected/multi_modifications.out
+++ b/src/test/regress/expected/multi_modifications.out
@@ -814,40 +814,31 @@ SELECT * FROM app_analytics_events ORDER BY id;
 DROP TABLE app_analytics_events;
 -- test function call in UPDATE SET
 -- https://github.com/citusdata/citus/issues/7676
-CREATE TABLE test_ref_multiexpr (
-    id bigint primary key
-  , col_int integer
-  , col_bool bool
-  , col_text text
-  , col_timestamp timestamp
+CREATE FUNCTION citus_is_coordinator_stable() returns bool as $$
+  select citus_is_coordinator();
+$$ language sql stable;
+CREATE TABLE bool_test (
+  id bigint primary key,
+  col_bool bool
   );
-SELECT create_reference_table('test_ref_multiexpr');
+SELECT create_reference_table('bool_test');
  create_reference_table
 ---------------------------------------------------------------------
 
 (1 row)
 
-/* TODO how to ensure in test that 'now()' is correctly pre-executed */
-INSERT INTO test_ref_multiexpr VALUES (1, 1, true, 'one', now());
-UPDATE test_ref_multiexpr
-SET (col_timestamp)
-  = (SELECT now())
-RETURNING id, col_int, col_bool;
- id | col_int | col_bool
+INSERT INTO bool_test values (1, true);
+UPDATE bool_test
+SET (col_bool)
+  = (SELECT citus_is_coordinator_stable())
+RETURNING id, col_bool;
+ id | col_bool
 ---------------------------------------------------------------------
-  1 |       1 | t
+  1 | t
 (1 row)
 
-UPDATE test_ref_multiexpr
-SET (col_bool, col_timestamp)
-  = (SELECT true, now())
-RETURNING id, col_int, col_bool;
- id | col_int | col_bool
----------------------------------------------------------------------
-  1 |       1 | t
-(1 row)
-
-DROP TABLE test_ref_multiexpr;
+DROP TABLE bool_test;
+DROP FUNCTION citus_is_coordinator_stable();
 -- Test multi-row insert with serial in a non-partition column
 CREATE TABLE app_analytics_events (id int, app_id serial, name text);
 SELECT create_distributed_table('app_analytics_events', 'id');
diff --git a/src/test/regress/sql/multi_modifications.sql b/src/test/regress/sql/multi_modifications.sql
index 25282c996..2a00e7992 100644
--- a/src/test/regress/sql/multi_modifications.sql
+++ b/src/test/regress/sql/multi_modifications.sql
@@ -507,29 +507,25 @@ DROP TABLE app_analytics_events;
 
 -- test function call in UPDATE SET
 -- https://github.com/citusdata/citus/issues/7676
-CREATE TABLE test_ref_multiexpr (
-    id bigint primary key
-  , col_int integer
-  , col_bool bool
-  , col_text text
-  , col_timestamp timestamp
+CREATE FUNCTION citus_is_coordinator_stable() returns bool as $$
+  select citus_is_coordinator();
+$$ language sql stable;
+
+CREATE TABLE bool_test (
+  id bigint primary key,
+  col_bool bool
   );
-SELECT create_reference_table('test_ref_multiexpr');
+SELECT create_reference_table('bool_test');
 
-/* TODO how to ensure in test that 'now()' is correctly pre-executed */
-INSERT INTO test_ref_multiexpr VALUES (1, 1, true, 'one', now());
+INSERT INTO bool_test values (1, true);
 
-UPDATE test_ref_multiexpr
-SET (col_timestamp)
-  = (SELECT now())
-RETURNING id, col_int, col_bool;
+UPDATE bool_test
+SET (col_bool)
+  = (SELECT citus_is_coordinator_stable())
+RETURNING id, col_bool;
 
-UPDATE test_ref_multiexpr
-SET (col_bool, col_timestamp)
-  = (SELECT true, now())
-RETURNING id, col_int, col_bool;
-
-DROP TABLE test_ref_multiexpr;
+DROP TABLE bool_test;
+DROP FUNCTION citus_is_coordinator_stable();
 
 -- Test multi-row insert with serial in a non-partition column
 CREATE TABLE app_analytics_events (id int, app_id serial, name text);