Merge pull request #2776 from coreos-inc/joseph.schorr/QUAY-652/servicekeyworker-data-interface

Change service key worker to use a data interface

conf/init/service/servicekeyworker/run

@@ -4,6 +4,6 @@ echo 'Starting service key worker'
 
 QUAYPATH=${QUAYPATH:-"."}
 cd ${QUAYDIR:-"/"}
-PYTHONPATH=$QUAYPATH venv/bin/python -m workers.service_key_worker 2>&1
+PYTHONPATH=$QUAYPATH venv/bin/python -m workers.servicekeyworker.servicekeyworker 2>&1
 
 echo 'Service key worker exited'

workers/servicekeyworker/models_interface.py

@@ -0,0 +1,28 @@
+from abc import ABCMeta, abstractmethod
+from six import add_metaclass
+
+
+@add_metaclass(ABCMeta)
+class ServiceKeyWorkerDataInterface(object):
+  """
+  Interface that represents all data store interactions required by the service key worker.
+  """
+
+  @abstractmethod
+  def set_key_expiration(self, key_id, expiration_date):
+    """ Sets the expiration date of the service key with the given key ID to that given. """
+    pass
+
+  @abstractmethod
+  def create_service_key_for_testing(self, expiration):
+    """ Creates a service key for testing with the given expiration. Returns the KID for
+        key.
+    """
+    pass
+
+  @abstractmethod
+  def get_service_key_expiration(self, key_id):
+    """ Returns the expiration date for the key with the given ID. If the key doesn't exist or
+        does not have an expiration, returns None.
+    """
+    pass

workers/servicekeyworker/models_pre_oci.py

@@ -0,0 +1,21 @@
+from data import model
+from workers.servicekeyworker.models_interface import ServiceKeyWorkerDataInterface
+
+
+class PreOCIModel(ServiceKeyWorkerDataInterface):
+  def set_key_expiration(self, kid, expiration_date):
+    model.service_keys.set_key_expiration(kid, expiration_date)
+
+  def create_service_key_for_testing(self, expiration):
+    key = model.service_keys.create_service_key('test', 'somekid', 'quay', '', {}, expiration)
+    return key.kid
+
+  def get_service_key_expiration(self, kid):
+    try:
+      key = model.service_keys.get_service_key(kid, approved_only=False)
+      return key.expiration_date
+    except model.ServiceKeyDoesNotExist:
+      return None
+
+
+pre_oci_model = PreOCIModel()

workers/servicekeyworker/servicekeyworker.py

@@ -2,26 +2,28 @@ import logging
 from datetime import datetime, timedelta
 
 from app import app, instance_keys
-from data.model.service_keys import set_key_expiration
+from workers.servicekeyworker.models_pre_oci import pre_oci_model as model
 from workers.worker import Worker
 
 logger = logging.getLogger(__name__)
 
+
 class ServiceKeyWorker(Worker):
   def __init__(self):
     super(ServiceKeyWorker, self).__init__()
-    self.add_operation(self._refresh_service_keys,
-                       app.config.get('INSTANCE_SERVICE_KEY_REFRESH', 60)*60)
+    self.add_operation(self._refresh_service_key,
+                       app.config.get('INSTANCE_SERVICE_KEY_REFRESH', 60) * 60)
 
-  def _refresh_service_keys(self):
+  def _refresh_service_key(self):
     """
-    Refreshes active service keys so they don't get garbage collected.
+    Refreshes the instance's active service key so it doesn't get garbage collected.
     """
     expiration = timedelta(minutes=instance_keys.service_key_expiration)
     logger.debug('Starting refresh of automatic service keys')
-    set_key_expiration(instance_keys.local_key_id, datetime.now() + expiration)
+    model.set_key_expiration(instance_keys.local_key_id, datetime.now() + expiration)
     logger.debug('Finished refresh of automatic service keys')
 
+
 if __name__ == "__main__":
   worker = ServiceKeyWorker()
   worker.start()

workers/servicekeyworker/test/test_servicekeyworker.py

@@ -0,0 +1,22 @@
+from datetime import datetime, timedelta
+from mock import patch
+
+from data import model
+from workers.servicekeyworker.servicekeyworker import ServiceKeyWorker
+from util.morecollections import AttrDict
+
+from test.fixtures import *
+from workers.servicekeyworker.models_pre_oci import pre_oci_model as model
+
+def test_refresh_service_key(initialized_db):
+  # Create a service key for testing.
+  original_expiration = datetime.now() + timedelta(minutes=10)
+  test_key_kid = model.create_service_key_for_testing(original_expiration)
+
+  instance_keys = AttrDict(dict(local_key_id=test_key_kid, service_key_expiration=30))
+  with patch('workers.servicekeyworker.servicekeyworker.instance_keys', instance_keys):
+    worker = ServiceKeyWorker()
+    worker._refresh_service_key()
+
+  # Ensure the key's expiration was changed.
+  assert model.get_service_key_expiration(test_key_kid) > original_expiration