Repository: canonical/operator
Pull request #501 (2021-04-06 02:38:24)

Make get_services() return mapping; add ServiceInfo.is_running
Fixes https://github.com/canonical/operator/issues/494 -- see discussion there.


Patch:
diff --git a/ops/model.py b/ops/model.py
index ae9af50..9f9a41f 100644
--- a/ops/model.py
+++ b/ops/model.py
@@ -346,8 +346,8 @@ class Unit:
         self._backend.application_version_set(version)
 
     @property
-    def containers(self) -> typing.Dict[str, 'Container']:
-        """Return a dict of containers indexed by name."""
+    def containers(self) -> 'ContainerMapping':
+        """Return a mapping of containers indexed by name."""
         if not self._is_our_unit:
             raise RuntimeError('cannot get container for a remote unit {}'.format(self))
         return self._containers
@@ -1068,13 +1068,14 @@ class Container:
         """Get the current effective pebble configuration."""
         return self._pebble.get_plan()
 
-    def get_services(self, *service_names: str) -> typing.List['pebble.ServiceInfo']:
-        """Get a list of service status information.
+    def get_services(self, *service_names: str) -> 'ServiceInfoMapping':
+        """Fetch and return a mapping of status information indexed by service name.
 
         If no service names are specified, return status information for all
         services, otherwise return information for only the given services.
         """
-        return self._pebble.get_services(service_names)
+        services = self._pebble.get_services(service_names)
+        return ServiceInfoMapping(services)
 
     def get_service(self, service_name: str) -> 'pebble.ServiceInfo':
         """Get status information for a single named service.
@@ -1086,7 +1087,7 @@ class Container:
             raise ModelError('service {!r} not found'.format(service_name))
         if len(services) > 1:
             raise RuntimeError('expected 1 service, got {}'.format(len(services)))
-        return services[0]
+        return services[service_name]
 
 
 class ContainerMapping(Mapping):
@@ -1112,6 +1113,29 @@ class ContainerMapping(Mapping):
         return repr(self._containers)
 
 
+class ServiceInfoMapping(Mapping):
+    """Map of service names to pebble.ServiceInfo objects.
+
+    This is done as a mapping object rather than a plain dictionary so that we
+    can extend it later, and so it's not mutable.
+    """
+
+    def __init__(self, services: typing.Iterable['pebble.ServiceInfo']):
+        self._services = {s.name: s for s in services}
+
+    def __getitem__(self, key: str):
+        return self._services[key]
+
+    def __iter__(self):
+        return iter(self._services)
+
+    def __len__(self):
+        return len(self._services)
+
+    def __repr__(self):
+        return repr(self._services)
+
+
 class ModelError(Exception):
     """Base class for exceptions raised when interacting with the Model."""
     pass
diff --git a/ops/pebble.py b/ops/pebble.py
index 0d00049..0e8ada7 100644
--- a/ops/pebble.py
+++ b/ops/pebble.py
@@ -509,6 +509,10 @@ class ServiceInfo:
         self.startup = startup
         self.current = current
 
+    def is_running(self) -> bool:
+        """Return True if this service is running (in the active state)."""
+        return self.current == ServiceStatus.ACTIVE
+
     @classmethod
     def from_dict(cls, d: Dict) -> 'ServiceInfo':
         """Create new object from dict parsed from JSON."""


Test patch:
diff --git a/test/test_model.py b/test/test_model.py
index 80b5137..a18fc28 100755
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -885,22 +885,24 @@ containers:
         self.pebble.responses.append(two_services)
         services = self.container.get_services()
         self.assertEqual(len(services), 2)
-        self.assertEqual(services[0].name, 's1')
-        self.assertEqual(services[0].startup, ops.pebble.ServiceStartup.ENABLED)
-        self.assertEqual(services[0].current, ops.pebble.ServiceStatus.ACTIVE)
-        self.assertEqual(services[1].name, 's2')
-        self.assertEqual(services[1].startup, ops.pebble.ServiceStartup.DISABLED)
-        self.assertEqual(services[1].current, ops.pebble.ServiceStatus.INACTIVE)
+        self.assertEqual(set(services), {'s1', 's2'})
+        self.assertEqual(services['s1'].name, 's1')
+        self.assertEqual(services['s1'].startup, ops.pebble.ServiceStartup.ENABLED)
+        self.assertEqual(services['s1'].current, ops.pebble.ServiceStatus.ACTIVE)
+        self.assertEqual(services['s2'].name, 's2')
+        self.assertEqual(services['s2'].startup, ops.pebble.ServiceStartup.DISABLED)
+        self.assertEqual(services['s2'].current, ops.pebble.ServiceStatus.INACTIVE)
 
         self.pebble.responses.append(two_services)
         services = self.container.get_services('s1', 's2')
         self.assertEqual(len(services), 2)
-        self.assertEqual(services[0].name, 's1')
-        self.assertEqual(services[0].startup, ops.pebble.ServiceStartup.ENABLED)
-        self.assertEqual(services[0].current, ops.pebble.ServiceStatus.ACTIVE)
-        self.assertEqual(services[1].name, 's2')
-        self.assertEqual(services[1].startup, ops.pebble.ServiceStartup.DISABLED)
-        self.assertEqual(services[1].current, ops.pebble.ServiceStatus.INACTIVE)
+        self.assertEqual(set(services), {'s1', 's2'})
+        self.assertEqual(services['s1'].name, 's1')
+        self.assertEqual(services['s1'].startup, ops.pebble.ServiceStartup.ENABLED)
+        self.assertEqual(services['s1'].current, ops.pebble.ServiceStatus.ACTIVE)
+        self.assertEqual(services['s2'].name, 's2')
+        self.assertEqual(services['s2'].startup, ops.pebble.ServiceStartup.DISABLED)
+        self.assertEqual(services['s2'].current, ops.pebble.ServiceStatus.INACTIVE)
 
         self.assertEqual(self.pebble.requests, [
             ('get_services', ()),
diff --git a/test/test_pebble.py b/test/test_pebble.py
index 9a0c1a2..aab2345 100644
--- a/test/test_pebble.py
+++ b/test/test_pebble.py
@@ -547,6 +547,13 @@ class TestServiceInfo(unittest.TestCase):
         self.assertEqual(s.startup, 'thingy')
         self.assertEqual(s.current, 'bob')
 
+    def test_is_running(self):
+        s = pebble.ServiceInfo('s', pebble.ServiceStartup.ENABLED, pebble.ServiceStatus.ACTIVE)
+        self.assertTrue(s.is_running())
+        for current in [pebble.ServiceStatus.INACTIVE, pebble.ServiceStatus.ERROR, 'other']:
+            s = pebble.ServiceInfo('s', pebble.ServiceStartup.ENABLED, current)
+            self.assertFalse(s.is_running())
+
 
 class MockClient(pebble.Client):
     """Mock Pebble client that simply records reqeusts and returns stored responses."""