Repository: ncar-xdev/jupyter-forward (language: python)

Issue:
Support for more shells
1. besides `tcsh` and `bash`, #133 makes it clear we need to support `sh` as well
2. Can we figure out remote shell from `$SHELL` (and raise exception mandating use of `--shell` if that variable is undefined)
3. one possible path forward would be to create a parent class that contains wrappers for running commands and then shell-specific child classes with correct functionality (e.g. `sh -l -c "[run this command]"` for `sh` and `bash` but `sh -c "[run this command]"` for `tcsh`; also `'{command} > {self.log_file} 2>&1'` for `sh` and `bash` but `'{command} >& {self.log_file}'` for `tcsh`)
4. May also need template for expanding variables (e.g. `hostname` in `'jupyter lab --no-browser --ip=\$(hostname -f)'`)

Ideally this would let us set up simple tests launching commands other than `jupyter` so that we can use Github Actions to run these tests and make sure we are supporting as many shells as possible (and that future PRs don't break existing support)
Better documentation of --shell
Following the incorporation of #67 we are no longer including `shell` in `self.run_kwargs` but it is still used in the context of `--launch-command`: https://github.com/NCAR/jupyter-forward/blob/ea8a0ec3de63730e96c081a051a60680a44a7ce5/jupyter_forward/core.py#L175

I think there are two options:

1. Update the documentation of `shell` in `cli.py`: https://github.com/NCAR/jupyter-forward/blob/ea8a0ec3de63730e96c081a051a60680a44a7ce5/jupyter_forward/cli.py#L68
2. Come up with a shell-agnostic way to incorporate `--launch-command` and then drop `--shell` altogether


Discussion / hints:

@andersy005 -- why does `--launch-command` generate a shell script to run rather than just running the command directly? (I'm not familiar with `--launch-command`, but I think the answer is important for deciding between the two different approaches to addressing this issue)
> why does --launch-command generate a shell script to run rather than just running the command directly? 

As you may recall, to launch jupyter lab we retrieve the `hostname` or `ip` address information and pass this to jupyter lab command: 

https://github.com/NCAR/jupyter-forward/blob/20096779c4d562fcd096b6b09b91ccc8b0851146/jupyter_forward/core.py#L178-L182



I tried running the command directly. However, the hostname information I was getting back were for the login node instead of the compute node. For instance, if you were to run this command via `srun` or `qsub`:

```bash
jupyter lab --no-browser --ip=$(hostname) > $HOME/.jupyter_forward/log.txt 2>&1
```




The `$(hostname)` would evaluate to the hostname of the login node or whichever node the `qsub` or `srun` command was issued from. 

The simplest remedy for this issue was to generate a shell script, and then execute the script once we are on the compute node. My knowledge of batch schedulers is limited :). So, if you know how to get the direct command approach to work, please let me know. 


Another thing to keep in mind is that sometimes you need to configure the shell invocation arguments so as to **make the shell act as if it had been invoked as a login shell** (this is necessary in order to get the `conda activate ENV_NAME` command to work). I don't know how to do this in a clean and generic way via the `direct command` approach. For example, here's a sample script generated by jupyter-forward:

```bash
$ cat batch-script.2020-11-28T05-01-25
#!/bin/bash -l

conda activate playground && jupyter lab --no-browser --ip=$(hostname) > /glade/u/home/abanihi/.jupyter_forward/log.2020-11-28T05-01-25 2>&1
```

What would be its equivalent in the `direct command` approach? 

Fix patch:
diff --git a/.github/scripts/ssh_setup.sh b/.github/scripts/ssh_setup.sh
index aac44ba..c23673c 100755
--- a/.github/scripts/ssh_setup.sh
+++ b/.github/scripts/ssh_setup.sh
@@ -2,7 +2,7 @@
 
 ssh-keygen -t ed25519 -f ~/.ssh/my_key -N ''
 cat > ~/.ssh/config <<EOF
-Host eniac.local
+Host localhost.local
     User $USER
     HostName 127.0.0.1
     IdentityFile ~/.ssh/my_key
@@ -10,4 +10,4 @@ EOF
 echo -n 'from="127.0.0.1" ' | cat - ~/.ssh/my_key.pub > ~/.ssh/authorized_keys
 chmod og-rw ~
 ls -la ~/.ssh
-ssh -o 'StrictHostKeyChecking no' eniac.local id
+ssh -o 'StrictHostKeyChecking no' localhost.local id
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index c790d92..089ebe3 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -12,6 +12,12 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
   cancel-in-progress: true
 
+env:
+  JUPYTER_FORWARD_ENABLE_SSH_TESTS: true
+  JUPYTER_FORWARD_SSH_TEST_USER: runner
+  JUPYTER_FORWARD_SSH_TEST_HOSTNAME: localhost.local
+  JUPYTER_FORWARD_TEST_SHELLS: '["bash", "tcsh", "zsh", null]'
+
 jobs:
   build:
     name: python-${{ matrix.python-version }}
@@ -31,7 +37,13 @@ jobs:
       - name: Configure SSH key for localhost
         run: |
           .github/scripts/ssh_setup.sh
-      - uses: conda-incubator/setup-miniconda@v2
+
+      - name: Install missing shells e.g. zsh and tcsh
+        run: |
+          sudo apt-get update
+          sudo apt-get install zsh tcsh
+          cat /etc/shells
+      - uses: conda-incubator/setup-miniconda@master
         with:
           channels: conda-forge
           channel-priority: strict
diff --git a/.github/workflows/upstream-dev-ci.yaml b/.github/workflows/upstream-dev-ci.yaml
index 2f15511..6219d47 100644
--- a/.github/workflows/upstream-dev-ci.yaml
+++ b/.github/workflows/upstream-dev-ci.yaml
@@ -11,6 +11,12 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
   cancel-in-progress: true
 
+env:
+  JUPYTER_FORWARD_ENABLE_SSH_TESTS: true
+  JUPYTER_FORWARD_SSH_TEST_USER: runner
+  JUPYTER_FORWARD_SSH_TEST_HOSTNAME: localhost.local
+  JUPYTER_FORWARD_TEST_SHELLS: '["bash", "tcsh", "zsh", null]'
+
 jobs:
   upstream-dev:
     name: upstream-dev
@@ -29,6 +35,11 @@ jobs:
       - name: Configure SSH key for localhost
         run: |
           .github/scripts/ssh_setup.sh
+      - name: Install missing shells e.g. zsh and tcsh
+        run: |
+          sudo apt-get update
+          sudo apt-get install zsh tcsh
+          cat /etc/shells
       - uses: conda-incubator/setup-miniconda@v2
         id: conda
         with:
diff --git a/.gitignore b/.gitignore
index ededb36..637f3b7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -107,3 +107,4 @@ venv.bak/
 .vscode/
 .vagrant/
 vagrant-ssh
+.envrc
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6a52efe..315f52d 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -16,11 +16,12 @@ repos:
     hooks:
       - id: pyupgrade
         args:
-          - '--py37-plus'
+          - '--py38-plus'
 
   - repo: https://github.com/psf/black
     rev: 22.1.0
     hooks:
+      - id: black
       - id: black-jupyter
 
   - repo: https://github.com/keewis/blackdoc
diff --git a/ci/environment-upstream-dev.yml b/ci/environment-upstream-dev.yml
index f4e39fc..90c9859 100644
--- a/ci/environment-upstream-dev.yml
+++ b/ci/environment-upstream-dev.yml
@@ -3,6 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - fabric
+  - jupyterlab
   - numpydoc
   - paramiko
   - pip
diff --git a/ci/environment.yml b/ci/environment.yml
index 8bf1b6c..f58c9f1 100644
--- a/ci/environment.yml
+++ b/ci/environment.yml
@@ -2,8 +2,8 @@ name: jupyter-forward-dev
 channels:
   - conda-forge
 dependencies:
-  - codecov
   - fabric
+  - jupyterlab
   - numpydoc
   - paramiko
   - pip
diff --git a/jupyter_forward/cli.py b/jupyter_forward/cli.py
index f5bb638..07ac03a 100644
--- a/jupyter_forward/cli.py
+++ b/jupyter_forward/cli.py
@@ -1,5 +1,5 @@
+import typing
 from pathlib import Path
-from typing import Optional
 
 import typer
 
@@ -66,13 +66,13 @@ def start(
         ),
     ),
     shell: str = typer.Option(
-        '/usr/bin/sh -l',
+        None,
         '--shell',
         '-s',
         show_default=True,
         help='Which remote shell binary to use.',
     ),
