Enable parallel multi-user integration tests with Kerberos (#707)

integration/README.md

@@ -91,9 +91,9 @@ and `{{COOK_SCHEDULER_URL}}`,
 which will be replaced with the Cook Scheduler URL being used by the current tests.
 The command must print (to `stdout`) the `Authorization` header value for the target user.
 
-When using Kerberos, multi-user integration tests must *not* be run in parallel.
-This restriction is due to only having 10 unique test user names,
-which are reused across each of the multi-user tests.
+* `COOK_MAX_TEST_USERS`:
+Setting this environment variable limits the number of test users allocated to each test runner process.
+This is most likely necessary in any scenario where users are actually authenticated by the underlying OS.
 
 ## Credits
 

integration/tests/cook/conftest.py

@@ -0,0 +1,30 @@
+# This file is automatically loaded and run by pytest during its setup process,
+# meaning it happens before any of the tests in this directory are run.
+# See the pytest documentation on conftest files for more information:
+# https://docs.pytest.org/en/2.7.3/plugins.html#conftest-py-plugins
+
+import os
+import subprocess
+import threading
+import time
+
+from tests.cook import util
+
+def _sudo_check(username):
+    """
+    Check if the current user can sudo as a test user.
+    This is necessary to obtain Kerberos auth headers for multi-user tests.
+    """
+    sudo_ok = (0 == subprocess.call(f'sudo -nu {username} echo CACHED SUDO', shell=True))
+    assert sudo_ok, "You need to pre-cache your sudo credentials. (Run a simple sudo command as a test user.)"
+
+def _sudo_checker_task(username):
+    """Periodically check sudo ability to ensure the credentials stay cached."""
+    while True:
+        _sudo_check(username)
+        time.sleep(60)
+
+if util.kerberos_enabled() and os.getenv('COOK_MAX_TEST_USERS'):
+    username = next(util._test_user_names())
+    _sudo_check(username)
+    threading.Thread(target=_sudo_checker_task, args=[username], daemon=True).start()

integration/tests/cook/util.py

@@ -27,9 +27,10 @@
 # Name of our custom HTTP header for user impersonation
 IMPERSONATION_HEADER = 'X-Cook-Impersonate'
 
+
 def continuous_integration():
     """Returns true if the CONTINUOUS_INTEGRATION environment variable is set, as done by Travis-CI."""
-    return os.environ.get('CONTINUOUS_INTEGRATION')
+    return os.getenv('CONTINUOUS_INTEGRATION')
 
 
 def has_docker_service():
@@ -42,9 +43,41 @@ def has_docker_service():
 _default_admin_name = 'root'
 _default_impersonator_name = 'poser'
 
+
 def _get_default_user_name():
     return os.getenv('USER', _default_user_name)
 
+
+@functools.lru_cache()
+def _test_user_ids():
+    """
+    Get the numeric user suffixes for this test worker.
+    Returns the range 0 to 1 million if COOK_MAX_TEST_USERS is not set.
+    If this is a distributed run with a limited number of users,
+    e.g., 10 per worker, then this function returns range(0, 10) for worker 0,
+    or range(20, 30) for worker 2.
+    """
+    pytest_worker = os.getenv('PYTEST_XDIST_WORKER')
+    max_test_users = int(os.getenv('COOK_MAX_TEST_USERS', 0))
+    if pytest_worker and max_test_users:
+        pytest_worker_id = int(pytest_worker[2:]) # e.g., "gw4" -> 4
+        test_user_min_id = max_test_users * pytest_worker_id
+        test_user_max_id = test_user_min_id + max_test_users
+        return range(test_user_min_id, test_user_max_id)
+    else:
+        return range(1000000)
+
+
+def _test_user_names(test_name_prefix=None):
+    """
+    Returns a generator of unique test user names, with form {PREFIX}{ID}.
+    The COOK_TEST_USER_PREFIX environment variable is used by default;
+    otherwise, the test_name_prefix value is used as the PREFIX.
+    """
+    name_prefix = os.getenv('COOK_TEST_USER_PREFIX', test_name_prefix)
+    return (f'{name_prefix}{i}' for i in _test_user_ids())
+
+
 # Shell command used to obtain Kerberos credentials for a given test user
 _kerberos_missing_cmd =  'echo "MISSING COOK_KERBEROS_TEST_AUTH_CMD" && exit 1'
 _kerberos_auth_cmd = os.getenv('COOK_KERBEROS_TEST_AUTH_CMD', _kerberos_missing_cmd)
@@ -107,17 +140,6 @@ def __exit__(self, ex_type, ex_val, ex_trace):
         self.previous_auth = None
 
 
-@functools.lru_cache()
-def _generate_kerberos_ticket_for_user(username):
-    """
-    Get a Kerberos authentication ticket for the given user.
-    Depends on COOK_KERBEROS_TEST_AUTH_CMD being set in the environment.
-    """
-    subcommand = (_kerberos_auth_cmd
-                  .replace('{{COOK_USER}}', username)
-                  .replace('{{COOK_SCHEDULER_URL}}', retrieve_cook_url()))
-    return subprocess.check_output(subcommand, shell=True).rstrip()
-
 
 class _KerberosUser(_AuthenticatedUser):
     """
@@ -129,9 +151,20 @@ def __init__(self, name, impersonatee=None):
         subcommand = (_kerberos_auth_cmd
                       .replace('{{COOK_USER}}', name)
                       .replace('{{COOK_SCHEDULER_URL}}', retrieve_cook_url()))
-        self.auth_token = _generate_kerberos_ticket_for_user(name)
+        self.auth_token = self._generate_kerberos_ticket_for_user(name)
         self.previous_token = None
 
+    @functools.lru_cache()
+    def _generate_kerberos_ticket_for_user(self, username):
+        """
+        Get a Kerberos authentication ticket for the given user.
+        Depends on COOK_KERBEROS_TEST_AUTH_CMD being set in the environment.
+        """
+        subcommand = (_kerberos_auth_cmd
+                      .replace('{{COOK_USER}}', username)
+                      .replace('{{COOK_SCHEDULER_URL}}', retrieve_cook_url()))
+        return subprocess.check_output(subcommand, shell=True).rstrip()
+
     def __enter__(self):
         global session
         super().__enter__()
@@ -167,16 +200,15 @@ def __init__(self, test_handle):
         if test_handle:
             test_id = test_handle.id()
             test_base_name = test_id[test_id.rindex('.test_')+6:].lower()
-            base_name = os.getenv('COOK_TEST_USER_PREFIX', f'{test_base_name}_')
-            self.__user_generator = (f'{base_name}{i}' for i in range(1000000))
+            self.__user_generator = _test_user_names(test_base_name)
 
     def new_user(self):
         """Return a fresh user object."""
         return self.user_class(next(self.__user_generator))
 
     def new_users(self, count=None):
         """Return a sequence of `count` fresh user objects."""
-        return map(self.user_class, itertools.islice(self.__user_generator, 0, count))
+        return [self.user_class(x) for x in itertools.islice(self.__user_generator, 0, count)]
 
     @functools.lru_cache()
     def default(self):