From afbab1658180c8e3c87041adc25ed24d466d35d0 Mon Sep 17 00:00:00 2001
From: Will Fitch <will.fitch@canonical.com>
Date: Fri, 16 Dec 2022 14:18:52 +0000
Subject: [PATCH] updated service test (#43)

## Proposal
Minor cleanup to pgbouncer start-service test.

## Commits
* updated service test

* lint

* fixed unit import
---
 src/pgbouncer@.service               |  1 +
 tests/integration/helpers/helpers.py | 10 ++++------
 tests/integration/test_charm.py      |  4 ++--
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/src/pgbouncer@.service b/src/pgbouncer@.service
index 6c2b443d5..3891e0649 100644
--- a/src/pgbouncer@.service
+++ b/src/pgbouncer@.service
@@ -10,6 +10,7 @@ User=postgres
 ExecStart=/sbin/pgbouncer -R /var/lib/postgresql/pgbouncer/instance_%i/pgbouncer.ini
 ExecReload=/bin/kill -HUP $MAINPID
 KillSignal=SIGINT
+PIDFile=/var/lib/postgresql/pgbouncer/instance_%i/pgbouncer.pid
 Restart=on-failure
 RestartSec=5s
 
diff --git a/tests/integration/helpers/helpers.py b/tests/integration/helpers/helpers.py
index fe03bd21e..21957c471 100644
--- a/tests/integration/helpers/helpers.py
+++ b/tests/integration/helpers/helpers.py
@@ -11,6 +11,7 @@
 
 import yaml
 from charms.pgbouncer_k8s.v0 import pgb
+from juju.unit import Unit
 from pytest_operator.plugin import OpsTest
 from tenacity import RetryError, Retrying, stop_after_delay, wait_fixed
 
@@ -72,7 +73,7 @@ async def get_unit_cores(unit: str) -> int:
         raise Exception(get_cores_from_unit.results)
 
 
-async def get_running_instances(unit: str, service: str) -> int:
+async def get_running_instances(unit: Unit, service: str) -> int:
     """Returns the number of running instances of the given service.
 
     Uses `ps` to find the number of instances of a given service.
@@ -84,11 +85,8 @@ async def get_running_instances(unit: str, service: str) -> int:
     Returns:
         an integer defining the number of running instances.
     """
-    get_running_instances = await unit.run(f"ps aux | grep {service}")
-    ps_output = get_running_instances.results.get("Stdout")
-    num_of_ps_lines = len(ps_output.split("\n"))
-    # one extra for grep process, and one for a blank line at the end
-    return num_of_ps_lines - 2
+    get_running_instances = await unit.run(f"pgrep -c {service}")
+    return int(get_running_instances.results.get("Stdout"))
 
 
 async def get_unit_info(ops_test: OpsTest, unit_name: str) -> Dict:
diff --git a/tests/integration/test_charm.py b/tests/integration/test_charm.py
index 903afb117..4928f8aa6 100644
--- a/tests/integration/test_charm.py
+++ b/tests/integration/test_charm.py
@@ -82,9 +82,9 @@ async def test_systemd_restarts_pgbouncer_processes(ops_test: OpsTest):
     assert await helpers.get_running_instances(unit, "pgbouncer") == expected_processes
 
     # Kill pgbouncer process and wait for it to restart
-    await unit.run("kill $(ps aux | grep pgbouncer | awk '{print $2}')")
+    await unit.run("kill $(pgrep pgbouncer)")
     async with ops_test.fast_forward():
-        await ops_test.model.wait_for_idle(apps=[PGB], status="blocked", timeout=300)
+        await ops_test.model.wait_for_idle(apps=[PGB], status="blocked", timeout=(3 * 60))
     assert ops_test.model.units[f"{PGB}/0"].workload_status_message == WAIT_MSG
 
     # verify all processes start again