Dataclass for passing around test config data

Author: ddworken

tests/bot-entrypoint.py

@@ -28,7 +28,7 @@ def load_env():
         "bin/keybaseca service &"
     ) % (shlex.quote(path)))
     # Sleep so keybaseca has time to start
-    time.sleep(2)
+    time.sleep(3)
     return "OK"
 
 if __name__ == '__main__':

tests/tests/conftest.py

@@ -1,21 +1,19 @@
 import pytest
 
-from lib import UtilitiesLib
-from lib import SUBTEAM, SUBTEAM_SECONDARY, USERNAME, BOT_USERNAME, EXPECTED_HASH
+from lib import TestConfig, run_command, clear_keys, clear_local_config
 
 @pytest.fixture(autouse=True)
 def run_around_tests():
-    lib = UtilitiesLib(SUBTEAM, SUBTEAM_SECONDARY, USERNAME, BOT_USERNAME, EXPECTED_HASH)
-    lib.clear_keys()
-    lib.clear_local_config()
+    clear_keys()
+    clear_local_config()
     # Calling yield triggers the test execution
     yield
 
 def pytest_sessionfinish(session, exitstatus):
     # Automatically run after all tests in order to ensure that no kssh-client config files stick around
-    lib = UtilitiesLib(SUBTEAM, SUBTEAM_SECONDARY, USERNAME, BOT_USERNAME, EXPECTED_HASH)
-    lib.run_command(f"keybase fs rm /keybase/team/{lib.subteam}.ssh/kssh-client.config || true" )
-    lib.run_command(f"keybase fs rm /keybase/team/{lib.subteam}.ssh.staging/kssh-client.config || true" )
-    lib.run_command(f"keybase fs rm /keybase/team/{lib.subteam}.ssh.prod/kssh-client.config || true" )
-    lib.run_command(f"keybase fs rm /keybase/team/{lib.subteam}.ssh.root_everywhere/kssh-client.config || true" )
-    lib.run_command(f"keybase fs rm /keybase/team/{lib.subteam_secondary}/kssh-client.config || true" )
+    tc = TestConfig.getDefaultTestConfig()
+    run_command(f"keybase fs rm /keybase/team/{tc.subteam}.ssh/kssh-client.config || true" )
+    run_command(f"keybase fs rm /keybase/team/{tc.subteam}.ssh.staging/kssh-client.config || true" )
+    run_command(f"keybase fs rm /keybase/team/{tc.subteam}.ssh.prod/kssh-client.config || true" )
+    run_command(f"keybase fs rm /keybase/team/{tc.subteam}.ssh.root_everywhere/kssh-client.config || true" )
+    run_command(f"keybase fs rm /keybase/team/{tc.subteam_secondary}/kssh-client.config || true" )

tests/tests/lib.py

@@ -4,113 +4,123 @@
 import signal
 import subprocess
 import time
+from typing import List
 
 import requests
 
-SUBTEAM = os.environ['SUBTEAM']
-SUBTEAM_SECONDARY = os.environ['SUBTEAM_SECONDARY']
-USERNAME = os.environ['KSSH_USERNAME']
-BOT_USERNAME = os.environ['BOT_USERNAME']
+def getDefaultExpectedHash() -> bytes:
+    # "uniquestring" is stored in /etc/unique of the SSH server. We then run the command `sha1sum /etc/unique` via kssh
+    # and assert that the output contains the sha1 hash of uniquestring. This checks to make sure the command given to
+    # kssh is actually executing on the remote server.
+    return hashlib.sha1(b"uniquestring").hexdigest().encode('utf-8')
 
-# "uniquestring" is stored in /etc/unique of the SSH server. We then run the command `sha1sum /etc/unique` via kssh
-# and assert that the output contains the sha1 hash of uniquestring. This checks to make sure the command given to
-# kssh is actually executing on the remote server.
-EXPECTED_HASH = hashlib.sha1(b"uniquestring").hexdigest().encode('utf-8')
+class TestConfig:
+    # Not actually a test class so mark it to be skipped
+    __test__ = False
 
-class UtilitiesLib:
     def __init__(self, subteam, subteam_secondary, username, bot_username, expected_hash):
         self.subteam = subteam
         self.subteam_secondary = subteam_secondary
         self.username = username
         self.bot_username = bot_username
         self.expected_hash = expected_hash
 
