Adding a test

Author: Invictus17

leaderelection/electionconfig.py

@@ -30,16 +30,16 @@ def __init__(self, lock, lease_duration, renew_deadline, retry_period, onstarted
             sys.exit("lease_duration must be greater than renew_deadline")
 
         if renew_deadline <= self.jitter_factor * retry_period:
-            sys.exit("lease_duration must be greater than renew_deadline")
+            sys.exit("renewDeadline must be greater than retry_period*jitter_factor")
 
         if lease_duration < 1:
-            sys.exit("lease_duration must be greater than zero")
+            sys.exit("lease_duration must be greater than one")
 
         if renew_deadline < 1:
-            sys.exit("renew_deadline must be greater than zero")
+            sys.exit("renew_deadline must be greater than one")
 
         if retry_period < 1:
-            sys.exit("retry_period must be greater than zero")
+            sys.exit("retry_period must be greater than one")
 
         self.lease_duration = lease_duration
         self.renew_deadline = renew_deadline

leaderelection/leaderelection.py

@@ -17,11 +17,22 @@
 import time
 import json
 import threading
+from http import HTTPStatus
 from leaderelectionrecord import LeaderElectionRecord
 import logging
 
 logging.basicConfig(level=logging.INFO)
 
+"""
+This package implements leader election using an annotation in a Kubernetes object.
+The onstarted_leading function is run in a thread and when it returns, if it does 
+it might not be safe to run it again in a process.
+
+At first all candidates are considered followers. The one to create a lock or update
+an existing lock first becomes the leader and remains so until it keeps renewing its
+lease.
+"""
+
 
 class LeaderElection:
     def __init__(self, election_config):
@@ -102,7 +113,7 @@ def try_acquire_or_renew(self):
 
         # A lock is not created with that name, try to create one
         if not lock_status:
-            if json.loads(old_election_record.body)['code'] != 404:
+            if json.loads(old_election_record.body)['code'] != HTTPStatus.NOT_FOUND:
                 logging.info("Error retrieving resource lock {} as {}".format(self.election_config.lock.name,
                                                                               old_election_record.reason))
                 return False

leaderelection/tests/__init__.py

@@ -0,0 +1,247 @@
+# Copyright 2020 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import leaderelection
+from leaderelectionrecord import LeaderElectionRecord
+from kubernetes.client.rest import ApiException
+import electionconfig
+import unittest
+import threading
+import json
+import pytest
+
+class LeaderElectionTest(unittest.TestCase):
+    def test_simple_leader_election(self):
+        election_history = []
+        leadership_history = []
+
+        def on_create():
+            election_history.append("create record")
+            leadership_history.append("get leadership")
+
+        def on_update():
+            election_history.append("update record")
+
+        def on_change():
+            election_history.append("change record")
+
+        mock_lock = MockResourceLock("mock", "mock_namespace", "mock", on_create, on_update, on_change, None)
+
+        def on_started_leading():
+            leadership_history.append("start leading")
+
+        def on_stopped_leading():
+            leadership_history.append("stop leading")
+
+        # Create config 4.5 4 3
+        config = electionconfig.Config(lock=mock_lock, lease_duration=2.5,
+                                       renew_deadline=2, retry_period=1.5, onstarted_leading=on_started_leading,
+                                       onstopped_leading=on_stopped_leading)
+
+        # Enter leader election
+        leaderelection.LeaderElection(config).run()
+
+        self.assert_history(election_history, ["create record", "update record", "update record", "update record"])
+        self.assert_history(leadership_history, ["get leadership", "start leading", "stop leading"])
+
+    def test_leader_election(self):
+        election_history = []
+        leadership_history = []
+
+        def on_create_A():
+            election_history.append("A creates record")
+            leadership_history.append("A gets leadership")
+
+        def on_update_A():
+            election_history.append("A updates record")
+
+        def on_change_A():
+            election_history.append("A gets leadership")
+
+        mock_lock_A = MockResourceLock("mock", "mock_namespace", "MockA", on_create_A, on_update_A, on_change_A, None)
+        mock_lock_A.renew_count_max = 3
+
+        def on_started_leading_A():
+            leadership_history.append("A starts leading")
+
+        def on_stopped_leading_A():
+            leadership_history.append("A stops leading")
+
+        config_A = electionconfig.Config(lock=mock_lock_A, lease_duration=2.5,
+                                         renew_deadline=2, retry_period=1.5, onstarted_leading=on_started_leading_A,
+                                         onstopped_leading=on_stopped_leading_A)
+
+        # Enter leader election
+        leaderelection.LeaderElection(config_A).run()
+
+        def on_create_B():
+            election_history.append("B creates record")
+            leadership_history.append("B gets leadership")
+
+        def on_update_B():
+            election_history.append("B updates record")
+
+        def on_change_B():
+            leadership_history.append("B gets leadership")
+
+        mock_lock_B = MockResourceLock("mock", "mock_namespace", "MockB", on_create_B, on_update_B, on_change_B, None)
+        mock_lock_B.renew_count_max = 4
+
+        def on_started_leading_B():
+            leadership_history.append("B starts leading")
+
+        def on_stopped_leading_B():
+            leadership_history.append("B stops leading")
+
+        config_B = electionconfig.Config(lock=mock_lock_B, lease_duration=2.5,
+                                         renew_deadline=2, retry_period=1.5, onstarted_leading=on_started_leading_B,
+                                         onstopped_leading=on_stopped_leading_B)
+
+        mock_lock_B.leader_record = mock_lock_A.leader_record
+        leaderelection.LeaderElection(config_B).run()
+
+        self.assert_history(election_history,
+                            ["A creates record",
+                             "A updates record",
+                             "A updates record",
+                             "B updates record",
+                             "B updates record",
+                             "B updates record",
+                             "B updates record"])
+        self.assert_history(leadership_history,
+                            ["A gets leadership",
+                             "A starts leading",
+                             "A stops leading",
+                             "B gets leadership",
+                             "B starts leading",
+                             "B stops leading"])
+
+    def test_Leader_election_with_renew_deadline(self):
+        election_history = []
+        leadership_history = []
+
+        def on_create():
+            election_history.append("create record")
+            leadership_history.append("get leadership")
+
+        def on_update():
+            election_history.append("update record")
+
+        def on_change():
+            election_history.append("change record")
+
+        def on_try_update():
+            election_history.append("try update record")
+
+        mock_lock = MockResourceLock("mock", "mock_namespace", "mock", on_create, on_update, on_change, on_try_update)
+        mock_lock.renew_count_max = 3
+
+        def on_started_leading():
+            leadership_history.append("start leading")
+
+        def on_stopped_leading():
+            leadership_history.append("stop leading")
+
+        # Create config
+        config = electionconfig.Config(lock=mock_lock, lease_duration=2.5,
+                                       renew_deadline=2, retry_period=1.5, onstarted_leading=on_started_leading,
+                                       onstopped_leading=on_stopped_leading)
+
+        # Enter leader election
+        leaderelection.LeaderElection(config).run()
+
+        self.assert_history(election_history,
+                            ["create record",
+                             "try update record",
+                             "update record",
+                             "try update record",
+                             "update record",
+                             "try update record",
+                             "try update record"])
+
+        self.assert_history(leadership_history, ["get leadership", "start leading", "stop leading"])
+
+    def assert_history(self, history, expected):
+        self.assertIsNotNone(expected)
+        self.assertIsNotNone(history)
+        self.assertEquals(len(expected), len(history))
+
+        for idx in range(len(history)):
+            self.assertEquals(history[idx], expected[idx],
+                              msg="Not equal at index {}, expected {}, got {}".format(idx, expected[idx],
+                                                                                      history[idx]))
+
+
+class MockResourceLock:
+    def __init__(self, name, namespace, identity, on_create=None, on_update=None, on_change=None, on_try_update=None):
+        self.leader_record = None
+        self.renew_count = 0
+        self.renew_count_max = 4
+        self.name = name
+        self.namespace = namespace
+        self.identity = str(identity)
+        self.lock = threading.RLock()
+
+        self.on_create = on_create
+        self.on_update = on_update
+        self.on_change = on_change
+        self.on_try_update = on_try_update
+
+    def get(self, name, namespace):
+        self.lock.acquire()
+        try:
+            if self.leader_record:
+                return True, self.leader_record
+            ApiException.body = json.dumps({'code': 404})
+            return False, ApiException
+        finally:
+            self.lock.release()
+
+    def create(self, name, namespace, election_record):
+        self.lock.acquire()
+        try:
+            if self.leader_record is not None:
+                return False
+            self.leader_record = election_record
+            self.on_create()
+            self.renew_count += 1
+            return True
+        finally:
+            self.lock.release()
+
+    def update(self, name, namespace, updated_record):
+        self.lock.acquire()
+        try:
+            if self.on_try_update:
+                self.on_try_update()
+            if self.renew_count >= self.renew_count_max:
+                return False
+
+            old_record = self.leader_record
+            self.leader_record = updated_record
+
+            self.on_update()
+
+            if old_record.holder_identity != updated_record.holder_identity:
+                self.on_change()
+
+            self.renew_count += 1
+            return True
+        finally:
+            self.lock.release()
+
+
+if __name__ == '__main__':
+    unittest.main()