Add utility classes for indication testing

1 files changed, 53 insertions, 0 deletions

diff --git a/src/account/test/common.py b/src/account/test/common.py
index 2818006..07543b3 100644
--- a/src/account/test/common.py
+++ b/src/account/test/common.py
@@ -210,3 +210,56 @@ class UserOps(object):
         Return true/false if user does/does not exist
         """
         return name in cls.list_users()
+
+
+class TestUserSet(object):
+    """
+    Class to hold list of "testing" users, able to make up new names
+    """
+
+    def __init__(self, prefix="test_user_", strength=8):
+        self.prefix = prefix
+        self.strength = strength
+        self._our_users = set()
+
+    def _new_name(self):
+        """
+        Make up a name that is not yet on the system
+        """
+        name = None
+        existing = UserOps.list_users()
+        rs = lmi.test.util.random_string
+        while not name or name in existing:
+            name = rs(strength=self.strength, prefix=self.prefix)
+        return name
+
+    def add(self, name=None):
+        """
+        Create a testing user; return the name
+        """
+        name = name if name else self._new_name()
+        UserOps.create_account(name)
+        self._our_users.add(name)
+        return name
+
+    def remove(self, name=None):
+        """
+        Remove given user or a random one; return the removed name
+        """
+        try:
+            name = name if name else next(iter(self._our_users))
+            assert name in self._our_users
+        except StopIteration:   # from above next() call
+            raise ValueError("no more testing users")
+        except AssertionError:
+            raise ValueError("not our testing user: %s" % name)
+        UserOps.clean_account(name)
+        self._our_users.remove(name)
+        return name
+
+    def remove_all(self):
+        """
+        Remove all testing users
+        """
+        [UserOps.clean_account(n) for n in self._our_users]
+        self._our_users = set()