-    def run_command(self, cmd, timeout=10):
-        # In order to properly run a command with a timeout and shell=True, we use Popen with a shell and group all child
-        # processes so we can kill all of them. See:
-        # - https://stackoverflow.com/questions/36952245/subprocess-timeout-failure
-        # - https://stackoverflow.com/questions/4789837/how-to-terminate-a-python-subprocess-launched-with-shell-true
-        with subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, preexec_fn=os.setsid) as process:
-            try:
-                stdout, stderr = process.communicate(timeout=timeout)
-                if process.returncode != 0:
-                    raise subprocess.CalledProcessError(process.returncode, cmd, stdout, stderr)
-                return stdout
-            except subprocess.TimeoutExpired as e:
-                os.killpg(process.pid, signal.SIGINT)
-                print(f"Output before timeout: {process.communicate()[0]}")
-                raise e
-
-    def read_file(self, filename):
-        """
-        Read the contents of the given filename to a list of strings. If it is a normal file,
-        uses the standard open() function. Otherwise, uses `keybase fs read`.
-        :param filename:    The name of the file to read
-        :return:            A list of lines in the file
-        """
-        if filename.startswith("/keybase/"):
-            return self.run_command(f"keybase fs read {filename}").splitlines()
-        with open(filename, 'rb') as f:
-            return f.readlines()
-
-    def clear_keys(self):
-        # Clear all keys generated by kssh
+    @staticmethod
+    def getDefaultTestConfig():
+        return TestConfig(
+            os.environ['SUBTEAM'],
+            os.environ['SUBTEAM_SECONDARY'],
+            os.environ['KSSH_USERNAME'],
+            os.environ['BOT_USERNAME'],
+            getDefaultExpectedHash()
+        )
+
+def run_command(cmd: str, timeout: int=10) -> bytes:
+    # In order to properly run a command with a timeout and shell=True, we use Popen with a shell and group all child
+    # processes so we can kill all of them. See:
+    # - https://stackoverflow.com/questions/36952245/subprocess-timeout-failure
+    # - https://stackoverflow.com/questions/4789837/how-to-terminate-a-python-subprocess-launched-with-shell-true
+    with subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, preexec_fn=os.setsid) as process:
         try:
-            self.run_command("rm -rf ~/.ssh/keybase-signed-key*")
-        except subprocess.CalledProcessError:
-            pass
-
-    def clear_local_config(self):
-        # Clear kssh's local config file
-        try:
-            self.run_command("rm -rf ~/.ssh/kssh.config")
-        except subprocess.CalledProcessError:
-            pass
-
-    def load_env(self, filename):
-        # Load the environment based off of the given filename which is the path to the python test script
-        env_name = os.path.basename(filename).split(".")[0]
-        return requests.get(f"http://ca-bot:8080/load_env?filename={env_name}").content == b"OK"
-
-    def assert_contains_hash(self, output):
-        assert EXPECTED_HASH in output
-
-    @contextmanager
-    def simulate_two_teams(self):
-        # A context manager that simulates running the given function in an environment with two teams set up
-        self.run_command(f"keybase fs read /keybase/team/{self.subteam}.ssh.staging/kssh-client.config | "
-                         f"sed 's/{self.subteam}.ssh.staging/{self.subteam_secondary}/g' | "
-                         f"sed 's/{self.bot_username}/otherbotname/g' | "
-                         f"keybase fs write /keybase/team/{self.subteam_secondary}/kssh-client.config")
-        try:
-            yield
-        finally:
-            self.run_command(f"keybase fs rm /keybase/team/{self.subteam_secondary}/kssh-client.config")
+            stdout, stderr = process.communicate(timeout=timeout)
+            if process.returncode != 0:
+                raise subprocess.CalledProcessError(process.returncode, cmd, stdout, stderr)
+            return stdout
+        except subprocess.TimeoutExpired as e:
+            os.killpg(process.pid, signal.SIGINT)
+            print(f"Output before timeout: {process.communicate()[0]}")
+            raise e
+
+def read_file(filename: str) -> List[bytes]:
+    """
+    Read the contents of the given filename to a list of strings. If it is a normal file,
+    uses the standard open() function. Otherwise, uses `keybase fs read`.
+    :param filename:    The name of the file to read
+    :return:            A list of lines in the file
+    """
+    if filename.startswith("/keybase/"):
+        return run_command(f"keybase fs read {filename}").splitlines()
+    with open(filename, 'rb') as f:
+        return f.readlines()
+
+def clear_keys():
+    # Clear all keys generated by kssh
+    try:
+        run_command("rm -rf ~/.ssh/keybase-signed-key*")
+    except subprocess.CalledProcessError:
+        pass
+
+def clear_local_config():
+    # Clear kssh's local config file
+    try:
+        run_command("rm -rf ~/.ssh/kssh.config")
+    except subprocess.CalledProcessError:
+        pass
+
+def load_env(filename: str):
+    # Load the environment based off of the given filename which is the path to the python test script
+    env_name = os.path.basename(filename).split(".")[0]
+    return requests.get(f"http://ca-bot:8080/load_env?filename={env_name}").content == b"OK"
+
+def assert_contains_hash(hash: bytes, output: bytes):
+    assert hash in output
+
+@contextmanager
+def simulate_two_teams(tc: TestConfig):
+    # A context manager that simulates running the given function in an environment with two teams set up
+    run_command(f"keybase fs read /keybase/team/{tc.subteam}.ssh.staging/kssh-client.config | "
+                     f"sed 's/{tc.subteam}.ssh.staging/{tc.subteam_secondary}/g' | "
+                     f"sed 's/{tc.bot_username}/otherbotname/g' | "
+                     f"keybase fs write /keybase/team/{tc.subteam_secondary}/kssh-client.config")
+    try:
+        yield
+    finally:
+        run_command(f"keybase fs rm /keybase/team/{tc.subteam_secondary}/kssh-client.config")
 