-    version: Optional[bool] = typer.Option(
+    version: typing.Optional[bool] = typer.Option(
         None,
         '--version',
         callback=version_callback,
diff --git a/jupyter_forward/console.py b/jupyter_forward/console.py
new file mode 100644
index 0000000..70df95b
--- /dev/null
+++ b/jupyter_forward/console.py
@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+import rich.console
+
+console = rich.console.Console()
diff --git a/jupyter_forward/core.py b/jupyter_forward/core.py
index 0387a0c..5b6e952 100644
--- a/jupyter_forward/core.py
+++ b/jupyter_forward/core.py
@@ -1,29 +1,24 @@
+from __future__ import annotations
+
+import dataclasses
 import datetime
 import getpass
 import pathlib
 import socket
 import sys
 import time
-from dataclasses import dataclass
 
 import invoke
 import paramiko
 from fabric import Connection
-from rich.console import Console
-
-timestamp = datetime.datetime.now().strftime('%Y-%m-%dT%H-%M-%S')
 
-console = Console()
+from .console import console
+from .helpers import _authentication_handler, is_port_available, open_browser, parse_stdout
 
-
-def _authentication_handler(title, instructions, prompt_list):
-    """
-    Handler for paramiko auth_interactive_dumb
-    """
-    return [getpass.getpass(str(pr[0])) for pr in prompt_list]
+timestamp = datetime.datetime.now().strftime('%Y-%m-%dT%H-%M-%S')
 
 
-@dataclass
+@dataclasses.dataclass
 class RemoteRunner:
     """Starts Jupyter lab on a remote resource and port forwards session to
     local machine.
@@ -47,7 +42,7 @@ class RemoteRunner:
     port_forwarding: bool = True
     launch_command: str = None
     identity: str = None
-    shell: str = '/usr/bin/sh -l'
+    shell: str = None
 
     def __post_init__(self):
         if self.notebook_dir is not None and self.notebook is not None:
@@ -56,12 +51,17 @@ def __post_init__(self):
             self.notebook = pathlib.Path(self.notebook)
             self.notebook_dir = str(self.notebook.parent)
             self.notebook = self.notebook.name
-        console.rule('[bold green]Authentication', characters='*')
+
         if self.port_forwarding and not is_port_available(self.port):
             console.print(
-                f'''[bold red]Specified port={self.port} is already in use on your local machine. Try a different port'''
+                f'''[bold red]:x: Specified port={self.port} is already in use on your local machine. Try a different port'''
             )
             sys.exit(1)
+        self._authenticate()
+        self._check_shell()
+
+    def _authenticate(self):
+        console.rule('[bold green]Authenticating', characters='*')
 
         connect_kwargs = {}
         if self.identity:
@@ -100,13 +100,23 @@ def __post_init__(self):
                     console.log('[bold red]:x: Failed to Authenticate your connection')
         if not self.session.is_connected:
             sys.exit(1)
-
         console.print('[bold cyan]:white_check_mark: The client is authenticated successfully')
 
+    def _check_shell(self):
+        console.rule('[bold green]Verifying shell location', characters='*')
+        if self.shell is None:
+            shell = self.session.run('echo $SHELL || echo $0', hide='out').stdout.strip()
+            if not shell:
+                raise ValueError('Could not determine shell. Please specify one using --shell.')
+            self.shell = shell
+        else:
+            # Get the full path to the shell in case the user specified a shell name
+            self.shell = self.run_command(f'which {self.shell}').stdout.strip()
+        console.print(f'[bold cyan]:white_check_mark: Using shell: {self.shell}')
+
     def run_command(
         self,
         command,
-        force_login_shell=False,
         exit=True,
         warn=True,
         pty=True,
@@ -114,12 +124,13 @@ def run_command(
         echo=True,
         **kwargs,
     ):
-        command = f'sh -l -c "{command}"' if force_login_shell else command
+        if 'csh' in self.shell:
+            command = f'''{self.shell} -c "{command}"'''
+        else:
+            command = f'''{self.shell} -lc "{command}"'''
         out = self.session.run(command, warn=warn, pty=pty, hide=hide, echo=echo, **kwargs)
-        if out.failed:
-            console.print(f'[bold red] {out.stderr}')
-            if exit:
-                sys.exit(1)
+        if out.failed and exit:
+            sys.exit(1)
         return out
 
     def setup_port_forwarding(self):
@@ -155,51 +166,71 @@ def start(self):
         try:
             self._launch_jupyter()
         except Exception as exc:
-            console.print(f'[bold red] {exc}')
+            console.print(f'[bold red]:x: {exc}')
             self.close()
         finally:
             console.rule(
-                '[bold red]Terminated the network 📡 connection to the remote end', characters='*'
+                f'[bold red]:x: Terminated the network 📡 connection to {self.session.host}',
+                characters='*',
             )
 
+    def _get_hostname(self):
+        if self.launch_command:
+            return r'\$(hostname -f)'
+        else:
+            return self.session.run('hostname -f').stdout.strip()
+
     def _launch_jupyter(self):
+        conda_activate_cmd = self._conda_activate_cmd()
+        self._set_log_directory()
+        self._set_log_file()
+        command = rf'jupyter lab --no-browser --ip={self._get_hostname()}'
+        if self.notebook_dir:
+            command = f'{command} --notebook-dir={self.notebook_dir}'
+        command = self._generate_redirect_command(command=command, log_file=self.log_file)
+        if self.conda_env:
+            command = f'{conda_activate_cmd} {self.conda_env} && {command}'
+
+        if self.launch_command:
+            command = f'{self.launch_command} {self._prepare_batch_job_script(command)}'
+
+        console.rule('[bold green]Launching Jupyter Lab', characters='*')
+        self.session.run(f'{self.shell} -c "{command}"', asynchronous=True, pty=True, echo=True)
+        self.parsed_result = self._parse_log_file()
 
+        if self.port_forwarding:
+            self.setup_port_forwarding()
+        else:
+            open_browser(url=self.parsed_result['url'], path=self.notebook)
+            self.run_command(command=f'tail -f {self.log_file}')
+
+    def _generate_redirect_command(self, *, log_file: str, command: str) -> str:
+        if 'csh' in self.shell:
+            return f'{command} >& {log_file}'
+        else:
+            return f'{command} > {log_file} 2>&1'
+
+    def _conda_activate_cmd(self):
         console.rule(
-            '[bold green]Running jupyter sanity checks (ensuring `jupyter` is in `$PATH`)',
+            '[bold green]Running jupyter sanity checks',
             characters='*',
         )
-        check_jupyter_status = 'sh -l -c "command -v jupyter"'
+        check_jupyter_status = 'which jupyter'
         conda_activate_cmd = 'source activate'
         if self.conda_env:
             try:
                 self.run_command(f'{conda_activate_cmd} {self.conda_env} && {check_jupyter_status}')
             except SystemExit:
-                console.print(f'`{conda_activate_cmd}` failed. Trying `conda activate`...')
+                console.print(
+                    f'[bold red]:x: `{conda_activate_cmd}` failed. Trying `conda activate`...'
+                )
                 self.run_command(f'conda activate {self.conda_env} && {check_jupyter_status}')
                 conda_activate_cmd = 'conda activate'
-
         else:
             self.run_command(check_jupyter_status)
-        self._set_log_directory()
-        self._set_log_file()
-        command = r'jupyter lab --no-browser --ip=\$(hostname -f)'
-        if self.notebook_dir:
-            command = f'{command} --notebook-dir={self.notebook_dir}'
-        command = f'{command} > {self.log_file} 2>&1'
-        if self.conda_env:
-            command = f'{conda_activate_cmd} {self.conda_env} && {command}'
-
-        if self.launch_command:
-            console.rule('[bold green]Preparing Batch Job script', characters='*')
-            script_file = f'{self.log_dir}/batch-script.{timestamp}'
-            cmd = f"""echo "#!{self.shell}\n\n{command}" > {script_file}"""
-            self.run_command(command=cmd)
-            self.run_command(command=f'chmod +x {script_file}')
-            command = f'{self.launch_command} {script_file}'
-
-        console.rule('[bold green]Launching Jupyter Lab', characters='*')
-        self.session.run(f'sh -l -c "{command}"', asynchronous=True, pty=True, echo=True)
+        return conda_activate_cmd
 
+    def _parse_log_file(self):
         # wait for logfile to contain access info, then write it to screen
         condition = True
         stdout = None
@@ -207,22 +238,32 @@ def _launch_jupyter(self):
             f'[bold cyan]Parsing {self.log_file} log file on {self.session.host} for jupyter information',
             spinner='weather',
         ):
-            pattern = 'is running at:'
+            # TODO: Ensure this loop doesn't run forever if the log file is not found or empty
             while condition:
                 try:
                     result = self.run_command(f'cat {self.log_file}', echo=False, hide='out')
-                    if pattern in result.stdout:
+                    if 'is running at:' in result.stdout.strip():
                         condition = False
                         stdout = result.stdout
                 except invoke.exceptions.UnexpectedExit:
                     pass
-        self.parsed_result = parse_stdout(stdout)
-
-        if self.port_forwarding:
-            self.setup_port_forwarding()
-        else:
-            open_browser(url=self.parsed_result['url'], path=self.notebook)
-            self.run_command(command=f'tail -f {self.log_file}')
+        return parse_stdout(stdout)
+
+    def _prepare_batch_job_script(self, command):
+        console.rule('[bold green]Preparing Batch Job script', characters='*')
+        script_file = f'{self.log_dir}/batch_job_script_{timestamp}'
+        shell = self.shell
+        if 'csh' not in shell:
+            shell = f'{shell} -l'
+        for command in [
+            f"echo -n '#!' > {script_file}",
+            f'echo {shell} >> {script_file}',
+            f"echo '{command}' >> {script_file}",
+            f'chmod +x {script_file}',
+        ]:
+            self.run_command(command=command, exit=True)
+        console.print(f'[bold cyan]:white_check_mark: Batch Job script resides in {script_file}')
+        return script_file
 
     def _set_log_file(self):
         log_file = f'{self.log_dir}/log_{timestamp}.txt'
@@ -233,103 +274,32 @@ def _set_log_file(self):
 
     def _set_log_directory(self):
         def _check_log_file_dir(directory):
-            check_dir_command = f'touch {directory}/foobar && rm -rf {directory}/foobar && echo "{directory} is WRITABLE" || echo "{directory} is NOT WRITABLE"'
+            check_dir_command = f"touch {directory}/foobar && rm -rf {directory}/foobar && echo '{directory} is WRITABLE' || echo '{directory} is NOT WRITABLE'"
             _tmp_dir_status = self.run_command(command=check_dir_command, exit=False)
             return directory if 'is WRITABLE' in _tmp_dir_status.stdout.strip() else None
 
         console.rule(f'[bold green] Creating log file on {self.session.host}', characters='*')
+        # TODO: Allow users to override this via a `--log-dir`
         log_dir = None
+        # Try TMPDIR first if defined
         tmp_dir_env_status = self.run_command(command='printenv TMPDIR', exit=False)
-        home_dir_env_status = self.run_command(command='printenv HOME', exit=False)
         if not tmp_dir_env_status.failed:
             log_dir = _check_log_file_dir('$TMPDIR')
-        elif not home_dir_env_status.failed:
-            log_dir = _check_log_file_dir('$HOME')
         else:
-            tmp_dir_error_message = '$TMPDIR is not defined'
-            home_dir_error_message = '$HOME is not defined'
-            console.print(
-                f'[bold red]Can not determine directory for log file:\n{home_dir_error_message}\n{tmp_dir_error_message}'
-            )
-            sys.exit(1)
+            # Try HOME if TMPDIR is not defined
+            home_dir_env_status = self.run_command(command='printenv HOME', exit=False)
+            if not home_dir_env_status.failed:
+                log_dir = _check_log_file_dir('$HOME')
+            else:
+                # Raise an error if neither TMPDIR or HOME are defined
+                tmp_dir_error_message = '$TMPDIR is not defined'
+                home_dir_error_message = '$HOME is not defined'
+                console.print(
+                    f'[bold red]:x: Can not determine directory for log file:\n{home_dir_error_message}\n{tmp_dir_error_message}'
+                )
+                sys.exit(1)
         log_dir = f'{log_dir}/.jupyter_forward'
         self.run_command(command=f'mkdir -p {log_dir}')
         console.print(f'[bold cyan]:white_check_mark: Log directory is set to {log_dir}')
         self.log_dir = log_dir
         return self
-
-
-def open_browser(port: int = None, token: str = None, url: str = None, path=None):
-    """Opens notebook interface in a new browser window.
-
-    Parameters
-    ----------
-    port : int, optional
-        Port number to use, by default None
-    token : str, optional
-        token used for authentication, by default None
-    url : str, optional
-        Notebook url, by default None
-    path : str, optional
-        Notebook path
-
-    Raises
-    ------
-    ValueError
-        If url is None and port is None
-    """
-
-    import webbrowser
-
-    if not url:
-        if port is None:
-            raise ValueError('Please specify port number to use.')
-        url = f'http://localhost:{port}'
-        if token:
-            url = f'{url}/?token={token}'
-        url = f'{url}/lab/tree/{path}' if path else url
-
-    console.rule('[bold green]Opening Jupyter Lab interface in a browser', characters='*')
-    console.print(f'Jupyter Lab URL: {url}')
-    console.rule('[bold green]', characters='*')
-    webbrowser.open(url, new=2)
-
-
-def is_port_available(port):
-    socket_for_port_check = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    status = socket_for_port_check.connect_ex(('localhost', int(port)))
-    return status != 0
-
-
-def parse_stdout(stdout: str):
-    """Parses stdout to determine remote_hostname, port, token, url
-
-    Parameters
-    ----------
-    stdout : str
-        Contents of the log file/stdout
-
-    Returns
-    -------
-    dict
-        A dictionary containing hotname, port, token, and url
-    """
-    import re
-    import urllib.parse
-
-    hostname, port, token, url = None, None, None, None
-    urls = set(
-        re.findall(
-            r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+',
-            stdout,
-        )
-    )
-    for url in urls:
-        url = url.strip()
-        if '127.0.0.1' not in url:
-            result = urllib.parse.urlparse(url)
-            hostname, port = result.netloc.split(':')
-            if 'token' in result.query:
-                token = result.query.split('token=')[-1].strip()
-            break
-    return {'hostname': hostname, 'port': port, 'token': token, 'url': url}
diff --git a/jupyter_forward/helpers.py b/jupyter_forward/helpers.py
new file mode 100644
index 0000000..6ce392c
--- /dev/null
+++ b/jupyter_forward/helpers.py
@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import getpass
+import re
+import socket
+import urllib.parse
+
+from .console import console
+
+
+def open_browser(port: int = None, token: str = None, url: str = None, path=None) -> None:
+    """Opens notebook interface in a new browser window.
+
+    Parameters
+    ----------
+    port : int, optional
+        Port number to use, by default None
+    token : str, optional
+        token used for authentication, by default None
+    url : str, optional
+        Notebook url, by default None
+    path : str, optional
+        Notebook path
+
+    Raises
+    ------
+    ValueError
+        If url is None and port is None
+    """
+
+    import webbrowser
+
+    if not url:
+        if port is None:
+            raise ValueError('Please specify port number to use.')
+        url = f'http://localhost:{port}'
+        if token:
+            url = f'{url}/?token={token}'
+        url = f'{url}/lab/tree/{path}' if path else url
+
+    console.rule('[bold green]Opening Jupyter Lab interface in a browser', characters='*')
+    console.print(f'Jupyter Lab URL: {url}')
+    console.rule('[bold green]', characters='*')
+    webbrowser.open(url, new=2)
+
+
+def is_port_available(port) -> bool:
+    socket_for_port_check = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    status = socket_for_port_check.connect_ex(('localhost', int(port)))
+    return status != 0
+
+
+def parse_stdout(stdout: str) -> dict[str, str]:
+    """Parses stdout to determine remote_hostname, port, token, url
+
+    Parameters
+    ----------
+    stdout : str
+        Contents of the log file/stdout
+
+    Returns
+    -------
+    dict
+        A dictionary containing hotname, port, token, and url
+    """
+
+    hostname, port, token, url = None, None, None, None
+    urls = set(
+        re.findall(
+            r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+',
+            stdout,
+        )
+    )
+    for url in urls:
+        url = url.strip()
+        if '127.0.0.1' not in url:
+            result = urllib.parse.urlparse(url)
+            hostname, port = result.netloc.split(':')
+            if 'token' in result.query:
+                token = result.query.split('token=')[-1].strip()
+            break
+    return {'hostname': hostname, 'port': port, 'token': token, 'url': url}
+
+
+def _authentication_handler(title, instructions, prompt_list):
+    """
+    Handler for paramiko auth_interactive_dumb
+    """
+    return [getpass.getpass(str(pr[0])) for pr in prompt_list]
diff --git a/setup.cfg b/setup.cfg
index 2643827..f01ee48 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -20,4 +20,4 @@ skip=
 
 [tool:pytest]
 console_output_style = count
-addopts = --cov=./ --cov-report=xml --verbose -s
+addopts = --cov=./ --cov-report=xml --verbose -s -rx


Test patch:
diff --git a/tests/misc.py b/tests/misc.py
new file mode 100644
index 0000000..31a2422
--- /dev/null
+++ b/tests/misc.py
@@ -0,0 +1,14 @@
+sample_log_file_contents = [
+    '[I 15:46:27.590 LabApp] JupyterLab extension loaded from /jupyterlab',
+    '[I 15:46:27.590 LabApp] JupyterLab application directory is /jupyter/lab',
+    '[I 15:46:27.594 LabApp] Serving notebooks from local directory: /glade',
+    '[I 15:46:27.594 LabApp] The Jupyter Notebook is running at:',
+    '[I 15:46:27.594 LabApp] http://eniac01:59628/?token=Loremipsumdolorsitamet',
+    '[I 15:46:27.594 LabApp]  or http://127.0.0.1:59628/?token=Loremipsumdolorsitamet',
+    '[I 15:46:27.594 LabApp] Use Control-C to stop this server',
+    '[C 15:46:27.604 LabApp] To access the notebook, open this file in a browser:',
+    'file:///.local/share/jupyter/runtime/nbserver-14905-open.html',
+    'Or copy and paste one of these URLs:',
+    'http://eniac01:59628/?token=Loremipsumdolorsitamet',
+    'or http://127.0.0.1:59628/?token=Loremipsumdolorsitamet',
+]
diff --git a/tests/test_cli.py b/tests/test_cli.py
new file mode 100644
index 0000000..83cf510
--- /dev/null
+++ b/tests/test_cli.py
@@ -0,0 +1,17 @@
+import typer.testing
+
+from jupyter_forward.cli import app
+
+runner = typer.testing.CliRunner()
+
+
+def test_help():
+    result = runner.invoke(app, ['--help'])
+    assert result.exit_code == 0
+    assert 'show this message and exit'.lower() in result.stdout.lower()
+
+
+def test_version():
+    result = runner.invoke(app, ['--version'])
+    assert result.exit_code == 0
+    assert 'Jupyter Forward CLI Version'.lower() in result.stdout.lower()
diff --git a/tests/test_core.py b/tests/test_core.py
index 10ab21e..5eb9818 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -1,119 +1,180 @@
 import datetime
+import json
 import os
-import socket
-from unittest import mock as mock
 
 import pytest
 
 import jupyter_forward
-from jupyter_forward.core import is_port_available, open_browser, parse_stdout
 
-NOT_GITHUB_ACTIONS = os.environ.get('GITHUB_ACTIONS') is None
-requires_gha = pytest.mark.skipif(NOT_GITHUB_ACTIONS, reason='requires GITHUB_ACTIONS')
+from .misc import sample_log_file_contents
 
+SHELLS = json.loads(os.environ.get('JUPYTER_FORWARD_TEST_SHELLS', '["bash", null]'))
+JUPYTER_FORWARD_ENABLE_SSH_TESTS = os.environ.get('JUPYTER_FORWARD_ENABLE_SSH_TESTS') is None
+requires_ssh = pytest.mark.skipif(JUPYTER_FORWARD_ENABLE_SSH_TESTS, reason='SSH tests disabled')
+ON_GITHUB_ACTIONS = os.environ.get('GITHUB_ACTIONS') is not None
 
-@pytest.fixture(scope='session', params=[f"{os.environ['USER']}@eniac.local"])
+
+@pytest.fixture(scope='package')
 def runner(request):
-    remote = jupyter_forward.RemoteRunner(request.param)
+    remote = jupyter_forward.RemoteRunner(
+        f"{os.environ['JUPYTER_FORWARD_SSH_TEST_USER']}@{os.environ['JUPYTER_FORWARD_SSH_TEST_HOSTNAME']}",
+        shell=request.param,
+    )
     yield remote
     remote.close()
 
 
+@requires_ssh
 @pytest.mark.parametrize(
-    'stdout, expected',
+    'port, conda_env, notebook, notebook_dir, port_forwarding, identity, shell',
     [
-        (
-            """[I 15:46:27.590 LabApp] JupyterLab extension loaded from /jupyterlab\n
-            [I 15:46:27.590 LabApp] JupyterLab application directory is /jupyter/lab\n
-            [I 15:46:27.594 LabApp] Serving notebooks from local directory: /glade\n
-            [I 15:46:27.594 LabApp] The Jupyter Notebook is running at:\n
-            [I 15:46:27.594 LabApp] http://eniac01:59628/?token=Loremipsumdolorsitamet\n
-            [I 15:46:27.594 LabApp]  or http://127.0.0.1:59628/?token=Loremipsumdolorsitamet\n
-            [I 15:46:27.594 LabApp] Use Control-C to stop this server\n
-            [C 15:46:27.604 LabApp]\n\n
-            To access the notebook, open this file in a browser:\n
-            file:///.local/share/jupyter/runtime/nbserver-14905-open.html\n
-            Or copy and paste one of these URLs:\n
-            http://eniac01:59628/?token=Loremipsumdolorsitamet\n
-            or http://127.0.0.1:59628/?token=Loremipsumdolorsitamet\n     """,
-            {
-                'hostname': 'eniac01',
-                'port': '59628',
-                'token': 'Loremipsumdolorsitamet',
-                'url': 'http://eniac01:59628/?token=Loremipsumdolorsitamet',
-            },
-        ),
-        ('', {'hostname': None, 'port': None, 'token': None, 'url': None}),
+        (8888, None, None, None, True, None, None),
+        (8888, None, None, '~/notebooks/', False, None, None),
+        (8888, None, '~/my_notebook.ipynb', None, True, None, 'bash'),
+        (8888, 'base', None, None, False, None, 'bash'),
     ],
 )
