From 2a6a115a7ce793a3a036dc6fbc2d704df766cead Mon Sep 17 00:00:00 2001
From: Hugo Hromic <hhromic@users.noreply.github.com>
Date: Sat, 6 Jan 2024 20:13:21 +0000
Subject: [PATCH] [FAST_BUILD] Fix Docker healthcheck when using custom runtime
 dirs (#2074)

* Fix Docker healthcheck when using custom runtime dirs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use a writable directory for healthcheck test

* Allow missing import for `jupyter_core` in mypy

* Set HOME according to NB_USER in healthcheck script

* Add custom runtime dir an NB_USER case to healthcheck test

* Call `jupyter --runtime-dir` directly in healthcheck script

* Update docker_healthcheck.py

* Update docker_healthcheck.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ayaz Salikhov <mathbunnyru@users.noreply.github.com>
---
 images/base-notebook/docker_healthcheck.py | 15 ++++++++++++++-
 tests/base-notebook/test_healthcheck.py    | 10 ++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/images/base-notebook/docker_healthcheck.py b/images/base-notebook/docker_healthcheck.py
index 8f3338e8..7dd3de02 100755
--- a/images/base-notebook/docker_healthcheck.py
+++ b/images/base-notebook/docker_healthcheck.py
@@ -3,6 +3,7 @@
 # Distributed under the terms of the Modified BSD License.
 import json
 import os
+import subprocess
 from pathlib import Path
 
 import requests
@@ -10,7 +11,19 @@ import requests
 # Several operations below deliberately don't check for possible errors
 # As this is a healthcheck, it should succeed or raise an exception on error
 
-runtime_dir = Path("/home/") / os.environ["NB_USER"] / ".local/share/jupyter/runtime/"
+# Docker runs healtchecks using an exec
+# It uses the default user configured when running the image: root for the case of a custom NB_USER or jovyan for the case of the default image user.
+# We manually change HOME to make `jupyter --runtime-dir` report a correct path
+# More information: <https://github.com/jupyter/docker-stacks/pull/2074#issuecomment-1879778409>
+result = subprocess.run(
+    ["jupyter", "--runtime-dir"],
+    check=True,
+    capture_output=True,
+    text=True,
+    env=dict(os.environ) | {"HOME": "/home/" + os.environ["NB_USER"]},
+)
+runtime_dir = Path(result.stdout.rstrip())
+
 json_file = next(runtime_dir.glob("*server-*.json"))
 
 url = json.loads(json_file.read_bytes())["url"]
diff --git a/tests/base-notebook/test_healthcheck.py b/tests/base-notebook/test_healthcheck.py
index 50d83c27..450a0111 100644
--- a/tests/base-notebook/test_healthcheck.py
+++ b/tests/base-notebook/test_healthcheck.py
@@ -42,6 +42,16 @@ LOGGER = logging.getLogger(__name__)
             ["start-notebook.py", "--ServerApp.base_url=/test"],
             "root",
         ),
+        (["JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime"], ["start-notebook.sh"], None),
+        (
+            [
+                "NB_USER=testuser",
+                "CHOWN_HOME=1",
+                "JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime",
+            ],
+            ["start-notebook.sh"],
+            "root",
+        ),
     ],
 )
 def test_health(