From 5082317f2f2eeed73b124966fc346424bb60c475 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 19 Nov 2024 17:24:13 +0200
Subject: [PATCH] Update sync nodes during upgrade

---
 src/charm.py               | 4 ++--
 src/upgrade.py             | 1 +
 tests/unit/test_charm.py   | 2 +-
 tests/unit/test_upgrade.py | 5 +++++
 4 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 37fbfb1552..2a2699b834 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -405,7 +405,7 @@ def _on_get_primary(self, event: ActionEvent) -> None:
         except RetryError as e:
             logger.error(f"failed to get primary with error {e}")
 
-    def _updated_synchronous_node_count(self, num_units: int | None = None) -> bool:
+    def updated_synchronous_node_count(self, num_units: int | None = None) -> bool:
         """Tries to update synchronous_node_count configuration and reports the result."""
         try:
             self._patroni.update_synchronous_node_count(num_units)
@@ -439,7 +439,7 @@ def _on_peer_relation_departed(self, event: RelationDepartedEvent) -> None:
 
         if "cluster_initialised" not in self._peers.data[
             self.app
-        ] or not self._updated_synchronous_node_count(len(self._units_ips)):
+        ] or not self.updated_synchronous_node_count(len(self._units_ips)):
             logger.debug("Deferring on_peer_relation_departed: cluster not initialized")
             event.defer()
             return
diff --git a/src/upgrade.py b/src/upgrade.py
index 629ba06fa8..2a20787981 100644
--- a/src/upgrade.py
+++ b/src/upgrade.py
@@ -146,6 +146,7 @@ def _on_upgrade_granted(self, event: UpgradeGrantedEvent) -> None:
         # Update the configuration.
         self.charm.unit.status = MaintenanceStatus("updating configuration")
         self.charm.update_config()
+        self.charm.updated_synchronous_node_count(len(self.charm._units_ips))
 
         self.charm.unit.status = MaintenanceStatus("refreshing the snap")
         self.charm._install_snap_packages(packages=SNAP_PACKAGES, refresh=True)
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 650dbe689d..4bbe46d23c 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -2189,7 +2189,7 @@ def test_on_peer_relation_departed(harness):
         patch("charm.Patroni.are_all_members_ready") as _are_all_members_ready,
         patch("charm.PostgresqlOperatorCharm._get_ips_to_remove") as _get_ips_to_remove,
         patch(
-            "charm.PostgresqlOperatorCharm._updated_synchronous_node_count"
+            "charm.PostgresqlOperatorCharm.updated_synchronous_node_count"
         ) as _updated_synchronous_node_count,
         patch("charm.Patroni.remove_raft_member") as _remove_raft_member,
         patch("charm.PostgresqlOperatorCharm._unit_ip") as _unit_ip,
diff --git a/tests/unit/test_upgrade.py b/tests/unit/test_upgrade.py
index 7dfbfc521b..97c4b5e560 100644
--- a/tests/unit/test_upgrade.py
+++ b/tests/unit/test_upgrade.py
@@ -106,6 +106,9 @@ def test_on_upgrade_granted(harness):
         patch("charm.Patroni.start_patroni") as _start_patroni,
         patch("charm.PostgresqlOperatorCharm._install_snap_packages") as _install_snap_packages,
         patch("charm.PostgresqlOperatorCharm.update_config") as _update_config,
+        patch(
+            "charm.PostgresqlOperatorCharm.updated_synchronous_node_count"
+        ) as _updated_synchronous_node_count,
     ):
         # Test when the charm fails to start Patroni.
         mock_event = MagicMock()
@@ -174,6 +177,7 @@ def test_on_upgrade_granted(harness):
         _member_started.reset_mock()
         _cluster_members.reset_mock()
         mock_event.defer.reset_mock()
+        _updated_synchronous_node_count.reset_mock()
         _is_replication_healthy.return_value = True
         with harness.hooks_disabled():
             harness.set_leader(True)
@@ -184,6 +188,7 @@ def test_on_upgrade_granted(harness):
         _set_unit_completed.assert_called_once()
         _set_unit_failed.assert_not_called()
         _on_upgrade_changed.assert_called_once()
+        _updated_synchronous_node_count.assert_called_once_with(2)
 
 
 def test_pre_upgrade_check(harness):