Merge ~vultaire/charm-graylog:upgrade-tests-and-reworks-20200724 into charm-graylog:master

[Paul Goins]

[Paul Goins (vultaire)]

Reworked my pending test code, pulling all of the non-pipeline-related tests together, and rebasing onto a more recent branch.

I have not run tests yet on this, hence still marked as WIP.

[Paul Goins (vultaire)]

Unit and functional tests should all now be passing. I'm doing one more clean run of functional tests, after which I'll mark this as ready for review.

[Paul Goins (vultaire)]

Tests are confirmed passing. Ready for review.

[Drew Freiberger (afreiberger)]

Much of this code appears to have landed since this MR. can you please review for bits added in this that may be missing from master? pathing has changed during 20.08 release cycle and will require major rebase and rework.

[Drew Freiberger (afreiberger)]

review: Needs Fixing

[Paul Goins (vultaire)]

Abandoning in favor of pipelines-2021.

review: Disapprove

Unmerged commits

5714aec... by Paul Goins on 2020-07-27

Oops; rename problem. Fixed.

9663997... by Paul Goins on 2020-07-27

Fixed unit test breakage

Yes, I could have got fancy and tried to avoid the netifaces dependency,
but there seemed little point; the library is small and I'd rather keep
the code looking clean without doing tricks such as performing imports
in the functions which use the libraries.

3fb2191... by Paul Goins on 2020-07-27

Minor changes to test_upgrade.py

79f3a4a... by Paul Goins on 2020-07-27

Added another prereq

2b1aebd... by Paul Goins on 2020-07-27

Removed some test helper bits in favor of installing libs for testing

fc1d113... by Paul Goins on 2020-07-27

Adding testing prereqs

2795245... by Paul Goins on 2020-02-29

Fixed mistake in reworked is_v2() helper

ea64132... by Paul Goins on 2020-02-29

Updates to shared test code:

* is_v2() should be based upon the config of the graylog unit. This
  is just in case a bug installs the wrong version of Graylog - if we
  test based upon the installed version, we'll miss bugs, whereas if
  we test based upon the config, that is less likely.

* I also tried to clean things up to remove the hardcoded 'graylog'
  and 'graylog/0' dependencies.

504b1f4... by Paul Goins on 2020-02-14

Fixed test_graylog_charm.py, updated shared code

* Set @staticmethod/@classmethod decorators on base class methods
  to enable use during setUpClass().

* Updated test_graylog_charm.py to leverage the shared helpers to
  simplify setup and replace API constant use with get_api_port().

336e271... by Paul Goins on 2020-02-14

Extracted new shared test logic module

* Extracted BaseTest superclass for sharing several helper functions.

* DEFAULT_API_PORT and DEFAULT_WEB_PORT also moved here.

