From a2d8e96d8e7e9b58244ba0885826af2ad02c88d1 Mon Sep 17 00:00:00 2001
From: Samuel Walladge <samuel.walladge@canonical.com>
Date: Wed, 24 Aug 2022 22:34:47 +0930
Subject: [PATCH] Fix test to work with vault with raft backend (#827)

The tests here run a series of commands on the vault client,
without waiting in between.
This is usually fine, but with the raft backend,
there can be delays between initialising, unsealing,
and the raft cluster settling.
During this time, most vault commands will return errors.
---
 zaza/openstack/charm_tests/vault/tests.py | 10 ++++
 zaza/openstack/charm_tests/vault/utils.py | 70 +++++++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/zaza/openstack/charm_tests/vault/tests.py b/zaza/openstack/charm_tests/vault/tests.py
index 4ba2e58..c31a06b 100644
--- a/zaza/openstack/charm_tests/vault/tests.py
+++ b/zaza/openstack/charm_tests/vault/tests.py
@@ -63,8 +63,18 @@ class BaseVaultTest(test_utils.OpenStackBaseTest):
         if cls.vip_client:
             cls.clients.append(cls.vip_client)
         cls.vault_creds = vault_utils.get_credentials()
+
+        # This little dance is to ensure a correct init and unseal sequence,
+        # for the case of vault with the raft backend.
+        # It will also work fine in other cases.
+        # The wait functions will raise AssertionErrors on timeouts.
+        init_client = vault_utils.wait_and_get_initialized_client(cls.clients)
+        vault_utils.unseal_all([init_client], cls.vault_creds['keys'][0])
+        vault_utils.wait_until_all_initialised(cls.clients)
         vault_utils.unseal_all(cls.clients, cls.vault_creds['keys'][0])
+
         vault_utils.auth_all(cls.clients, cls.vault_creds['root_token'])
+        vault_utils.wait_for_ha_settled(cls.clients)
         vault_utils.ensure_secret_backend(cls.clients[0])
 
     def tearDown(self):
diff --git a/zaza/openstack/charm_tests/vault/utils.py b/zaza/openstack/charm_tests/vault/utils.py
index c29c745..280eba4 100644
--- a/zaza/openstack/charm_tests/vault/utils.py
+++ b/zaza/openstack/charm_tests/vault/utils.py
@@ -21,6 +21,7 @@ import hvac
 import logging
 import requests
 import tempfile
+import time
 import urllib3
 import yaml
 import tenacity
@@ -280,6 +281,75 @@ def ensure_secret_backend(client):
         pass
 
 
+def wait_for_ha_settled(clients):
+    """Wait until vault ha is settled (for all passed clients).
+
+    Raise an AssertionError if any are not settled within 2 minutes.
+    This function is effectively a no-op for non-ha vault.
+    Requires all vault units to be unsealed.
+
+    :param clients: Clients to use to talk to vault
+    :type clients: List[CharmVaultClient]
+    :raises: AssertionError
+    """
+    for client in clients:
+        for attempt in tenacity.Retrying(
+            reraise=True,
+            wait=tenacity.wait_fixed(10),
+            stop=tenacity.stop_after_attempt(12),  # wait for max 2 minutes
+        ):
+            with attempt:
+                # ha_status could also raise other errors,
+                # eg. if unsealing still in progress.
+                # This is why we're using tenacity here;
+                # avoids needing to manually handle other exceptions.
+                ha_status = client.hvac_client.ha_status
+                if (
+                    not ha_status.get('leader_address') and
+                    ha_status.get('ha_enabled')
+                ):
+                    raise AssertionError('Timeout waiting for ha to settle')
+
+
+def wait_until_all_initialised(clients):
+    """Wait until vault is initialized (for all passed clients).
+
+    Raise an AssertionError if any are not initialized within 2 minutes.
+
+    :param clients: Clients to use to talk to vault
+    :type clients: List[CharmVaultClient]
+    :raises: AssertionError
+    """
+    for client in clients:
+        for _ in range(12):
+            if is_initialized(client):
+                break
+            time.sleep(10)  # max 2 minutes (12 x 10s)
+        else:
+            raise AssertionError("Timeout waiting for vault to initialize")
+
+
+def wait_and_get_initialized_client(clients):
+    """Wait until at least one vault unit is initialized.
+
+    And return the initialized client.
+    Raise an AssertionError
+    if no initialized clients are found within 2 minutes.
+
+    :param clients: Clients to use to talk to vault
+    :type clients: List[CharmVaultClient]
+    :raises: AssertionError
+    :returns: an initialized client
+    :rtype: CharmVaultClient
+    """
+    for _ in range(12):
+        for client in clients:
+            if is_initialized(client):
+                return client
+        time.sleep(10)  # max 2 minutes (12 x 10s)
+    raise AssertionError("Timeout waiting for vault to initialize")
+
+
 def find_unit_with_creds():
     """Find the unit thats has stored the credentials.