2853: Merge branch 'master' into 2853-rendezvous

[arvados.git] / sdk / python / tests / test_keep_client.py

diff --git a/sdk/python/tests/test_keep_client.py b/sdk/python/tests/test_keep_client.py
index c724efc1875329f5a7a78fef10af9dfca58a2674..bb0022eb432628dbf0c21b9b9fc2bd20a73f1fb1 100644 (file)
--- a/sdk/python/tests/test_keep_client.py
+++ b/sdk/python/tests/test_keep_client.py
@@ -1,5 +1,7 @@
+import hashlib
 import mock
 import os
+import re
 import socket
 import unittest
 import urlparse
@@ -9,6 +11,84 @@ import arvados.retry
 import arvados_testutil as tutil
 import run_test_server
 
+class KeepRendezvousWeightTestCase(unittest.TestCase):
+    def setUp(self):
+        self.keep_client = arvados.KeepClient(
+            api_client=mock.MagicMock(name='api_client'),
+            proxy='', local_store='')
+        self.keep_client._keep_services = []
+        self.n_services = 0
+
+    def addServices(self, n):
+        for x in range(n):
+            uuid = "zzzzz-bi6l4-{:015x}".format(self.n_services)
+            uri = "https://keep0x{:x}.zzzzz.arvadosapi.com:25107/".format(self.n_services)
+            self.keep_client._keep_services.append(
+                {'uuid': uuid, '_service_root': uri})
+            self.n_services += 1
+
+    def test_ProbeOrderReferenceSet(self):
+        # expected_order[i] is the probe order for
+        # hash=md5(sprintf("%064x",i)) where there are 16 services
+        # with uuid sprintf("anything-%015x",j) with j in 0..15. E.g.,
+        # the first probe for the block consisting of 64 "0"
+        # characters is the service whose uuid is
+        # "zzzzz-bi6l4-000000000000003", so expected_order[0][0]=='3'.
+        expected_order = [
+            list('3eab2d5fc9681074'),
+            list('097dba52e648f1c3'),
+            list('c5b4e023f8a7d691'),
+            list('9d81c02e76a3bf54'),
+            ]
+        hashes = [
+            hashlib.md5("{:064x}".format(x)).hexdigest()
+            for x in range(len(expected_order))]
+        self.addServices(16)
+        for i, hash in enumerate(hashes):
+            roots = self.keep_client.weighted_service_roots(hash)
+            got_order = [
+                re.search(r'//keep0x([0-9a-f]+)', root).group(1)
+                for root in roots]
+            self.assertEqual(got_order, expected_order[i])
+
+    def test_ProbeWasteAddingOneServer(self):
+        hashes = [
+            hashlib.md5("{:064x}".format(x)).hexdigest() for x in range(100)]
+        initial_services = 12
+        self.addServices(initial_services)
+        probes_before = [
+            self.keep_client.weighted_service_roots(hash) for hash in hashes]
+        for added_services in range(1, 12):
+            self.addServices(1)
+            total_penalty = 0
+            for hash_index in range(len(hashes)):
+                probe_after = self.keep_client.weighted_service_roots(
+                    hashes[hash_index])
+                penalty = probe_after.index(probes_before[hash_index][0])
+                self.assertLessEqual(penalty, added_services)
+                total_penalty += penalty
+            # Average penalty per block should not exceed
+            # N(added)/N(orig) by more than 20%, and should get closer
+            # to the ideal as we add data points.
+            expect_penalty = (
+                added_services *
+                len(hashes) / initial_services)
+            max_penalty = (
+                expect_penalty *
+                (120 - added_services)/100)
+            min_penalty = (
+                expect_penalty * 8/10)
+            self.assertTrue(
+                min_penalty <= total_penalty <= max_penalty,
+                "With {}+{} services, {} blocks, penalty {} but expected {}..{}".format(
+                    initial_services,
+                    added_services,
+                    len(hashes),
+                    total_penalty,
+                    min_penalty,
+                    max_penalty))
+
+
 class KeepTestCase(run_test_server.TestCaseWithServers):
     MAIN_SERVER = {}
     KEEP_SERVER = {}
@@ -249,7 +329,7 @@ class KeepClientServiceTestCase(unittest.TestCase):
     def get_service_roots(self, *services):
         api_client = self.mock_keep_services(*services)
         keep_client = arvados.KeepClient(api_client=api_client)
-        services = keep_client.shuffled_service_roots('000000')
+        services = keep_client.weighted_service_roots('000000')
         return [urlparse.urlparse(url) for url in sorted(services)]
 
     def test_ssl_flag_respected_in_roots(self):