Fix integration tests (#14)

* fix integration tests

Co-authored-by: Marc Oppenheimer <[email protected]>

src/charm.py

@@ -40,9 +40,6 @@ def __init__(self, *args):
         self.framework.observe(
             self.on[ZOOKEEPER_REL_NAME].relation_changed, self._on_kafka_pebble_ready
         )
-        self.framework.observe(
-            self.on[ZOOKEEPER_REL_NAME].relation_departed, self._on_zookeeper_broken
-        )
         self.framework.observe(
             self.on[ZOOKEEPER_REL_NAME].relation_broken, self._on_zookeeper_broken
         )
@@ -163,8 +160,12 @@ def _on_zookeeper_joined(self, event: RelationJoinedEvent) -> None:
         if self.unit.is_leader():
             event.relation.data[self.app].update({"chroot": "/" + self.app.name})
 
-    def _on_zookeeper_broken(self, _: RelationEvent) -> None:
+    def _on_zookeeper_broken(self, event: RelationEvent) -> None:
         """Handler for `zookeeper_relation_departed/broken` events."""
+        if not self.container.can_connect():
+            event.defer()
+            return
+
         logger.info("stopping kafka service")
         self.container.stop(CHARM_KEY)
         self.unit.status = BlockedStatus("missing required zookeeper relation")

tests/integration/helpers.py

@@ -7,14 +7,18 @@
 from typing import Any, List, Tuple
 
 import yaml
+from pytest_operator.plugin import OpsTest
 
 METADATA = yaml.safe_load(Path("./metadata.yaml").read_text())
+KAFKA_CONTAINER = METADATA["resources"]["kafka-image"]["upstream-source"]
 APP_NAME = METADATA["name"]
+ZK_NAME = "zookeeper-k8s"
 
 
 def check_user(model_full_name: str, username: str, zookeeper_uri: str) -> None:
+    container_command = f"KAFKA_OPTS=-Djava.security.auth.login.config=/data/kafka/config/kafka-jaas.cfg ./opt/kafka/bin/kafka-configs.sh --zookeeper {zookeeper_uri} --describe --entity-type users --entity-name {username}"
     result = check_output(
-        f"JUJU_MODEL={model_full_name} juju ssh kafka-k8s/0 'kafka.configs --zookeeper {zookeeper_uri} --describe --entity-type users --entity-name {username}'",
+        f"JUJU_MODEL={model_full_name} juju ssh --container kafka kafka-k8s/0 '{container_command}'",
         stderr=PIPE,
         shell=True,
         universal_newlines=True,
@@ -53,3 +57,22 @@ def get_zookeeper_connection(unit_name: str, model_full_name: str) -> Tuple[List
         return usernames, zookeeper_uri
     else:
         raise Exception("config not found")
+
+
+def check_application_status(ops_test: OpsTest, app_name: str) -> str:
+    """Return the application status for an app name."""
+    model_name = ops_test.model.info.name
+    proc = check_output(f"juju status --model={model_name}".split())
+    proc = proc.decode("utf-8")
+
+    statuses = {"active", "maintenance", "waiting", "blocked"}
+    for line in proc.splitlines():
+        parts = line.split()
+        # first line with app name will have application status
+        if parts and parts[0] == app_name:
+            # NOTE: intersects possible statuses with the list of values:
+            # this is done because sometimes version number exists and
+            # sometimes it doesn't on juju status output
+            find_status = list(statuses & set(parts))
+            if find_status:
+                return find_status[0]

tests/integration/test_charm.py

@@ -6,33 +6,49 @@
 import logging
 
 import pytest
+from helpers import APP_NAME, KAFKA_CONTAINER, ZK_NAME, check_application_status
 from pytest_operator.plugin import OpsTest
 
 logger = logging.getLogger(__name__)
 
-APP_NAME = "kafka"
-ZK = "zookeeper-k8s"
-
 
 @pytest.mark.abort_on_fail
 async def test_build_and_deploy(ops_test: OpsTest):
     kafka_charm = await ops_test.build_charm(".")
     await asyncio.gather(
-        ops_test.model.deploy("zookeeper-k8s", channel="edge", application_name=ZK, num_units=1),
-        ops_test.model.deploy(kafka_charm, application_name=APP_NAME, num_units=1),
+        ops_test.model.deploy(
+            "zookeeper-k8s",
+            channel="edge",
+            application_name=ZK_NAME,
+            num_units=3,
+        ),
+        ops_test.model.deploy(
+            kafka_charm,
+            application_name=APP_NAME,
+            num_units=1,
+            resources={"kafka-image": KAFKA_CONTAINER},
+        ),
     )
-    await ops_test.model.wait_for_idle(apps=[APP_NAME, ZK])
-    assert ops_test.model.applications[APP_NAME].status == "waiting"
-    assert ops_test.model.applications[ZK].status == "active"
+    await ops_test.model.block_until(lambda: len(ops_test.model.applications[ZK_NAME].units) == 3)
+    await ops_test.model.wait_for_idle(apps=[APP_NAME, ZK_NAME], timeout=1000)
+
+    assert check_application_status(ops_test, APP_NAME) == "waiting"
+    assert ops_test.model.applications[ZK_NAME].status == "active"
+
+    await ops_test.model.add_relation(APP_NAME, ZK_NAME)
+
+    async with ops_test.fast_forward():
+        await ops_test.model.wait_for_idle(apps=[APP_NAME, ZK_NAME])
 
-    await ops_test.model.add_relation(APP_NAME, ZK)
-    await ops_test.model.wait_for_idle(apps=[APP_NAME, ZK])
     assert ops_test.model.applications[APP_NAME].status == "active"
-    assert ops_test.model.applications[ZK].status == "active"
+    assert ops_test.model.applications[ZK_NAME].status == "active"
 
 
 @pytest.mark.abort_on_fail
 async def test_blocks_without_zookeeper(ops_test: OpsTest):
-    await asyncio.gather(ops_test.model.applications[ZK].remove())
-    await ops_test.model.wait_for_idle(apps=[APP_NAME])
-    assert ops_test.model.applications[ZK].status == "blocked"
+    async with ops_test.fast_forward():
+        await ops_test.model.applications[ZK_NAME].remove()
+        await ops_test.model.wait_for_idle(apps=[APP_NAME], raise_on_error=False, timeout=1000)
+
+    # Unit is on 'blocked' but whole app is on 'waiting'
+    assert check_application_status(ops_test, APP_NAME) == "waiting"

tests/integration/test_kafka_provider.py

@@ -6,14 +6,17 @@
 import logging
 
 import pytest
+from helpers import (
+    APP_NAME,
+    KAFKA_CONTAINER,
+    ZK_NAME,
+    check_user,
+    get_zookeeper_connection,
+)
 from pytest_operator.plugin import OpsTest
 
-from tests.integration.helpers import check_user, get_zookeeper_connection
-
 logger = logging.getLogger(__name__)
 
-APP_NAME = "kafka"
-ZK = "zookeeper"
 DUMMY_NAME_1 = "app"
 DUMMY_NAME_2 = "appii"
 
@@ -25,27 +28,36 @@ def usernames():
 
 @pytest.mark.abort_on_fail
 async def test_deploy_charms_relate_active(ops_test: OpsTest, usernames):
-    zk_charm = await ops_test.build_charm(".")
+    kafka_charm = await ops_test.build_charm(".")
     app_charm = await ops_test.build_charm("tests/integration/app-charm")
 
     await asyncio.gather(
         ops_test.model.deploy(
-            "zookeeper", channel="edge", application_name="zookeeper", num_units=1
+            "zookeeper-k8s", channel="edge", application_name=ZK_NAME, num_units=3
+        ),
+        ops_test.model.deploy(
+            kafka_charm,
+            application_name=APP_NAME,
+            num_units=1,
+            resources={"kafka-image": KAFKA_CONTAINER},
         ),
-        ops_test.model.deploy(zk_charm, application_name=APP_NAME, num_units=1),
         ops_test.model.deploy(app_charm, application_name=DUMMY_NAME_1, num_units=1),
     )
-    await ops_test.model.wait_for_idle(apps=[APP_NAME, DUMMY_NAME_1, ZK])
-    await ops_test.model.add_relation(APP_NAME, ZK)
-    await ops_test.model.wait_for_idle(apps=[APP_NAME, ZK])
+    await ops_test.model.block_until(lambda: len(ops_test.model.applications[ZK_NAME].units) == 3)
+    await ops_test.model.wait_for_idle(apps=[APP_NAME, DUMMY_NAME_1, ZK_NAME])
+
+    await ops_test.model.add_relation(APP_NAME, ZK_NAME)
+    await ops_test.model.wait_for_idle(apps=[APP_NAME, ZK_NAME])
+
     await ops_test.model.add_relation(APP_NAME, DUMMY_NAME_1)
     await ops_test.model.wait_for_idle(apps=[APP_NAME, DUMMY_NAME_1])
+
     assert ops_test.model.applications[APP_NAME].status == "active"
     assert ops_test.model.applications[DUMMY_NAME_1].status == "active"
 
     # implicitly tests setting of kafka app data
     returned_usernames, zookeeper_uri = get_zookeeper_connection(
-        unit_name="kafka/0", model_full_name=ops_test.model_full_name
+        unit_name="kafka-k8s/0", model_full_name=ops_test.model_full_name
     )
     usernames.update(returned_usernames)
 
@@ -64,12 +76,13 @@ async def test_deploy_multiple_charms_relate_active(ops_test: OpsTest, usernames
     await ops_test.model.wait_for_idle(apps=[DUMMY_NAME_2])
     await ops_test.model.add_relation(APP_NAME, DUMMY_NAME_2)
     await ops_test.model.wait_for_idle(apps=[APP_NAME, DUMMY_NAME_2])
+
     assert ops_test.model.applications[APP_NAME].status == "active"
     assert ops_test.model.applications[DUMMY_NAME_1].status == "active"
     assert ops_test.model.applications[DUMMY_NAME_2].status == "active"
 
     returned_usernames, zookeeper_uri = get_zookeeper_connection(
-        unit_name="kafka/0", model_full_name=ops_test.model_full_name
+        unit_name="kafka-k8s/0", model_full_name=ops_test.model_full_name
     )
     usernames.update(returned_usernames)
 
@@ -84,18 +97,12 @@ async def test_deploy_multiple_charms_relate_active(ops_test: OpsTest, usernames
 @pytest.mark.abort_on_fail
 async def test_remove_application_removes_user(ops_test: OpsTest, usernames):
     await ops_test.model.applications[DUMMY_NAME_1].remove()
+    await ops_test.model.applications[DUMMY_NAME_2].remove()
     await ops_test.model.wait_for_idle(apps=[APP_NAME])
     assert ops_test.model.applications[APP_NAME].status == "active"
 
-    _, zookeeper_uri = get_zookeeper_connection(
-        unit_name="kafka/0", model_full_name=ops_test.model_full_name
-    )
-
     # checks that past usernames no longer exist in ZooKeeper
-    with pytest.raises(AssertionError):
-        for username in usernames:
-            check_user(
-                username=username,
-                zookeeper_uri=zookeeper_uri,
-                model_full_name=ops_test.model_full_name,
-            )
+    with pytest.raises(Exception):
+        _, _ = get_zookeeper_connection(
+            unit_name="kafka-k8s/0", model_full_name=ops_test.model_full_name
+        )