diff --git a/lib/charms/layer/graylog/api.py b/lib/charms/layer/graylog/api.py
new file mode 100644
index 0000000..825a70e
--- /dev/null
+++ b/lib/charms/layer/graylog/api.py
@@ -0,0 +1,231 @@
+<<<<<<< lib/charms/layer/graylog/api.py
+=======
+#
+#    "             "
+#  mmm   m   m   mmm   m   m
+#    #   #   #     #   #   #
+#    #   #   #     #   #   #
+#    #   "mm"#     #   "mm"#
+#    #             #
+#  ""            ""
+# This file is managed by Juju. Local changes will be lost.
+import json
+import os
+import requests
+
+from charmhelpers.core import hookenv
+from charmhelpers.core.hookenv import log
+
+
+def get_ignore_indexer_failures_file():
+    return "/usr/local/lib/nagios/plugins/ignore_indexer_failures.timestamp"
+
+
+class GraylogApi:
+
+    def __init__(self, base_url, username, password, token_name='graylog-api'):
+        self.base_url = base_url
+        if not base_url.endswith('/'):
+            self.base_url += '/'
+        self.username = username
+        self.password = password
+        self.auth = (self.username, self.password)
+        self.token_name = token_name
+        self.token = None
+        self.input_types = None
+        self.req_timeout = 3
+        self.req_retries = 4
+
+    def request(self, path, method='GET', data={}, params=None):
+        if path.startswith('/'):
+            path = path[1:]
+        url = '{}{}?pretty=true'.format(self.base_url, path)
+        headers = {
+            'Accept': 'application/json',
+            'Content-Type': 'application/json',
+            'Accept-Encoding': '',
+            'X-Requested-By': 'Graylog Charm'
+        }
+        if data:
+            data = json.dumps(data, indent=True)
+        tries = 0
+        while tries < self.req_retries:
+            tries += 1
+            try:
+                resp = requests.api.request(method, url, auth=self.auth,
+                                            data=data, params=params,
+                                            headers=headers, timeout=self.req_timeout)
+                if resp.ok:
+                    if method == 'DELETE':
+                        return True
+                    return json.loads(resp.content.decode('utf-8'))
+                else:
+                    msg = 'API error code: {}'.format(resp.status_code)
+                    log(msg)
+                    hookenv.status_set('blocked', msg)
+            except Exception as ex:
+                msg = 'Error calling graylog api: {}'.format(ex)
+                log(msg)
+
+    def token_get(self, token_name=None, halt=False):
+        if self.token:
+            return self.token
+        if not token_name:
+            token_name = self.token_name
+        if self.password == 'token':
+            self.token = self.username
+            self.auth = (self.token, 'token')
+            return self.token
+        url = 'users/{}/tokens'.format(self.username)
+        resp = self.request(url)
+        if not resp:
+            return
+        for token in resp['tokens']:
+            if token['name'] == token_name:
+                self.token = token['token']
+                self.auth = (token['token'], 'token')
+                return self.token
+
+        # None found, let's create a new
+        url = '{}/{}'.format(url, token_name)
+        resp = self.request(url, method='POST')
+        if not halt:
+            return self.token_get(token_name=token_name, halt=True)
+        else:
+            hookenv.status_set('blocked', 'Cannot get API token.')
+            return None  # Don't loop until we run out of memory.
+
+    def index_set_get(self, index_id=None):
+        if not self.token:
+            self.token_get()
+        url = 'system/indices/index_sets'
+        if index_id:
+            return self.request('{}/{}'.format(url, index_id))
+        resp = self.request(url)
+        if resp:
+            return resp.get('index_sets', None)
+
+    def index_set_update(self, index_id, data):
+        if not self.token:
+            self.token_get()
+        url = 'system/indices/index_sets/{}'.format(index_id)
+        return self.request(url, method='PUT', data=data)
+
+    def log_input_get(self, input_id=None):
+        if not self.token:
+            self.token_get()
+        url = 'system/inputs'
+        if input_id:
+            return self.request('{}/{}'.format(url, input_id))
+        resp = self.request(url)
+        if resp:
+            return resp.get('inputs', None)
+
+    def log_input_update(self, input_id=None, data=None):
+        if not self.token:
+            self.token_get()
+        url = 'system/inputs/'
+        if input_id:
+            method = 'PUT'
+            url += input_id
+        else:
+            method = 'POST'
+        return self.request(url, method=method, data=data)
+
+    def log_input_remove(self, input_id):
+        if not self.token:
+            self.token_get()
+        url = 'system/inputs/{}'.format(input_id)
+        return self.request(url, method='DELETE')
+
+    def map_log_input_type(self, type):
+        """
+        Map specified int types to Gray's input type objects. Examples:
+
+          "Beats" -> org.graylog.plugins.beats.BeatsInput
+          "Raw TCP" -> org.graylog2.inputs.raw.tcp.RawTCPInput
+          "Syslog UDP" -> org.graylog2.inputs.syslog.udp.SyslogUDPInput
+        """
+        if not self.input_types:
+            url = 'system/inputs/types'
+            result = self.request(url)
+            if result is not None:
+                self.input_types = result['types']
+            else:
+                return None
+        for k, v in self.input_types.items():
+            if v.lower().find(type.lower()) != -1:
+                return k
+            # For raw input types, map something like "Raw TCP" and "Raw UDP"
+            # instead of requiring the exact name, "Raw/Plaintext TCP".
+            if k.lower().find('.'.join(type.split()).lower()) != -1:
+                return k
+        return None
+
+    def cluster_get(self):
+        if not self.token:
+            self.token_get()
+        url = 'cluster'
+        return self.request(url)
+
+    def user_get(self, username):
+        if not self.token:
+            self.token_get()
+        url = 'users/{}'.format(username)
+        return self.request(url)
+
+    def user_create(self, username, password, read_only=True):
+        if not self.token:
+            self.token_get()
+        user = self.user_get(username)
+        if user:
+            return user
+        url = 'users/'
+        d = {
+            'username': username,
+            'password': password,
+            'full_name': username,
+            'email': '{}@local'.format(username),
+            'permissions': [],
+        }
+        self.request(url, method='POST', data=d)
+        return self.user_get(username)
+
+    def user_permissions_clear(self, username):
+        if not self.token:
+            self.token_get()
+        user = self.user_get(username)
+        if not user:
+            return
+        url = 'users/{}/permissions'.format(username)
+        self.request(url, method='DELETE')
+
+    def user_permissions_set(self, username, permission):
+        if not self.token:
+            self.token_get()
+        user = self.user_get(username)
+        if not user:
+            return
+        url = 'users/{}/permissions'.format(username)
+        self.request(url, method='PUT', data={'permissions': permission})
+
+    def indexer_cluster_health(self):
+        if not self.token:
+            self.token_get()
+        return self.request('system/indexer/cluster/health')
+
+    def indexer_failures_count(self, since):
+        if not self.token:
+            self.token_get()
+        return self.request('system/indexer/failures/count', params={'since': since})
+
+    def journal_state(self):
+        if not self.token:
+            self.token_get()
+        return self.request('system/journal')
+
+    def notifications_get(self):
+        if not self.token:
+            self.token_get()
+        return self.request('system/notifications')
+>>>>>>> lib/charms/layer/graylog/api.py
diff --git a/lib/charms/layer/graylog/utils.py b/lib/charms/layer/graylog/utils.py
index 7e9021e..29ec527 100644
--- a/lib/charms/layer/graylog/utils.py
+++ b/lib/charms/layer/graylog/utils.py
@@ -5,7 +5,6 @@ from urllib.parse import urlparse
 import netifaces
 
 from charmhelpers.core import hookenv