-def test_parse_stdout(stdout, expected):
-    parsed_results = parse_stdout(stdout)
-    assert parsed_results == expected
-
-
-@pytest.mark.parametrize('port', [8888, 9999])
-def test_is_port_available(
-    port,
-):
-    @mock.create_autospec
-    def connect_ex(self, address):
-        # if address[1] == 8888:
-        return 0
-
-    with mock.patch.object(socket.socket, 'connect_ex', connect_ex) as m:
-        is_port_available(port)
-        m.assert_called_once()
+def test_runner_init(port, conda_env, notebook, notebook_dir, port_forwarding, identity, shell):
+    remote_runner = jupyter_forward.RemoteRunner(
+        f"{os.environ['JUPYTER_FORWARD_SSH_TEST_USER']}@{os.environ['JUPYTER_FORWARD_SSH_TEST_HOSTNAME']}",
+        port=port,
+        conda_env=conda_env,
+        notebook=notebook,
+        notebook_dir=notebook_dir,
+        identity=identity,
+        port_forwarding=port_forwarding,
+        shell=shell,
+    )
+
+    assert remote_runner.port == port
+    assert remote_runner.conda_env == conda_env
+
+
+@requires_ssh
+def test_runner_init_notebook_dir_error():
+    with pytest.raises(ValueError):
+        jupyter_forward.RemoteRunner(
+            f"{os.environ['JUPYTER_FORWARD_SSH_TEST_USER']}@{os.environ['JUPYTER_FORWARD_SSH_TEST_HOSTNAME']}",
+            notebook_dir='~/notebooks/',
+            notebook='~/my_notebook.ipynb',
+        )
 
 
-def test_open_browser_exception():
-    with pytest.raises(ValueError):
-        open_browser(token='ssh')
+@requires_ssh
+def test_runner_init_port_unavailable():
+    with pytest.raises(SystemExit):
+        jupyter_forward.RemoteRunner(
+            f"{os.environ['JUPYTER_FORWARD_SSH_TEST_USER']}@{os.environ['JUPYTER_FORWARD_SSH_TEST_HOSTNAME']}",
+            port=22,
+        )
 
 
-@pytest.mark.parametrize(
-    'port, token, url, expected',
-    [
-        (9999, 'ssh', None, 'http://localhost:9999/?token=ssh'),
-        (None, None, 'http://localhost:9999', 'http://localhost:9999'),
-    ],
-)
-def test_open_browser(port, token, url, expected):
-    with mock.patch('webbrowser.open') as mockwebopen:
-        open_browser(port, token, url)
-        mockwebopen.assert_called_once_with(expected, new=2)
+@requires_ssh
+def test_runner_authentication_error():
+    with pytest.raises(SystemExit):
+        jupyter_forward.RemoteRunner(
+            f"foobar@{os.environ['JUPYTER_FORWARD_SSH_TEST_HOSTNAME']}",
+        )
 
 
-@requires_gha
+@requires_ssh
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
 def test_connection(runner):
