From bd290582adcce018ea8b156e63c1bf91baca5f18 Mon Sep 17 00:00:00 2001
From: Lucas Gameiro <lucas.borges@canonical.com>
Date: Fri, 1 Nov 2024 10:09:02 +0100
Subject: [PATCH] [DPE-5689] Fix connection rejection rule in `pg_hba.conf`
 (#751)

* update check

* add unit test

* complement test
---
 src/backups.py             |  6 ++++--
 tests/unit/test_backups.py | 26 ++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/backups.py b/src/backups.py
index d7e31cd442..12f32798c1 100644
--- a/src/backups.py
+++ b/src/backups.py
@@ -806,10 +806,12 @@ def _on_create_backup_action(self, event) -> None:  # noqa: C901
             event.fail(error_message)
             return
 
+        disabled_connectivity = False
         if not self.charm.is_primary:
             # Create a rule to mark the cluster as in a creating backup state and update
             # the Patroni configuration.
             self._change_connectivity_to_database(connectivity=False)
+            disabled_connectivity = True
 
         self.charm.unit.status = MaintenanceStatus("creating backup")
         # Set flag due to missing in progress backups on JSON output
@@ -879,8 +881,8 @@ def _on_create_backup_action(self, event) -> None:  # noqa: C901
                 logger.info(f"Backup succeeded: with backup-id {datetime_backup_requested}")
                 event.set_results({"backup-status": "backup created"})
 
-        if not self.charm.is_primary:
-            # Remove the rule the marks the cluster as in a creating backup state
+        if disabled_connectivity:
+            # Remove the rule that marks the cluster as in a creating backup state
             # and update the Patroni configuration.
             self._change_connectivity_to_database(connectivity=True)
 
diff --git a/tests/unit/test_backups.py b/tests/unit/test_backups.py
index 28f9a70a0f..3945335046 100644
--- a/tests/unit/test_backups.py
+++ b/tests/unit/test_backups.py
@@ -1332,6 +1332,32 @@ def test_on_create_backup_action(harness):
         mock_event.fail.assert_not_called()
         mock_event.set_results.assert_called_once_with({"backup-status": "backup created"})
 
+        # Test when this unit is a replica but gets promoted to primary mid-way
+        mock_event.reset_mock()
+        mock_event.params = {"type": "full"}
+        _upload_content_to_s3.reset_mock()
+        _is_primary.return_value = (False, True)
+        harness.charm.backup._on_create_backup_action(mock_event)
+        _upload_content_to_s3.assert_has_calls([
+            call(
+                expected_metadata,
+                f"backup/{harness.charm.model.name}.{harness.charm.cluster_name}/latest",
+                mock_s3_parameters,
+            ),
+            call(
+                "Stdout:\nfake stdout\n\nStderr:\nfake stderr\n",
+                f"backup/{harness.charm.model.name}.{harness.charm.cluster_name}/2023-01-01T09:00:00Z/backup.log",
+                mock_s3_parameters,
+            ),
+        ])
+        assert _change_connectivity_to_database.call_count == 2
+        _change_connectivity_to_database.assert_has_calls([
+            call(connectivity=False),
+            call(connectivity=True),
+        ])
+        mock_event.fail.assert_not_called()
+        mock_event.set_results.assert_called_once_with({"backup-status": "backup created"})
+
 
 def test_on_list_backups_action(harness):
     with (