-from charms.layer import snap
 from .version_check import is_v2
 
 
diff --git a/lib/charms/layer/graylog/version_check.py b/lib/charms/layer/graylog/version_check.py
index 23aafd5..9ee5adc 100644
--- a/lib/charms/layer/graylog/version_check.py
+++ b/lib/charms/layer/graylog/version_check.py
@@ -1,13 +1,7 @@
 import os
 
-if os.environ.get('JUJU_UNIT_NAME'):
-    charm = True
-else:
-    charm = False
-
-if charm:
-    from charmhelpers.core import hookenv
-    from charms.layer import snap
+from charmhelpers.core import hookenv
+from charms.layer import snap
 
 
 SNAP_NAME = 'graylog'
diff --git a/src/tests/functional/tests/__init__.py b/src/tests/functional/tests/__init__.py
index 351a8a2..32b2652 100644
--- a/src/tests/functional/tests/__init__.py
+++ b/src/tests/functional/tests/__init__.py
@@ -11,4 +11,4 @@ def _add_path(path):
 
 built_charm = pathlib.Path(os.environ["CHARM_BUILD_DIR"]) / "graylog"
 # Need to avoid importing c.l.graylog as it wants charmhelpers etc.
-_add_path(str(built_charm / "lib/charms/layer/graylog"))
+_add_path(str(built_charm / "lib"))
diff --git a/src/tests/functional/tests/bundles/bionic-graylog2.yaml b/src/tests/functional/tests/bundles/bionic-graylog2.yaml
index bd256db..b1318f0 100644
--- a/src/tests/functional/tests/bundles/bionic-graylog2.yaml
+++ b/src/tests/functional/tests/bundles/bionic-graylog2.yaml
@@ -30,6 +30,19 @@ applications:
   nrpe:
     charm: cs:nrpe
 
