[CLIENT-3846] Speed up and optimize admin dev tests

test/new_tests/conftest.py

@@ -4,11 +4,13 @@
 import sys
 
 import pytest
-
 from . import invalid_data
 from .test_base_class import TestBaseClass
 
 import aerospike
+from aerospike import exception as e
+import logging
+
 
 # Comment this out because nowhere in the repository is using it
 '''
@@ -242,6 +244,95 @@ def invalid_key(request):
 # aerospike.set_log_level(aerospike.LOG_LEVEL_DEBUG)
 # aerospike.set_log_handler(None)
 
+HARD_LIMIT_SECS = 3
+# Server team states that this time interval is a safe amount
+POLL_INTERVAL_SECS = 0.1
+
+def poll_until_role_exists(role_name: str, client: aerospike.Client):
+    start = time.time()
+    while time.time() - start < HARD_LIMIT_SECS:
+        try:
+            client.admin_query_role(role=role_name)
+        except e.InvalidRole:
+            time.sleep(POLL_INTERVAL_SECS)
+            continue
+        logging.debug("Role now exists. Return early")
+        return
+    logging.debug("poll_until_role_exists timeout")
+
+def poll_until_role_doesnt_exist(role_name: str, client: aerospike.Client):
+    start = time.time()
+    try:
+        while time.time() - start < HARD_LIMIT_SECS:
+            client.admin_query_role(role=role_name)
+            time.sleep(POLL_INTERVAL_SECS)
+    except e.InvalidRole:
+        logging.debug("Role no longer exists. Return early")
+        return
+    logging.debug("poll_until_role_doesnt_exist timeout")
+
+def poll_until_user_exists(username: str, client: aerospike.Client, roles: list[str]):
+    start = time.time()
+    while time.time() - start < HARD_LIMIT_SECS:
+        try:
+            user_dict = client.admin_query_user_info(user=username)
+            if user_dict["roles"] != roles:
+                continue
+        except e.InvalidUser:
+            time.sleep(POLL_INTERVAL_SECS)
+            continue
+        logging.debug("User now exists with expected roles.")
+        return
+    logging.debug("poll_until_user_exists timeout")
+
+def poll_until_user_doesnt_exist(username: str, client: aerospike.Client):
+    start = time.time()
+    try:
+        while time.time() - start < HARD_LIMIT_SECS:
+            client.admin_query_user_info(user=username)
+            time.sleep(POLL_INTERVAL_SECS)
+    except e.InvalidUser:
+        logging.debug("User no longer exists. Return early")
+        return
+    logging.debug("poll_until_user_doesnt_exist timeout")
+
+def wait_for_job_completion(as_connection, job_id, job_module: int = aerospike.JOB_QUERY, time_limit_secs: float = float("inf")):
+    """
+    Blocks until the job has completed
+    """
+    start = time.time()
+    while time.time() - start < time_limit_secs:
+        response = as_connection.job_info(job_id, job_module)
+        if response["status"] != aerospike.JOB_STATUS_INPROGRESS:
+            break
+        time.sleep(0.1)
+
+# Monkeypatching the client class or instance isn't possible since it's immutable
+
+def admin_create_role_and_poll(client: aerospike.Client, role: str, *args, **kwargs):
+    retval = client.admin_create_role(role, *args, **kwargs)
+    poll_until_role_exists(role, client)
+    return retval
+
+def admin_create_user_and_poll(client: aerospike.Client, username: str, password: str, roles: list, *args, **kwargs):
+    retval = client.admin_create_user(username, password, roles, *args, **kwargs)
+
+    # The server creates a user and adds roles to it asynchronously, since security is a type of smd operation
+    # So client.admin_create_user() can return before all th
+    poll_until_user_exists(username, client, roles)
+
+    return retval
+
+def admin_drop_user_and_poll(client: aerospike.Client, username: str, *args, **kwargs):
+    retval = client.admin_drop_user(username, *args, **kwargs)
+    poll_until_user_doesnt_exist(username, client)
+    return retval
+
+def admin_drop_role_and_poll(client: aerospike.Client, role: str, *args, **kwargs):
+    retval = client.admin_drop_role(role, *args, **kwargs)
+    poll_until_role_doesnt_exist(role, client)
+    return retval
+
 def verify_record_ttl(client: aerospike.Client, key, expected_ttl: int):
     _, meta = client.exists(key)
     clock_skew_tolerance_secs = 50

test/new_tests/test_admin_change_password.py

@@ -3,25 +3,32 @@
 import pytest
 import time
 from .test_base_class import TestBaseClass
+from .conftest import admin_drop_user_and_poll, poll_until_user_doesnt_exist
 
 import aerospike
 
 
 @pytest.mark.usefixtures("connection_config")
-class TestChangePassword(object):
+class TestChangePassword:
     def setup_method(self, method):
         """
         Setup method
         """
         if TestBaseClass.auth_in_use() is False:
-            pytest.skip("No user specified, may not be a secured cluster", allow_module_level=True)
+            pytest.skip(
+                "No user specified, may not be a secured cluster",
+                allow_module_level=True,
+            )
 
         config = TestBaseClass.get_connection_config()
-        self.client = aerospike.client(config).connect(config["user"], config["password"])
+        self.client = aerospike.client(config).connect(
+            config["user"], config["password"]
+        )
 
         try:
-            self.client.admin_create_user("testchangepassworduser", "aerospike", ["read"])
-            time.sleep(2)
+            self.client.admin_create_user(
+                "testchangepassworduser", "aerospike", ["read"]
+            )
         except aerospike.exception.UserExistsError:
             pass  # we are good, no such role exists
         self.delete_users = []
@@ -32,7 +39,7 @@ def teardown_method(self, method):
         """
 
         try:
-            self.client.admin_drop_user("testchangepassworduser")
+            admin_drop_user_and_poll(self.client, "testchangepassworduser")
         except Exception:
             pass
 
@@ -56,13 +63,15 @@ def test_change_password_with_proper_parameters(self):
         status = self.clientreaduser.admin_change_password(user, password)
 
         assert status == 0
-
-        time.sleep(2)
         config = self.connection_config
 
         # Assert that connecting to the server with the old password fails
-        with pytest.raises((aerospike.exception.InvalidPassword, aerospike.exception.InvalidCredential)):
-            self.clientreaduserwrong = aerospike.client(config).connect(user, "aerospike")
+        with pytest.raises(
+            (aerospike.exception.InvalidPassword, aerospike.exception.InvalidCredential)
+        ):
+            self.clientreaduserwrong = aerospike.client(config).connect(
+                user, "aerospike"
+            )
 
         self.clientreaduserright = aerospike.client(config).connect(user, "newpassword")
 
@@ -101,8 +110,12 @@ def test_change_password_with_proper_timeout_policy_value(self):
         time.sleep(2)
         config = self.connection_config
 
-        with pytest.raises((aerospike.exception.InvalidPassword, aerospike.exception.InvalidCredential)):
-            self.clientreaduserwrong = aerospike.client(config).connect(user, "aerospike")
+        with pytest.raises(
+            (aerospike.exception.InvalidPassword, aerospike.exception.InvalidCredential)
+        ):
+            self.clientreaduserwrong = aerospike.client(config).connect(
+                user, "aerospike"
+            )
 
         self.clientreaduserright = aerospike.client(config).connect(user, "newpassword")