-    USER = os.environ['USER']
+    USER = os.environ['JUPYTER_FORWARD_SSH_TEST_USER']
     assert runner.session.is_connected
-    assert runner.session.host == '127.0.0.1'
+    assert runner.session.host in [
+        '127.0.0.1',
+        'localhost',
+        {os.environ['JUPYTER_FORWARD_SSH_TEST_HOSTNAME']},
+    ]
     assert runner.session.user == USER
 
 
-@requires_gha
+@requires_ssh
 @pytest.mark.parametrize('command', ['echo $HOME'])
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
 def test_run_command(runner, command):
     out = runner.run_command(command)
     assert not out.failed
-    assert out.stdout.strip() == f"{os.environ['HOME']}"
+    f"{os.environ['HOME']}" in out.stdout.strip()
 
 
-@requires_gha
+@requires_ssh
 @pytest.mark.parametrize('command', ['echod $HOME'])
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
 def test_run_command_failure(runner, command):
     out = runner.run_command(command, exit=False)
     assert out.failed
-    assert 'echod: command not found' in out.stdout.strip()
+    assert 'not found' in out.stdout.strip().lower()
 
     with pytest.raises(SystemExit):
         runner.run_command(command)
 
 
-@requires_gha
+@requires_ssh
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
 def test_set_logs(runner):
     runner._set_log_directory()
