From e27b168ec609f7f6a8c5f341905358ad1410ee56 Mon Sep 17 00:00:00 2001
From: Jusong Yu <jusong.yeu@gmail.com>
Date: Fri, 13 Oct 2023 16:31:52 +0200
Subject: [PATCH] f-dd

---
 stack/base/before-notebook.d/10_prepare-home-config.sh | 6 +++++-
 stack/base/load-singlesshagent.sh                      | 8 ++++----
 tests/test-common.py                                   | 8 ++++++++
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/stack/base/before-notebook.d/10_prepare-home-config.sh b/stack/base/before-notebook.d/10_prepare-home-config.sh
index 567d6659..a4d0ab70 100644
--- a/stack/base/before-notebook.d/10_prepare-home-config.sh
+++ b/stack/base/before-notebook.d/10_prepare-home-config.sh
@@ -42,7 +42,11 @@ if [[ -f /opt/bin/load-singlesshagent.sh ]] && ! grep -q "${header}" /home/${NB_
 ${header}
 if [ -f /opt/bin/load-singlesshagent.sh ]; then
     source /opt/bin/load-singlesshagent.sh
-    ssh-add /home/${NB_USER}/.ssh/id_rsa 2> /dev/null
 fi
 EOF
 fi
+
+# load the ssh-agent and add the default key generated
+# the return code can be non-zero if the ssh-agent is not running
+# which will cause the notebook to fail to start so we need to ignore the return code
+source /opt/bin/load-singlesshagent.sh || true
diff --git a/stack/base/load-singlesshagent.sh b/stack/base/load-singlesshagent.sh
index c2e8f9f0..d7dc56cd 100644
--- a/stack/base/load-singlesshagent.sh
+++ b/stack/base/load-singlesshagent.sh
@@ -26,13 +26,13 @@ load_singlesshagent() {
     SSH_ENV="$HOME/.ssh/agent-environment"
     # Source SSH settings, if applicable
     if [ -r "${SSH_ENV}" ]; then
-        source "${SSH_ENV}" > /dev/null
+        source "${SSH_ENV}" 2> /dev/null
         [ "$VERBOSE" == "true" ] && echo "- sourcing existing environment" >&2
     else
         [ "$VERBOSE" == "true" ] && echo "- no existing environment to source" >&2
     fi
 
-    SSH_ADD_OUTPUT=`ssh-add -l 2> /dev/null`
+    SSH_ADD_OUTPUT=`ssh-add -l`
     # Needed, the later 'test' calls will replace this
     SSHADD_RETVAL="$?"
     # Error code: 0: there are keys; 1: there are no keys; 2: cannot contact agent
@@ -44,8 +44,8 @@ load_singlesshagent() {
     elif [ "$SSHADD_RETVAL" == "1" ]
     then
         [ "$VERBOSE" == "true" ] && echo "  - ssh-agent found (${SSH_AGENT_PID}), no keys (I might want to add keys here)" >&2
-        # Maybe the user wants to do it by himself?
-        #ssh-add
+        # run ssh-add to add the default generate key `id_rsa` to the agent
+        ssh-add ~/.ssh/id_rsa 2> /dev/null
     elif [ "$SSHADD_RETVAL" == "0" ]
     then
         NUMKEYS=`echo "$SSH_ADD_OUTPUT" | wc -l`
diff --git a/tests/test-common.py b/tests/test-common.py
index a5ce02b8..acc63d8b 100644
--- a/tests/test-common.py
+++ b/tests/test-common.py
@@ -11,3 +11,11 @@ def test_verdi_status(aiidalab_exec, nb_user):
     output = aiidalab_exec("verdi status", user=nb_user).decode().strip()
     assert "Connected to RabbitMQ" in output
     assert "Daemon is running" in output
+
+
+def test_ssh_agent_is_running(aiidalab_exec, nb_user):
+    output = aiidalab_exec("ps aux | grep ssh-agent", user=nb_user).decode().strip()
+    assert "ssh-agent" in output
+
+    # also check only one ssh-agent process is running
+    assert len(output.splitlines()) == 1