+  old-graylog:
+    charm: cs:graylog-32
+    num_units: 1
+    series: bionic
+
+  elastic2:
+    charm: cs:elasticsearch
+    num_units: 1
+
+  mongo2:
+    charm: cs:mongodb
+    num_units: 1
+
 relations:
   - - ubuntu
     - filebeat
@@ -43,3 +56,7 @@ relations:
     - haproxy
   - - graylog
     - nrpe
+  - - old-graylog
+    - mongo2
+  - - old-graylog
+    - elastic2
diff --git a/src/tests/functional/tests/bundles/bionic-graylog3.yaml b/src/tests/functional/tests/bundles/bionic-graylog3.yaml
index 596cf7d..1f6020a 100644
--- a/src/tests/functional/tests/bundles/bionic-graylog3.yaml
+++ b/src/tests/functional/tests/bundles/bionic-graylog3.yaml
@@ -30,6 +30,19 @@ applications:
   nrpe:
     charm: cs:nrpe
 
+  old-graylog:
+    charm: cs:graylog-32
+    num_units: 1
+    series: bionic
+
+  elastic2:
+    charm: cs:elasticsearch
+    num_units: 1
+
+  mongo2:
+    charm: cs:mongodb
+    num_units: 1
+
 relations:
   - - ubuntu
     - filebeat
@@ -43,3 +56,7 @@ relations:
     - haproxy
   - - graylog
     - nrpe
+  - - old-graylog
+    - mongo2
+  - - old-graylog
+    - elastic2
diff --git a/src/tests/unit/requirements.txt b/src/tests/unit/requirements.txt
index c465d62..d39b4a3 100644
--- a/src/tests/unit/requirements.txt
+++ b/src/tests/unit/requirements.txt
@@ -5,3 +5,4 @@ netifaces
 pytest
 pytest-cov
 requests
