From 1290fb4e72eadcb71b97a877b558d05180a306b2 Mon Sep 17 00:00:00 2001
From: Kian Parvin <kian.parvin@canonical.com>
Date: Fri, 7 Jul 2023 09:45:50 +0200
Subject: [PATCH] JIMM k8s charm fixes

---
 charms/jimm-k8s/src/charm.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/charms/jimm-k8s/src/charm.py b/charms/jimm-k8s/src/charm.py
index b42a766fc..aa3526cfb 100755
--- a/charms/jimm-k8s/src/charm.py
+++ b/charms/jimm-k8s/src/charm.py
@@ -66,6 +66,7 @@
     "CANDID_URL",
 ]
 
+JIMM_SERVICE_NAME = "jimm"
 DATABASE_NAME = "jimm"
 LOG_FILE = "/var/log/jimm"
 # This likely will just be JIMM's port.
@@ -258,7 +259,7 @@ def _update_workload(self, event):
             "summary": "jimm layer",
             "description": "pebble config layer for jimm",
             "services": {
-                "jimm": {
+                JIMM_SERVICE_NAME: {
                     "override": "merge",
                     "summary": "JAAS Intelligent Model Manager",
                     "command": "/root/jimmsrv",
@@ -276,10 +277,10 @@ def _update_workload(self, event):
         }
         container.add_layer("jimm", pebble_layer, combine=True)
         if self._ready():
-            if container.get_service("jimm").is_running():
+            if container.get_service(JIMM_SERVICE_NAME).is_running():
                 container.replan()
             else:
-                container.start("jimm")
+                container.start(JIMM_SERVICE_NAME)
             self.unit.status = ActiveStatus("running")
         else:
             logger.info("workload container not ready - defering")
@@ -304,7 +305,7 @@ def _on_stop(self, _):
         """Stop JIMM."""
         container = self.unit.get_container(WORKLOAD_CONTAINER)
         if container.can_connect():
-            container.stop()
+            container.stop(JIMM_SERVICE_NAME)
         self._ready()
 
     def _on_update_status(self, _):
@@ -357,12 +358,12 @@ def _ready(self):
 
         if container.can_connect():
             plan = container.get_plan()
-            if plan.services.get("jimm") is None:
+            if plan.services.get(JIMM_SERVICE_NAME) is None:
                 logger.error("waiting for service")
                 self.unit.status = WaitingStatus("waiting for service")
                 return False
 
-            env_vars = plan.services.get("jimm").environment
+            env_vars = plan.services.get(JIMM_SERVICE_NAME).environment
 
             for setting in REQUIRED_SETTINGS:
                 if not env_vars.get(setting, ""):
@@ -371,7 +372,7 @@ def _ready(self):
                     )
                     return False
 
-            if container.get_service("jimm").is_running():
+            if container.get_service(JIMM_SERVICE_NAME).is_running():
                 self.unit.status = ActiveStatus("running")
             else:
                 self.unit.status = WaitingStatus("stopped")
@@ -555,7 +556,7 @@ def _get_dns_name(self, event):
         if not self._state.is_ready():
             event.defer()
             logger.warning("State is not ready")
-            return
+            return None
 
         default_dns_name = "{}.{}-endpoints.{}.svc.cluster.local".format(
             self.unit.name.replace("/", "-"),
@@ -601,7 +602,7 @@ def _on_certificate_expiring(self, event: CertificateExpiringEvent) -> None:
 
         new_csr = generate_csr(
             private_key=private_key.encode(),
-            subject=self.dns_name,
+            subject=dns_name,
         )
         self.certificates.request_certificate_renewal(
             old_certificate_signing_request=old_csr,
@@ -609,7 +610,7 @@ def _on_certificate_expiring(self, event: CertificateExpiringEvent) -> None:
         )
         self._state.csr = new_csr.decode()
 
-        self._update_workload()
+        self._update_workload(event)
 
     @requires_state_setter
     def _on_certificate_revoked(self, event: CertificateRevokedEvent) -> None:
@@ -634,7 +635,7 @@ def _on_certificate_revoked(self, event: CertificateRevokedEvent) -> None:
         del self._state.chain
 
         self.unit.status = WaitingStatus("Waiting for new certificate")
-        self._update_workload()
+        self._update_workload(event)
 
     @requires_state_setter
     def _on_ingress_ready(self, event: IngressPerAppReadyEvent):