98 lines
3.1 KiB
Go
98 lines
3.1 KiB
Go
package jupyter_test
|
|
|
|
import (
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/jupyter"
|
|
)
|
|
|
|
func TestPrepareContainerConfig_PublicJupyter_RegistersKernelAndStartsNotebook(t *testing.T) {
|
|
oldEnv := os.Getenv("FETCHML_JUPYTER_CONDA_ENV")
|
|
oldKernel := os.Getenv("FETCHML_JUPYTER_KERNEL_NAME")
|
|
_ = os.Setenv("FETCHML_JUPYTER_CONDA_ENV", "base")
|
|
_ = os.Setenv("FETCHML_JUPYTER_KERNEL_NAME", "python")
|
|
t.Cleanup(func() {
|
|
_ = os.Setenv("FETCHML_JUPYTER_CONDA_ENV", oldEnv)
|
|
_ = os.Setenv("FETCHML_JUPYTER_KERNEL_NAME", oldKernel)
|
|
})
|
|
|
|
req := &jupyter.StartRequest{
|
|
Name: "test",
|
|
Workspace: "/tmp/ws",
|
|
Image: "quay.io/jupyter/base-notebook:latest",
|
|
Network: jupyter.NetworkConfig{
|
|
HostPort: 8888,
|
|
ContainerPort: 8888,
|
|
EnableToken: false,
|
|
EnablePassword: false,
|
|
},
|
|
Security: jupyter.SecurityConfig{AllowNetwork: true},
|
|
}
|
|
|
|
cfg := jupyter.PrepareContainerConfig("svc", req)
|
|
if len(cfg.Command) < 3 {
|
|
t.Fatalf("expected bootstrap command, got %#v", cfg.Command)
|
|
}
|
|
if cfg.Command[0] != "bash" || cfg.Command[1] != "-lc" {
|
|
t.Fatalf("expected bash -lc wrapper, got %#v", cfg.Command)
|
|
}
|
|
script := cfg.Command[2]
|
|
if !strings.Contains(script, "ipykernel install") {
|
|
t.Fatalf("expected ipykernel install in script, got %q", script)
|
|
}
|
|
if !strings.Contains(script, "start-notebook.sh") {
|
|
t.Fatalf("expected start-notebook.sh in script, got %q", script)
|
|
}
|
|
if !strings.Contains(script, "--ServerApp.token=") {
|
|
t.Fatalf("expected token disable flags in script, got %q", script)
|
|
}
|
|
}
|
|
|
|
func TestPrepareContainerConfig_MLToolsRunner_RegistersKernelAndStartsNotebook(t *testing.T) {
|
|
oldEnv := os.Getenv("FETCHML_JUPYTER_CONDA_ENV")
|
|
oldKernel := os.Getenv("FETCHML_JUPYTER_KERNEL_NAME")
|
|
_ = os.Setenv("FETCHML_JUPYTER_CONDA_ENV", "ml_env")
|
|
_ = os.Setenv("FETCHML_JUPYTER_KERNEL_NAME", "ml")
|
|
t.Cleanup(func() {
|
|
_ = os.Setenv("FETCHML_JUPYTER_CONDA_ENV", oldEnv)
|
|
_ = os.Setenv("FETCHML_JUPYTER_KERNEL_NAME", oldKernel)
|
|
})
|
|
|
|
req := &jupyter.StartRequest{
|
|
Name: "test",
|
|
Workspace: "/tmp/ws",
|
|
Image: "localhost/ml-tools-runner:latest",
|
|
Network: jupyter.NetworkConfig{
|
|
HostPort: 8888,
|
|
ContainerPort: 8888,
|
|
EnableToken: false,
|
|
},
|
|
Security: jupyter.SecurityConfig{AllowNetwork: true},
|
|
}
|
|
|
|
cfg := jupyter.PrepareContainerConfig("svc", req)
|
|
if len(cfg.Command) < 3 {
|
|
t.Fatalf("expected bootstrap command, got %#v", cfg.Command)
|
|
}
|
|
if cfg.Command[0] != "bash" || cfg.Command[1] != "-lc" {
|
|
t.Fatalf("expected bash -lc wrapper, got %#v", cfg.Command)
|
|
}
|
|
script := cfg.Command[2]
|
|
if !strings.Contains(script, "ipykernel install") {
|
|
t.Fatalf("expected ipykernel install in script, got %q", script)
|
|
}
|
|
if !strings.Contains(script, "jupyter notebook") {
|
|
t.Fatalf("expected jupyter notebook in script, got %q", script)
|
|
}
|
|
if !strings.Contains(script, "conda run -n ml_env") {
|
|
t.Fatalf("expected conda run to use ml_env in script, got %q", script)
|
|
}
|
|
if strings.Contains(script, "--ServerApp.token=") {
|
|
return
|
|
}
|
|
if !strings.Contains(script, "--NotebookApp.token=") {
|
|
t.Fatalf("expected token disable flags in script, got %q", script)
|
|
}
|
|
}
|