+netifaces
diff --git a/tests/requirements.txt b/tests/requirements.txt
new file mode 100644
index 0000000..27331d7
--- /dev/null
+++ b/tests/requirements.txt
@@ -0,0 +1,7 @@
+<<<<<<< tests/requirements.txt
+=======
+git+https://github.com/openstack-charmers/zaza.git#egg=zaza
+charmhelpers
+charms.reactive
+netifaces
+>>>>>>> tests/requirements.txt
diff --git a/tests/shared.py b/tests/shared.py
new file mode 100644
index 0000000..9ea5c5a
--- /dev/null
+++ b/tests/shared.py
@@ -0,0 +1,40 @@
+from zaza import model
+from zaza.utilities import juju
+
+from charms.layer.graylog.api import GraylogApi
+
+
+DEFAULT_API_PORT = "9001"  # Graylog 2 only
+DEFAULT_WEB_PORT = "9000"
+
+
+class BaseTest:
+
+    app_name = 'graylog'
+
+    @classmethod
+    def is_v2(cls, app_name=app_name):
+        config = model.get_application_config(app_name)
+        return config['channel']['value'].startswith('2/')
+
+    @classmethod
+    def get_api_port(cls):
+        return DEFAULT_API_PORT if cls.is_v2() else DEFAULT_WEB_PORT
+
+    @classmethod
+    def get_graylog_clients(cls, app_name=app_name):
+        graylog_units = juju.get_full_juju_status().applications[app_name]["units"]
+        graylog_addrs = [unit["public-address"] for unit in graylog_units.values()]
+        port = cls.get_api_port()
+        password = cls.get_graylog_admin_password()
+        return [
+            GraylogApi("http://{}:{}/api".format(graylog_addr, port), "admin", password)
+            for graylog_addr in graylog_addrs
+        ]
+
+    @classmethod
+    def get_graylog_admin_password(cls, unit=None):
+        if unit is None:
+            unit = '{}/0'.format(cls.app_name)
+        result = model.run_action(unit, "show-admin-password")
+        return result.data["results"]["admin-password"]
diff --git a/tests/test_graylog_charm.py b/tests/test_graylog_charm.py
new file mode 100644
index 0000000..ccd8fed
--- /dev/null
+++ b/tests/test_graylog_charm.py
@@ -0,0 +1,117 @@
+<<<<<<< tests/test_graylog_charm.py
+=======
+"""Encapsulate graylog testing."""
+
+import logging
+import time
+import unittest
+
+import zaza.model as model
+
+from tests.shared import BaseTest
+
+
+CURL_TIMEOUT = 180
+REQ_TIMEOUT = 12
+
+
+class BaseGraylogTest(unittest.TestCase, BaseTest):
+    """Base for Graylog charm tests."""
+
+    @classmethod
+    def setUpClass(cls):
+        super(BaseGraylogTest, cls).setUpClass()
+        cls.model_name = model.get_juju_model()
+        cls.application_name = "graylog"
+        cls.lead_unit_name = model.get_lead_unit_name(
+            cls.application_name, model_name=cls.model_name
+        )
+        cls.units = model.get_units(
+            cls.application_name, model_name=cls.model_name
+        )
+        logging.debug("Leader unit is {}".format(cls.lead_unit_name))
+        cls.api = cls.get_graylog_clients()[0]
+        cls.api.req_timeout = (
+            REQ_TIMEOUT
+        )  # try harder to get an api connection to cater for overloaded testsystems
+        logging.debug("API at {}".format(cls.api.base_url))
+
+        # account for the differences between graylog2 and 3
+        if cls.is_v2():
+            cls.file_field = "file"
+            cls.expected_plugins = {
+                "org.graylog.plugins.pipelineprocessor.ProcessorPlugin",
+                "org.graylog.plugins.beats.BeatsInputPlugin",
+                "org.graylog.plugins.collector.CollectorPlugin",
+
+            }
+        else:  # graylog 3+
+            cls.file_field = "filebeat_log_file_path"
+            cls.expected_plugins = {
+                "org.graylog.aws.AWSPlugin",
+                "org.graylog.plugins.collector.CollectorPlugin",
+                "org.graylog.plugins.threatintel.ThreatIntelPlugin"
+            }
+
+class CharmOperationTest(BaseGraylogTest):
+    """Verify operations"""
+
+    def test_05_api_ready(self):
+        """Test if the API is ready
+
+        Curl the api endpoint on the sentried graylog unit. We'll retry
+        until the CURL_TIMEOUT because it may take a few seconds for the
+        graylog systemd service to start.
+        """
+        curl_command = "curl http://localhost:{}".format(self.get_api_port())
+        timeout = time.time() + CURL_TIMEOUT
+        while time.time() < timeout:
+            response = model.run_on_unit(self.lead_unit_name, curl_command)
+            if response["Code"] == "0":
+                return
+            logging.warning(
+                "Unexpected curl response: {}. Retrying in 30s.".format(
+                    response
+                )
+            )
+            time.sleep(30)
+
+        # we didn't get rc=0 in the alloted time, fail the test
+        self.fail(
+            "Graylog didn't respond to the command \n"
+            "'{curl_command}' as expected.\n"
+            "Result: {result}".format(
+                curl_command=curl_command, result=response
+            )
+        )
+
+    def test_06_elasticsearch_healthy(self):
+        resp = self.api.indexer_cluster_health()
+        self.assertEqual(resp["status"], "green")
+
+    def test_07_elasticsearch_no_faults(self):
+        resp = self.api.indexer_failures_count("2020-01-01")
+        self.assertEqual(int(resp["count"]), 0)
+
+    def test_08_plugins(self):
+        resp = self.api.request("/system/plugins")
+        plugin_set = set([p["unique_id"] for p in resp["plugins"]])
+
+        self.assertTrue(self.expected_plugins <= plugin_set)
+
+    def test_09_cluster(self):
+        resp = self.api.cluster_get()
+        self.assertGreater(len(resp.values()), 0)
+        node = list(resp.values()).pop()
+        self.assertTrue(node["is_processing"])
+        self.assertEqual(node["lifecycle"], "running")
+
+    def test_10_keyword_search(self):
+        params = {
+            "query": '{}: "/var/log/cloud-init-output.log" AND ci-info'.format(self.file_field),
+            "keyword": "4 hours ago",
+        }
+
+        resp = self.api.request("/search/universal/keyword", params=params)
+        self.assertTrue(int(resp["total_results"]) > 0)
+>>>>>>> tests/test_graylog_charm.py
diff --git a/tests/test_legacy.py b/tests/test_legacy.py
new file mode 100644
index 0000000..03294fb
--- /dev/null
+++ b/tests/test_legacy.py
@@ -0,0 +1,83 @@
+<<<<<<< tests/test_legacy.py
+=======
+import re
+import unittest
+import yaml
+from zaza.utilities import juju
+from zaza import model
+
+from tests.shared import BaseTest, DEFAULT_API_PORT, DEFAULT_WEB_PORT
+
+
+class LegacyTests(unittest.TestCase, BaseTest):
+
+    # These tests were ported from tests/test_10_basic.py in changeset a3b54c2.
+    # They were temporarily deleted during the conversion from Amulet to Zaza.
+    # Aside from porting to work with libjuju and some refactoring, these tests
+    # are largely as they were before.
+
+    def test_api_ready(self):
+        """Curl the api endpoint on the graylog unit."""
+        port = self.get_api_port()
+        curl_command = "curl http://localhost:{} --connect-timeout 30".format(port)
+        self.run_command("graylog/0", curl_command)
+
+    def test_elasticsearch_active(self):
+        self._assert_remote_address_in_server_conf("elastic/0", "client")
+        # This was also stuck in the same test...
+        api = self.get_graylog_clients()[0]
+        resp = api.indexer_cluster_health()
+        self.assertTrue(resp["status"] == "green")
+
+    def test_mongodb_active(self):
+        self._assert_remote_address_in_server_conf("mongo/0", "database")
+        # This was also stuck in the same test...
+        api = self.get_graylog_clients()[0]
+        resp = api.cluster_get()
+        self.assertTrue(resp[list(resp)[0]]["is_processing"])
+
+    def _assert_remote_address_in_server_conf(self, remote_unit, remote_relation_name):
+        # This is a very simplistic, naive check, but it's equivalent to
+        # that from tests/test_10_basic.py in a3b54c2.
+        source_unit = "graylog/0"
+        remote_addr = juju.get_relation_from_unit(
+            source_unit, remote_unit, remote_relation_name
+        )["private-address"]
+        assert remote_addr in self.get_file_contents(
+            source_unit, "/var/snap/graylog/common/server.conf"
+        )
+
+    def test_website_active(self):
+        data = juju.get_relation_from_unit("haproxy/0", "graylog/0", "website")
+        self.assertEqual(data["port"], DEFAULT_WEB_PORT)
+
+        service_ports = {}
+        for service in yaml.safe_load(data["all_services"]):
+            service_ports[service["service_name"]] = service["service_port"]
+
+        self.assertEqual(str(service_ports["web"]), DEFAULT_WEB_PORT)
+        if self.is_v2():
+            self.assertEqual(str(service_ports["api"]), DEFAULT_API_PORT)
+
+    def run_command(self, target, command, timeout=30):
+        result = model.run_on_unit(target, command, timeout=timeout)
+        assert result["Code"] == "0"
+        return result
+
+    def test_nrpe_config(self):
+        cfg = self.get_file_contents(
+            "nrpe/0", "/etc/nagios/nrpe.d/check_graylog_health.cfg"
+        )
+        self.assertTrue(re.search(r"command.*check_graylog_health", cfg))
+
+    def test_nrpe_health_check(self):
+        cfg = self.get_file_contents(
+            "nrpe/0", "/usr/local/lib/nagios/plugins/check_graylog_health.py"
+        )
+        self.assertTrue(re.search(r"CRITICAL", cfg))
+
+    def get_file_contents(self, unit, filename):
+        result = self.run_command(unit, "cat '{}'".format(filename))
+        assert result["Code"] == "0"
+        return result["Stdout"]
+>>>>>>> tests/test_legacy.py
diff --git a/tests/test_upgrade.py b/tests/test_upgrade.py
new file mode 100644
index 0000000..49b173a
--- /dev/null
+++ b/tests/test_upgrade.py
@@ -0,0 +1,65 @@
+import os
+import pathlib
+import subprocess
+import unittest
+import time
+from zaza.charm_lifecycle import utils as zlc_utils
+from zaza.utilities import juju
+from zaza import model
+
+DEFAULT_API_PORT = "9001"  # Graylog 2 only
+DEFAULT_WEB_PORT = "9000"
+
+
+built_charm = pathlib.Path(os.environ["CHARM_BUILD_DIR"]) / "graylog"
+
+
+class UpgradeTests(unittest.TestCase):
+
+    """Test for charm upgrading.
+
+    Intent: if the charm is upgraded from an old version (especially from a
+    GL2-only version to a GL2+3-supporting version), does it indeed still work?
+
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        # 2019-09-26: libjuju doesn't yet support charm upgrades with
+        # --switch <localdir> or --path <localdir>, so we need to fall back to
+        # the command line.
+        model_name = model.get_juju_model()
+        subprocess.check_call([
+            'juju', 'upgrade-charm',
+            '--model', model_name,
+            'old-graylog', '--path', built_charm])
+
+        # Ensure we wait enough time for the upgrade to have started.
+        time.sleep(60)
+
+        test_config = zlc_utils.get_charm_config()
+        model.wait_for_application_states(
+            model_name,
+            test_config.get('target_deploy_status', {}))
+
+    def test_api_ready(self):
+        """Curl the api endpoint on the graylog unit."""
+        port = self.get_api_port()
+        curl_command = "curl http://localhost:{} --connect-timeout 30".format(port)
+        self.run_command("graylog/0", curl_command)
+
+    def get_api_port(self):
+        port = DEFAULT_API_PORT if self.is_v2() else DEFAULT_WEB_PORT
+        return port
+
+    def run_command(self, target, command, timeout=30):
+        result = model.run_on_unit(target, command, timeout=timeout)
+        assert result["Code"] == "0"
+        return result
+
+    def is_v2(self):
+        return (
+            juju.get_full_juju_status()
+            .applications["graylog"]["workload-version"]
+            .startswith("2.")
+        )
diff --git a/tests/tests.yaml b/tests/tests.yaml
new file mode 100644
index 0000000..517d340
--- /dev/null
+++ b/tests/tests.yaml
@@ -0,0 +1,26 @@
+<<<<<<< tests/tests.yaml
+=======
+charm_name: graylog
+gate_bundles:
+  - bionic-graylog2
+  - bionic-graylog3
+  - focal-graylog2
+  - focal-graylog3
+  - bionic-graylog2-ha
+  - bionic-graylog3-ha
+smoke_bundles:
+  - bionic-graylog3
+dev_bundles:
+  - bionic-graylog3
+tests:
+  - tests.test_legacy.LegacyTests
+  - tests.test_graylog_charm.CharmOperationTest
+  - tests.test_upgrade.UpgradeTests
+target_deploy_status:
+  filebeat:
+    workload-status: active
+    workload-status-message: 'Filebeat ready.'
+  haproxy:
+    workload-status: unknown
+    workload-status-message: ""
+>>>>>>> tests/tests.yaml