-    @contextmanager
-    def outputs_audit_log(self, filename, expected_number):
-        # A context manager that asserts that the given function triggers expected_number of audit logs to be added to '/keybase/team/team.ssh.prod/ca.log'
-        # Note that fuse is not running in the container so this has to use `keybase fs read`
+@contextmanager
+def outputs_audit_log(tc: TestConfig, filename: str, expected_number: int):
+    # A context manager that asserts that the given function triggers expected_number of audit logs to be added to '/keybase/team/team.ssh.prod/ca.log'
+    # Note that fuse is not running in the container so this has to use `keybase fs read`
 
-        # Make a set of the lines in the audit log before we ran
-        before_lines = set(self.read_file(filename))
+    # Make a set of the lines in the audit log before we ran
+    before_lines = set(read_file(filename))
 
-        # Then run the code inside the context manager
-        yield
+    # Then run the code inside the context manager
+    yield
 
-        # And sleep to give KBFS some time
-        time.sleep(1.5)
+    # And sleep to give KBFS some time
+    time.sleep(1.5)
 
-        # Then see if there are new lines using set difference. This is only safe/reasonable since we include a
-        # timestamp in audit log lines.
-        after_lines = set(self.read_file(filename))
-        new_lines = after_lines - before_lines
+    # Then see if there are new lines using set difference. This is only safe/reasonable since we include a
+    # timestamp in audit log lines.
+    after_lines = set(read_file(filename))
+    new_lines = after_lines - before_lines
 
-        cnt = 0
-        for line in new_lines:
-            line = line.decode('utf-8')
-            if line and f"Processing SignatureRequest from user={self.username}" in line and f"principals:{self.subteam}.ssh.staging,{self.subteam}.ssh.root_everywhere, expiration:+1h, pubkey:ssh-ed25519" in line:
-                cnt += 1
+    cnt = 0
+    for line in new_lines:
+        line = line.decode('utf-8')
+        if line and f"Processing SignatureRequest from user={tc.username}" in line and f"principals:{tc.subteam}.ssh.staging,{tc.subteam}.ssh.root_everywhere, expiration:+1h, pubkey:ssh-ed25519" in line:
+            cnt += 1
 
-        if cnt != expected_number:
-            assert False, f"Found {cnt} audit log entries, expected {expected_number}! New audit logs: {new_lines}"
+    if cnt != expected_number:
+        assert False, f"Found {cnt} audit log entries, expected {expected_number}! New audit logs: {new_lines}"