refactored EtcdClient

syncthing_monitor/etcd_client.py

@@ -4,40 +4,32 @@ import etcd3 as etcd3
 from retrying import retry
 
 CLUSTER_INFO_KEY = '/syncthing_monitor/cluster_info'
+DEFAULT_PORT = 2379
 
 
 class EtcdClient:
-    def __init__(self, host):
-        self.host = host
+    def __init__(self, host, port=DEFAULT_PORT):
+        self.etcd = etcd3.client(host=host, port=port)
+
+    def load_cluster_info(self):
+        raw_value, meta = self.etcd.get(CLUSTER_INFO_KEY)
+        if raw_value is None:
+            return {'devices': []}
+        return json.loads(raw_value)
 
     @retry
-    def append_device_id(self, device_id, port=2379):
-        cluster_info = {'devices': []}
-
-        etcd = etcd3.client(host=self.host, port=port)
-
-        with etcd.lock('syncthing_monitor'):
-            raw_value = etcd.get(CLUSTER_INFO_KEY)[0]
-
-            if raw_value is not None:
-                cluster_info = json.loads(raw_value)
+    def append_device_id(self, device_id):
+        with self.etcd.lock('syncthing_monitor'):
+            cluster_info = self.load_cluster_info()
 
             if not any(device_id in element for element in cluster_info['devices']):
                 cluster_info['devices'].append(device_id)
 
-            print("Updating etcd value for '{0}': {1}".format(CLUSTER_INFO_KEY, json.dumps(cluster_info)), flush=True)
-            etcd.put(CLUSTER_INFO_KEY, json.dumps(cluster_info))
+            print("Updating etcd value for '{0}': {1}".format(CLUSTER_INFO_KEY, json.dumps(cluster_info)))
+            self.etcd.put(CLUSTER_INFO_KEY, json.dumps(cluster_info))
 
     @retry
-    def get_device_list(self, port=2379):
-        cluster_info = {'devices': []}
-
-        etcd = etcd3.client(host=self.host, port=port)
-
-        raw_value = etcd.get(CLUSTER_INFO_KEY)[0]
-
-        if raw_value is not None:
-            cluster_info = json.loads(raw_value)
-
-        print("Obtained cluster_info devices from etcd: {0}".format(json.dumps(cluster_info['devices'])))
-        return cluster_info['devices']
+    def get_device_list(self):
+        device_list = self.load_cluster_info()['devices']
+        print("Obtained device_list devices from etcd: {0}".format(json.dumps(device_list)))
+        return device_list