11import re
22import string
3+ import time
34import uuid
45
56import paramiko
1415TIMEOUT_SECS = 300
1516
1617
17- @pytest .fixture (scope = "function " )
18+ @pytest .fixture (scope = "session " )
1819def paramiko_object (jupyterhub_access_token ):
19- """Connects to JupyterHub ssh cluster from outside the cluster."""
20+ """Connects to JupyterHub SSH cluster from outside the cluster.
21+
22+ Ensures the JupyterLab pod is ready before attempting reauthentication
23+ by setting both `auth_timeout` and `banner_timeout` appropriately,
24+ and by retrying the connection until the pod is ready or a timeout occurs.
25+ """
2026 params = {
2127 "hostname" : constants .NEBARI_HOSTNAME ,
2228 "port" : 8022 ,
2329 "username" : constants .KEYCLOAK_USERNAME ,
2430 "password" : jupyterhub_access_token ,
2531 "allow_agent" : constants .PARAMIKO_SSH_ALLOW_AGENT ,
2632 "look_for_keys" : constants .PARAMIKO_SSH_LOOK_FOR_KEYS ,
27- "auth_timeout" : 5 * 60 ,
2833 }
2934
3035 ssh_client = paramiko .SSHClient ()
3136 ssh_client .set_missing_host_key_policy (paramiko .AutoAddPolicy ())
32- try :
33- ssh_client .connect (** params )
34- yield ssh_client
35- finally :
36- ssh_client .close ()
37-
38-
39- def run_command (command , stdin , stdout , stderr ):
40- delimiter = uuid .uuid4 ().hex
41- stdin .write (f"echo { delimiter } { command } { delimiter } \n " )
42-
43- output = []
44-
45- line = stdout .readline ()
46- while not re .match (f"^{ delimiter } , line .strip ()):
47- line = stdout .readline ()
4837
49- line = stdout .readline ()
50- if delimiter not in line :
51- output .append (line )
52-
53- while not re .match (f"^{ delimiter } , line .strip ()):
54- line = stdout .readline ()
55- if delimiter not in line :
56- output .append (line )
57-
58- return "" .join (output ).strip ()
59-
60-
61- @pytest .mark .timeout (TIMEOUT_SECS )
62- @pytest .mark .filterwarnings ("ignore::urllib3.exceptions.InsecureRequestWarning" )
63- @pytest .mark .filterwarnings ("ignore::ResourceWarning" )
64- def test_simple_jupyterhub_ssh (paramiko_object ):
65- stdin , stdout , stderr = paramiko_object .exec_command ("" )
38+ yield ssh_client , params
39+
40+ ssh_client .close ()
41+
42+
43+ def invoke_shell (
44+ client : paramiko .SSHClient , params : dict [str , any ]
45+ ) -> paramiko .Channel :
46+ client .connect (** params )
47+ return client .invoke_shell ()
48+
49+
50+ def extract_output (delimiter : str , output : str ) -> str :
51+ # Extract the command output between the start and end delimiters
52+ match = re .search (rf"{ delimiter } { delimiter } , output , re .DOTALL )
53+ if match :
54+ print (match .group (1 ).strip ())
55+ return match .group (1 ).strip ()
56+ else :
57+ return output .strip ()
58+
59+
60+ def run_command_list (
61+ commands : list [str ], channel : paramiko .Channel , wait_time : int = 0
62+ ) -> dict [str , str ]:
63+ command_delimiters = {}
64+ for command in commands :
65+ delimiter = uuid .uuid4 ().hex
66+ command_delimiters [command ] = delimiter
67+ b = channel .send (f"echo { delimiter } { command } { delimiter } \n " )
68+ if b == 0 :
69+ print (f"Command '{ command } )
70+ # Wait for the output to be ready before reading
71+ time .sleep (wait_time )
72+ while not channel .recv_ready ():
73+ time .sleep (1 )
74+ print ("Waiting for output" )
75+ output = ""
76+ while channel .recv_ready ():
77+ output += channel .recv (65535 ).decode ("utf-8" )
78+ outputs = {}
79+ for command , delimiter in command_delimiters .items ():
80+ command_output = extract_output (delimiter , output )
81+ outputs [command ] = command_output
82+ return outputs
6683
6784
6885@pytest .mark .timeout (TIMEOUT_SECS )
6986@pytest .mark .filterwarnings ("ignore::urllib3.exceptions.InsecureRequestWarning" )
7087@pytest .mark .filterwarnings ("ignore::ResourceWarning" )
7188def test_print_jupyterhub_ssh (paramiko_object ):
72- stdin , stdout , stderr = paramiko_object . exec_command ( "" )
73-
74- # commands to run and just print the output
89+ client , params = paramiko_object
90+ channel = invoke_shell ( client , params )
91+ # Commands to run and just print the output
7592 commands_print = [
7693 "id" ,
7794 "env" ,
@@ -80,52 +97,60 @@ def test_print_jupyterhub_ssh(paramiko_object):
8097 "ls -la" ,
8198 "umask" ,
8299 ]
83-
84- for command in commands_print :
85- print (f'COMMAND: "{ command } )
86- print (run_command (command , stdin , stdout , stderr ))
100+ outputs = run_command_list (commands_print , channel )
101+ for command , output in outputs .items ():
102+ print (f"COMMAND: { command } )
103+ print (f"OUTPUT: { output } )
104+ channel .close ()
87105
88106
89107@pytest .mark .timeout (TIMEOUT_SECS )
90108@pytest .mark .filterwarnings ("ignore::urllib3.exceptions.InsecureRequestWarning" )
91109@pytest .mark .filterwarnings ("ignore::ResourceWarning" )
92110def test_exact_jupyterhub_ssh (paramiko_object ):
93- stdin , stdout , stderr = paramiko_object .exec_command ("" )
94-
95- # commands to run and exactly match output
96- commands_exact = [
97- ("id -u" , "1000" ),
98- ("id -g" , "100" ),
99- ("whoami" , constants .KEYCLOAK_USERNAME ),
100- ("pwd" , f"/home/{ constants .KEYCLOAK_USERNAME } ),
101- ("echo $HOME" , f"/home/{ constants .KEYCLOAK_USERNAME } ),
102- ("conda activate default && echo $CONDA_PREFIX" , "/opt/conda/envs/default" ),
103- (
104- "hostname" ,
105- f"jupyter-{ escape_string (constants .KEYCLOAK_USERNAME , safe = set (string .ascii_lowercase + string .digits ), escape_char = '-' ).lower ()} ,
106- ),
107- ]
111+ client , params = paramiko_object
112+ channel = invoke_shell (client , params )
113+ # Commands to run and exactly match output
114+ commands_exact = {
115+ "id -u" : "1000" ,
116+ "id -g" : "100" ,
117+ "whoami" : constants .KEYCLOAK_USERNAME ,
118+ "pwd" : f"/home/{ constants .KEYCLOAK_USERNAME } ,
119+ "echo $HOME" : f"/home/{ constants .KEYCLOAK_USERNAME } ,
120+ "conda activate default && echo $CONDA_PREFIX" : "/opt/conda/envs/default" ,
121+ "hostname" : f"jupyter-{ escape_string (constants .KEYCLOAK_USERNAME , safe = set (string .ascii_lowercase + string .digits ), escape_char = '-' ).lower ()} ,
122+ }
123+ outputs = run_command_list (list (commands_exact .keys ()), channel )
124+ for command , output in outputs .items ():
125+ assert (
126+ output == outputs [command ]
127+ ), f"Command '{ command } { outputs [command ]} { output }
108128
109- for command , output in commands_exact :
110- assert output == run_command (command , stdin , stdout , stderr )
129+ channel .close ()
111130
112131
113132@pytest .mark .timeout (TIMEOUT_SECS )
114133@pytest .mark .filterwarnings ("ignore::urllib3.exceptions.InsecureRequestWarning" )
115134@pytest .mark .filterwarnings ("ignore::ResourceWarning" )
116135def test_contains_jupyterhub_ssh (paramiko_object ):
117- stdin , stdout , stderr = paramiko_object .exec_command ("" )
118-
119- # commands to run and string need to be contained in output
120- commands_contain = [
121- ("ls -la" , ".bashrc" ),
122- ("cat ~/.bashrc" , "Managed by Nebari" ),
123- ("cat ~/.profile" , "Managed by Nebari" ),
124- ("cat ~/.bash_logout" , "Managed by Nebari" ),
125- # ensure we don't copy over extra files from /etc/skel in init container
126- ("ls -la ~/..202*" , "No such file or directory" ),
127- ("ls -la ~/..data" , "No such file or directory" ),
128- ]
136+ client , params = paramiko_object
137+ channel = invoke_shell (client , params )
138+
139+ # Commands to run and check if the output contains specific strings
140+ commands_contain = {
141+ "ls -la" : ".bashrc" ,
142+ "cat ~/.bashrc" : "Managed by Nebari" ,
143+ "cat ~/.profile" : "Managed by Nebari" ,
144+ "cat ~/.bash_logout" : "Managed by Nebari" ,
145+ # Ensure we don't copy over extra files from /etc/skel in init container
146+ "ls -la ~/..202*" : "No such file or directory" ,
147+ "ls -la ~/..data" : "No such file or directory" ,
148+ }
149+
150+ outputs = run_command_list (commands_contain .keys (), channel , 30 )
151+ for command , expected_output in commands_contain .items ():
152+ assert (
153+ expected_output in outputs [command ]
154+ ), f"Command '{ command } { expected_output } { outputs [command ]}
129155
130- for command , output in commands_contain :
131- assert output in run_command (command , stdin , stdout , stderr )
156+ channel .close ()
0 commit comments