From b35247364f7934579bff9585a071ff9668157e9a Mon Sep 17 00:00:00 2001
From: Felipe Reyes <felipe.reyes@canonical.com>
Date: Fri, 17 Jun 2022 16:36:51 -0400
Subject: [PATCH] Set cluster-partition-handling on upgrade-charm.

For units deployed before the implementation of the
cluster-partition-handling strategy they won't have that key set in the
leader making the charm believe there are pending tasks, so this change
seeds the key when is not set with the value present in the charm's
configuration.

Change-Id: Ifdae35ffee1ad7a8f4e5248c817cca14b69d9566
Closes-Bug: #1979092
---
 hooks/rabbitmq_server_relations.py           |  6 +++
 unit_tests/test_rabbitmq_server_relations.py | 45 ++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/hooks/rabbitmq_server_relations.py b/hooks/rabbitmq_server_relations.py
index fd75a0ab..a8a089fe 100755
--- a/hooks/rabbitmq_server_relations.py
+++ b/hooks/rabbitmq_server_relations.py
@@ -819,6 +819,12 @@ def upgrade_charm():
     if missing_packages:
         apt_install(missing_packages, fatal=True)
 
+    if is_leader() and not leader_get(rabbit.CLUSTER_MODE_KEY):
+        log("Setting {} to {} on upgrade-charm.".format(
+            rabbit.CLUSTER_MODE_KEY,
+            config(rabbit.CLUSTER_MODE_KEY)), level='INFO')
+        leader_set({rabbit.CLUSTER_MODE_KEY: config(rabbit.CLUSTER_MODE_KEY)})
+
 
 MAN_PLUGIN = 'rabbitmq_management'
 PROM_PLUGIN = 'rabbitmq_prometheus'
diff --git a/unit_tests/test_rabbitmq_server_relations.py b/unit_tests/test_rabbitmq_server_relations.py
index aa3f59e0..094adcef 100644
--- a/unit_tests/test_rabbitmq_server_relations.py
+++ b/unit_tests/test_rabbitmq_server_relations.py
@@ -479,3 +479,48 @@ class RelationUtil(CharmTestCase):
             'rabbitmq-server')
         wait_app.assert_called_once_with()
         self.assertFalse(update_clients.called)
+
+    @patch.object(rabbitmq_server_relations, 'leader_set')
+    @patch.object(rabbitmq_server_relations, 'leader_get')
+    @patch.object(rabbitmq_server_relations, 'apt_install')
+    @patch.object(rabbitmq_server_relations, 'filter_installed_packages')
+    @patch.object(rabbitmq_server_relations, 'apt_update')
+    @patch.object(rabbitmq_server_relations, 'update_clients')
+    @patch('rabbit_utils.clustered_with_leader')
+    @patch('rabbit_utils.update_peer_cluster_status')
+    @patch('rabbit_utils.migrate_passwords_to_peer_relation')
+    @patch.object(rabbitmq_server_relations, 'is_elected_leader')
+    @patch('os.listdir')
+    @patch('charmhelpers.contrib.hardening.harden.config')
+    @patch.object(rabbitmq_server_relations, 'config')
+    def test_upgrade_charm(self, config, harden_config, listdir,
+                           is_elected_releader,
+                           migrate_passwords_to_peer_relation,
+                           update_peer_cluster_status,
+                           clustered_with_leader,
+                           update_clients,
+                           apt_update,
+                           filter_installed_packages,
+                           apt_install,
+                           leader_get,
+                           leader_set):
+        config.side_effect = self.test_config
+        harden_config.side_effect = self.test_config
+        is_elected_releader.return_value = True
+        clustered_with_leader.return_value = True
+        filter_installed_packages.side_effect = lambda x: x
+
+        leader_get.return_value = None
+
+        rabbitmq_server_relations.upgrade_charm()
+
+        migrate_passwords_to_peer_relation.assert_called()
+        update_peer_cluster_status.assert_called()
+        apt_update.assert_called_with(fatal=True)
+        apt_install.assert_called_with(['python3-amqplib',
+                                        'python3-croniter'],
+                                       fatal=True)
+        leader_set.assert_called_with(
+            {rabbit_utils.CLUSTER_MODE_KEY:
+             self.test_config.get(rabbit_utils.CLUSTER_MODE_KEY)}
+        )