-    assert runner.log_dir == '$HOME/.jupyter_forward'
+    assert '/.jupyter_forward' in runner.log_dir
     runner._set_log_file()
     now = datetime.datetime.now()
     assert f"log_{now.strftime('%Y-%m-%dT%H')}" in runner.log_file
+
+
+@requires_ssh
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
+def test_prepare_batch_job_script(runner):
+    if ON_GITHUB_ACTIONS and ('csh' in runner.shell):
+        pytest.xfail('Fails on GitHub Actions due to inconsistent shell behavior')
+    runner._set_log_directory()
+    script_file = runner._prepare_batch_job_script('echo hello world')
+    assert 'batch_job_script' in script_file
+    assert 'hello world' in runner.run_command(f'cat {script_file}').stdout.strip()
+
+
+@requires_ssh
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
+def test_parse_log_file(runner):
+    runner._set_log_directory()
+    runner._set_log_file()
+    runner.run_command(f"echo '{sample_log_file_contents[0]}' > {runner.log_file}")
+    for line in sample_log_file_contents[1:]:
+        runner.run_command(f"echo '{line}' >> {runner.log_file}")
+    out = runner._parse_log_file()
+    assert out == {
+        'hostname': 'eniac01',
+        'port': '59628',
+        'token': 'Loremipsumdolorsitamet',
+        'url': 'http://eniac01:59628/?token=Loremipsumdolorsitamet',
+    }
+
+
+@requires_ssh
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
+@pytest.mark.parametrize('environment', ['jupyter-forward-dev', None])
+def test_conda_activate_cmd(runner, environment):
+    if ON_GITHUB_ACTIONS and ('csh' in runner.shell or 'zsh' in runner.shell):
+        pytest.xfail('Fails on GitHub Actions due to inconsistent shell behavior')
+    runner.conda_env = environment
+    cmd = runner._conda_activate_cmd()
+    assert cmd in ['source activate', 'conda activate']
+
+
+@requires_ssh
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
+def test_conda_activate_cmd_error(runner):
+    runner.conda_env = 'DOES_NOT_EXIST'
+    with pytest.raises(SystemExit):
+        runner._conda_activate_cmd()
+
+
+@requires_ssh
+@pytest.mark.parametrize('runner', SHELLS, indirect=True)
+def test_generate_redirect_cmd(runner):
+    runner._set_log_directory()
+    runner._set_log_file()
+    cmd = runner._generate_redirect_command(command='echo "hello world"', log_file=runner.log_file)
+    if 'csh' in runner.shell:
+        assert cmd.endswith(runner.log_file)
+    else:
+        assert cmd.endswith(f'{runner.log_file} 2>&1')
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
new file mode 100644
index 0000000..7b05e7c
--- /dev/null
+++ b/tests/test_helpers.py
@@ -0,0 +1,61 @@
+import socket
+from unittest import mock
+
+import pytest
+
+import jupyter_forward.helpers
+
+from .misc import sample_log_file_contents
+
+
+@pytest.mark.parametrize(
+    'stdout, expected',
+    [
+        (
+            sample_log_file_contents,
+            {
+                'hostname': 'eniac01',
+                'port': '59628',
+                'token': 'Loremipsumdolorsitamet',
+                'url': 'http://eniac01:59628/?token=Loremipsumdolorsitamet',
+            },
+        ),
+        ('', {'hostname': None, 'port': None, 'token': None, 'url': None}),
+    ],
+)
+def test_parse_stdout(stdout, expected):
+    stdout = '\n'.join(stdout)
+    parsed_results = jupyter_forward.helpers.parse_stdout(stdout)
+    assert parsed_results == expected
+
+
+@pytest.mark.parametrize('port', [8888, 9999])
+def test_is_port_available(
+    port,
+):
+    @mock.create_autospec
+    def connect_ex(self, address):
+        # if address[1] == 8888:
+        return 0
+
+    with mock.patch.object(socket.socket, 'connect_ex', connect_ex) as m:
+        jupyter_forward.helpers.is_port_available(port)
+        m.assert_called_once()
+
+
+def test_open_browser_exception():
+    with pytest.raises(ValueError):
+        jupyter_forward.helpers.open_browser(token='ssh')
+
+
+@pytest.mark.parametrize(
+    'port, token, url, expected',
+    [
+        (9999, 'ssh', None, 'http://localhost:9999/?token=ssh'),
+        (None, None, 'http://localhost:9999', 'http://localhost:9999'),
+    ],
+)
+def test_open_browser(port, token, url, expected):
+    with mock.patch('webbrowser.open') as mockwebopen:
+        jupyter_forward.helpers.open_browser(port, token, url)
+        mockwebopen.assert_called_once_with(expected, new=2)