attempt to address paramiko connection errors (nebari-dev#2811)

Author: dcmcand

.github/workflows/test_local_integration.yaml

@@ -73,6 +73,11 @@ jobs:
           python-version: "3.11"
           miniconda-version: "latest"
 
+      - name: Install JQ
+        run: |
+          sudo apt-get update
+          sudo apt-get install jq -y
+
       - name: Install Nebari and playwright
         run: |
           pip install .[dev]
@@ -97,6 +102,14 @@ jobs:
           nebari keycloak adduser --user "${TEST_USERNAME}" "${TEST_PASSWORD}" --config ${{ steps.init.outputs.config }}
           nebari keycloak listusers --config ${{ steps.init.outputs.config }}
 
+      - name: Await Workloads
+        uses: jupyterhub/action-k8s-await-workloads@v3
+        with:
+          workloads: "" # all
+          namespace: "dev"
+          timeout: 60
+          max-restarts: 0
+
       ### DEPLOYMENT TESTS
       - name: Deployment Pytests
         env:

tests/tests_deployment/test_jupyterhub_ssh.py

@@ -1,5 +1,6 @@
 import re
 import string
+import time
 import uuid
 
 import paramiko
@@ -14,64 +15,80 @@
 TIMEOUT_SECS = 300
 
 
-@pytest.fixture(scope="function")
+@pytest.fixture(scope="session")
 def paramiko_object(jupyterhub_access_token):
-    """Connects to JupyterHub ssh cluster from outside the cluster."""
+    """Connects to JupyterHub SSH cluster from outside the cluster.
+
+    Ensures the JupyterLab pod is ready before attempting reauthentication
+    by setting both `auth_timeout` and `banner_timeout` appropriately,
+    and by retrying the connection until the pod is ready or a timeout occurs.
+    """
     params = {
         "hostname": constants.NEBARI_HOSTNAME,
         "port": 8022,
         "username": constants.KEYCLOAK_USERNAME,
         "password": jupyterhub_access_token,
         "allow_agent": constants.PARAMIKO_SSH_ALLOW_AGENT,
         "look_for_keys": constants.PARAMIKO_SSH_LOOK_FOR_KEYS,
-        "auth_timeout": 5 * 60,
     }
 
     ssh_client = paramiko.SSHClient()
     ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-    try:
-        ssh_client.connect(**params)
-        yield ssh_client
-    finally:
-        ssh_client.close()
-
-
-def run_command(command, stdin, stdout, stderr):
-    delimiter = uuid.uuid4().hex
-    stdin.write(f"echo {delimiter}start; {command}; echo {delimiter}end\n")
-
-    output = []
-
-    line = stdout.readline()
-    while not re.match(f"^{delimiter}start$", line.strip()):
-        line = stdout.readline()
 
-    line = stdout.readline()
-    if delimiter not in line:
-        output.append(line)
-
-    while not re.match(f"^{delimiter}end$", line.strip()):
-        line = stdout.readline()
-        if delimiter not in line:
-            output.append(line)
-
-    return "".join(output).strip()
-
-
-@pytest.mark.timeout(TIMEOUT_SECS)
-@pytest.mark.filterwarnings("ignore::urllib3.exceptions.InsecureRequestWarning")
-@pytest.mark.filterwarnings("ignore::ResourceWarning")
-def test_simple_jupyterhub_ssh(paramiko_object):
-    stdin, stdout, stderr = paramiko_object.exec_command("")
+    yield ssh_client, params
+
+    ssh_client.close()
+
+
+def invoke_shell(
+    client: paramiko.SSHClient, params: dict[str, any]
+) -> paramiko.Channel:
+    client.connect(**params)
+    return client.invoke_shell()
+
+
+def extract_output(delimiter: str, output: str) -> str:
+    # Extract the command output between the start and end delimiters
+    match = re.search(rf"{delimiter}start\n(.*)\n{delimiter}end", output, re.DOTALL)
+    if match:
+        print(match.group(1).strip())
+        return match.group(1).strip()
+    else:
+        return output.strip()
+
+
+def run_command_list(
+    commands: list[str], channel: paramiko.Channel, wait_time: int = 0
+) -> dict[str, str]:
+    command_delimiters = {}
+    for command in commands:
+        delimiter = uuid.uuid4().hex
+        command_delimiters[command] = delimiter
+        b = channel.send(f"echo {delimiter}start; {command}; echo {delimiter}end\n")
+        if b == 0:
+            print(f"Command '{command}' failed to send")
+    # Wait for the output to be ready before reading
+    time.sleep(wait_time)
+    while not channel.recv_ready():
+        time.sleep(1)
+        print("Waiting for output")
+    output = ""
+    while channel.recv_ready():
+        output += channel.recv(65535).decode("utf-8")
+    outputs = {}
+    for command, delimiter in command_delimiters.items():
+        command_output = extract_output(delimiter, output)
+        outputs[command] = command_output
+    return outputs
 
 
 @pytest.mark.timeout(TIMEOUT_SECS)
 @pytest.mark.filterwarnings("ignore::urllib3.exceptions.InsecureRequestWarning")
 @pytest.mark.filterwarnings("ignore::ResourceWarning")
 def test_print_jupyterhub_ssh(paramiko_object):
-    stdin, stdout, stderr = paramiko_object.exec_command("")
-
-    # commands to run and just print the output
+    client, params = paramiko_object
+    channel = invoke_shell(client, params)
+    # Commands to run and just print the output
     commands_print = [
         "id",
         "env",
@@ -80,52 +97,60 @@ def test_print_jupyterhub_ssh(paramiko_object):
         "ls -la",
         "umask",
     ]
-
-    for command in commands_print:
-        print(f'COMMAND: "{command}"')
-        print(run_command(command, stdin, stdout, stderr))
+    outputs = run_command_list(commands_print, channel)
+    for command, output in outputs.items():
+        print(f"COMMAND: {command}")
+        print(f"OUTPUT: {output}")
+    channel.close()
 
 
 @pytest.mark.timeout(TIMEOUT_SECS)
 @pytest.mark.filterwarnings("ignore::urllib3.exceptions.InsecureRequestWarning")
 @pytest.mark.filterwarnings("ignore::ResourceWarning")
 def test_exact_jupyterhub_ssh(paramiko_object):
-    stdin, stdout, stderr = paramiko_object.exec_command("")
-
-    # commands to run and exactly match output
-    commands_exact = [
-        ("id -u", "1000"),
-        ("id -g", "100"),
-        ("whoami", constants.KEYCLOAK_USERNAME),
-        ("pwd", f"/home/{constants.KEYCLOAK_USERNAME}"),
-        ("echo $HOME", f"/home/{constants.KEYCLOAK_USERNAME}"),
-        ("conda activate default && echo $CONDA_PREFIX", "/opt/conda/envs/default"),
-        (
-            "hostname",
-            f"jupyter-{escape_string(constants.KEYCLOAK_USERNAME, safe=set(string.ascii_lowercase + string.digits), escape_char='-').lower()}",
-        ),
-    ]
+    client, params = paramiko_object
+    channel = invoke_shell(client, params)
+    # Commands to run and exactly match output
+    commands_exact = {
+        "id -u": "1000",
+        "id -g": "100",
+        "whoami": constants.KEYCLOAK_USERNAME,
+        "pwd": f"/home/{constants.KEYCLOAK_USERNAME}",
+        "echo $HOME": f"/home/{constants.KEYCLOAK_USERNAME}",
+        "conda activate default && echo $CONDA_PREFIX": "/opt/conda/envs/default",
+        "hostname": f"jupyter-{escape_string(constants.KEYCLOAK_USERNAME, safe=set(string.ascii_lowercase + string.digits), escape_char='-').lower()}",
+    }
+    outputs = run_command_list(list(commands_exact.keys()), channel)
+    for command, output in outputs.items():
+        assert (
+            output == outputs[command]
+        ), f"Command '{command}' output '{outputs[command]}' does not match expected '{output}'"
 
-    for command, output in commands_exact:
-        assert output == run_command(command, stdin, stdout, stderr)
+    channel.close()
 
 
 @pytest.mark.timeout(TIMEOUT_SECS)
 @pytest.mark.filterwarnings("ignore::urllib3.exceptions.InsecureRequestWarning")
 @pytest.mark.filterwarnings("ignore::ResourceWarning")
 def test_contains_jupyterhub_ssh(paramiko_object):
-    stdin, stdout, stderr = paramiko_object.exec_command("")
-
-    # commands to run and string need to be contained in output
-    commands_contain = [
-        ("ls -la", ".bashrc"),
-        ("cat ~/.bashrc", "Managed by Nebari"),
-        ("cat ~/.profile", "Managed by Nebari"),
-        ("cat ~/.bash_logout", "Managed by Nebari"),
-        # ensure we don't copy over extra files from /etc/skel in init container
-        ("ls -la ~/..202*", "No such file or directory"),
-        ("ls -la ~/..data", "No such file or directory"),
-    ]
+    client, params = paramiko_object
+    channel = invoke_shell(client, params)
+
+    # Commands to run and check if the output contains specific strings
+    commands_contain = {
+        "ls -la": ".bashrc",
+        "cat ~/.bashrc": "Managed by Nebari",
+        "cat ~/.profile": "Managed by Nebari",
+        "cat ~/.bash_logout": "Managed by Nebari",
+        # Ensure we don't copy over extra files from /etc/skel in init container
+        "ls -la ~/..202*": "No such file or directory",
+        "ls -la ~/..data": "No such file or directory",
+    }
+
+    outputs = run_command_list(commands_contain.keys(), channel, 30)
+    for command, expected_output in commands_contain.items():
+        assert (
+            expected_output in outputs[command]
+        ), f"Command '{command}' output does not contain expected substring '{expected_output}'. Instead got '{outputs[command]}'"
 
-    for command, output in commands_contain:
-        assert output in run_command(command, stdin, stdout, stderr)
+    channel.close()