From 4d02556806b0aa0b3781dfb2bcea8c67f42f3479 Mon Sep 17 00:00:00 2001
From: Ian Costanzo <ian@anon-solutions.ca>
Date: Tue, 17 Dec 2024 08:04:19 -0800
Subject: [PATCH] Test cleanup

Signed-off-by: Ian Costanzo <ian@anon-solutions.ca>
---
 .../simple_restart/docker-compose.yml         |  5 ---
 scenarios/examples/simple_restart/example.py  | 40 +++++++++++--------
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/scenarios/examples/simple_restart/docker-compose.yml b/scenarios/examples/simple_restart/docker-compose.yml
index 6bf3c92208..7185096b05 100644
--- a/scenarios/examples/simple_restart/docker-compose.yml
+++ b/scenarios/examples/simple_restart/docker-compose.yml
@@ -6,8 +6,6 @@ services:
       - POSTGRES_PASSWORD=DB_PASSWORD
     ports:
       - 5433:5432
-    volumes:
-      - wallet-db-data:/var/lib/pgsql/data
     healthcheck:
       test: ["CMD-SHELL", "pg_isready -U DB_USER"]
       interval: 10s
@@ -124,6 +122,3 @@ services:
         condition: service_healthy
       bob:
         condition: service_healthy
-
-volumes:
-  wallet-db-data:
diff --git a/scenarios/examples/simple_restart/example.py b/scenarios/examples/simple_restart/example.py
index cb954db0cc..558d94285b 100644
--- a/scenarios/examples/simple_restart/example.py
+++ b/scenarios/examples/simple_restart/example.py
@@ -33,14 +33,12 @@
 def healthy(container: Container) -> bool:
     """Check if container is healthy."""
     inspect_results = container.attrs
-    print(">>> state:", inspect_results["State"])
     return inspect_results["State"]["Running"] and inspect_results["State"]["Health"]["Status"] == "healthy"
 
 
 def unhealthy(container: Container) -> bool:
     """Check if container is unhealthy."""
     inspect_results = container.attrs
-    print(">>> state:", inspect_results["State"])
     return not inspect_results["State"]["Running"]
 
 
@@ -133,21 +131,29 @@ async def main():
     print(">>> new container:", 'alice', json.dumps(new_alice_container.attrs))
     alice_id = new_alice_container.attrs['Id']
 
-    wait_until_healthy(client, alice_id)
-    print(">>> new alice container is healthy")
-
-    # run some more tests ...  alice should still be connected to bob for example ...
-    async with Controller(base_url=ALICE) as alice, Controller(base_url=BOB) as bob:
-        # Present the the credential's attributes
-        print(">>> present proof ... again ...")
-        await indy_present_proof_v2(
-            bob,
-            alice,
-            bob_conn.connection_id,
-            alice_conn.connection_id,
-            requested_attributes=[{"name": "firstname"}],
-        )
-        print(">>> Done! (again)")
+    try:
+        wait_until_healthy(client, alice_id)
+        print(">>> new alice container is healthy")
+
+        # run some more tests ...  alice should still be connected to bob for example ...
+        async with Controller(base_url=ALICE) as alice, Controller(base_url=BOB) as bob:
+            # Present the the credential's attributes
+            print(">>> present proof ... again ...")
+            await indy_present_proof_v2(
+                bob,
+                alice,
+                bob_conn.connection_id,
+                alice_conn.connection_id,
+                requested_attributes=[{"name": "firstname"}],
+            )
+            print(">>> Done! (again)")
+    finally:
+        # cleanup - shut down alice agent (not part of docker compose)
+        print(">>> shut down alice ...")
+        alice_container = client.containers.get(alice_id)
+        alice_container.stop()
+        wait_until_healthy(client, alice_id, is_healthy=False)
+        alice_container.remove()
 
 
 if __name__ == "__main__":