Add unit tests

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

tests/unit/test_charm.py

@@ -164,7 +164,9 @@ def test_patroni_scrape_config_tls(harness):
 
 
 def test_primary_endpoint(harness):
-    with patch(
+    with patch("charm.stop_after_delay", new_callable=PropertyMock) as _stop_after_delay, patch(
+        "charm.wait_fixed", new_callable=PropertyMock
+    ) as _wait_fixed, patch(
         "charm.PostgresqlOperatorCharm._units_ips",
         new_callable=PropertyMock,
         return_value={"1.1.1.1", "1.1.1.2"},
@@ -173,6 +175,10 @@ def test_primary_endpoint(harness):
         _patroni.return_value.get_primary.return_value = sentinel.primary
         assert harness.charm.primary_endpoint == "1.1.1.1"
 
+        # Check needed to ensure a fast charm deployment.
+        _stop_after_delay.assert_called_once_with(5)
+        _wait_fixed.assert_called_once_with(3)
+
         _patroni.return_value.get_member_ip.assert_called_once_with(sentinel.primary)
         _patroni.return_value.get_primary.assert_called_once_with()
 

tests/unit/test_cluster.py

@@ -42,6 +42,7 @@ def json(self):
         "http://server1/cluster": {
             "members": [{"name": "postgresql-0", "host": "1.1.1.1", "role": "leader", "lag": "1"}]
         },
+        "http://server1/health": {"state": "running"},
         "http://server4/cluster": {"members": []},
     }
     if args[0] in data:
@@ -128,6 +129,28 @@ def test_get_member_ip(peers_ips, patroni):
         tc.assertIsNone(ip)
 
 
+def test_get_patroni_health(peers_ips, patroni):
+    with patch("cluster.stop_after_delay", new_callable=PropertyMock) as _stop_after_delay, patch(
+        "cluster.wait_fixed", new_callable=PropertyMock
+    ) as _wait_fixed, patch(
+        "charm.Patroni._patroni_url", new_callable=PropertyMock
+    ) as _patroni_url, patch("requests.get", side_effect=mocked_requests_get) as _get:
+        # Test when the Patroni API is reachable.
+        _patroni_url.return_value = "http://server1"
+        health = patroni.get_patroni_health()
+
+        # Check needed to ensure a fast charm deployment.
+        _stop_after_delay.assert_called_once_with(60)
+        _wait_fixed.assert_called_once_with(7)
+
+        tc.assertEqual(health, {"state": "running"})
+
+        # Test when the Patroni API is not reachable.
+        _patroni_url.return_value = "http://server2"
+        with tc.assertRaises(tenacity.RetryError):
+            patroni.get_patroni_health()
+
+
 def test_get_postgresql_version(peers_ips, patroni):
     with patch("charm.snap.SnapClient") as _snap_client:
         # TODO test a real implementation

tests/unit/test_upgrade.py

@@ -60,6 +60,41 @@ def test_log_rollback(harness):
         )
 
 
+@pytest.mark.parametrize(
+    "unit_states,is_cluster_initialised,call",
+    [
+        (["ready"], False, False),
+        (["ready", "ready"], True, False),
+        (["idle"], False, False),
+        (["idle"], True, False),
+        (["ready"], True, True),
+    ],
+)
+def test_on_upgrade_charm_check_legacy(harness, unit_states, is_cluster_initialised, call):
+    with (
+        patch(
+            "charms.data_platform_libs.v0.upgrade.DataUpgrade.state",
+            new_callable=PropertyMock(return_value=None),
+        ) as _state,
+        patch(
+            "charms.data_platform_libs.v0.upgrade.DataUpgrade.unit_states",
+            new_callable=PropertyMock(return_value=unit_states),
+        ) as _unit_states,
+        patch(
+            "charm.PostgresqlOperatorCharm.is_cluster_initialised",
+            new_callable=PropertyMock(return_value=is_cluster_initialised),
+        ) as _is_cluster_initialised,
+        patch("charm.Patroni.member_started", new_callable=PropertyMock) as _member_started,
+        patch(
+            "upgrade.PostgreSQLUpgrade._prepare_upgrade_from_legacy"
+        ) as _prepare_upgrade_from_legacy,
+    ):
+        with harness.hooks_disabled():
+            harness.set_leader(True)
+        harness.charm.upgrade._on_upgrade_charm_check_legacy()
+        _member_started.assert_called_once() if call else _member_started.assert_not_called()
+
+
 @patch_network_get(private_address="1.1.1.1")
 def test_on_upgrade_granted(harness):
     with (