chore(test): Add test cases for shell and home manager (#163)

* chore: Add more test cases

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Fix

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

Makefile

@@ -91,7 +91,7 @@ export GOFLAGS ?= -count=1
 #
 
 # All targets.
-.PHONY: lint test build container push addlicense debug debug-local build-local generate clean
+.PHONY: lint test build container push addlicense debug debug-local build-local generate clean test-local
 
 build: build-local
 
@@ -131,8 +131,11 @@ debug-local:
 addlicense:
 	addlicense -c "The envd Authors" **/*.go **/**/*.go **/**/**/*.go
 
-test: generate
+test-local:
 	@go test -v -race -coverprofile=coverage.out ./...
+
+test: generate 
+	@go test -race -coverprofile=coverage.out ./...
 	@go tool cover -func coverage.out | tail -n 1 | awk '{ print "Total coverage: " $$3 }'
 
 clean:

cmd/envd/build_test.go

@@ -21,6 +21,7 @@ import (
 	. "github.com/onsi/gomega"
 
 	"github.com/tensorchord/envd/pkg/docker"
+	"github.com/tensorchord/envd/pkg/home"
 )
 
 var _ = Describe("build command", func() {
@@ -29,6 +30,7 @@ var _ = Describe("build command", func() {
 		"envd.test", "--debug", "build", "--path", buildContext,
 	}
 	BeforeEach(func() {
+		Expect(home.Initialize()).NotTo(HaveOccurred())
 		cli, err := docker.NewClient(context.TODO())
 		Expect(err).NotTo(HaveOccurred())
 		_ = cli.Destroy(context.TODO(), buildContext)

cmd/envd/up_test.go

@@ -21,6 +21,7 @@ import (
 	. "github.com/onsi/gomega"
 
 	"github.com/tensorchord/envd/pkg/docker"
+	"github.com/tensorchord/envd/pkg/home"
 )
 
 var _ = Describe("up command", func() {
@@ -29,6 +30,7 @@ var _ = Describe("up command", func() {
 		"envd.test", "--debug", "up", "--path", buildContext, "--detach",
 	}
 	BeforeEach(func() {
+		Expect(home.Initialize()).NotTo(HaveOccurred())
 		cli, err := docker.NewClient(context.TODO())
 		Expect(err).NotTo(HaveOccurred())
 		_ = cli.Destroy(context.TODO(), buildContext)

pkg/editor/jupyter/util_test.go

@@ -0,0 +1,71 @@
+package jupyter
+
+import (
+	"testing"
+
+	"github.com/tensorchord/envd/pkg/lang/ir"
+)
+
+func TestGenerateCommand(t *testing.T) {
+	testcases := []struct {
+		graph    ir.Graph
+		dir      string
+		expected []string
+	}{
+		{
+			graph: ir.Graph{
+				JupyterConfig: &ir.JupyterConfig{
+					Password: "",
+					Port:     8888,
+				},
+			},
+			dir: "test",
+			expected: []string{
+				"jupyter", "notebook", "--allow-root",
+				"--ip", "0.0.0.0", "--notebook-dir", "test",
+				"--NotebookApp.password", "''", "--NotebookApp.token", "''",
+				"--port", "8888",
+			},
+		},
+		{
+			graph: ir.Graph{
+				JupyterConfig: &ir.JupyterConfig{
+					Password: "test",
+					Port:     8888,
+				},
+			},
+			dir: "test",
+			expected: []string{
+				"jupyter", "notebook", "--allow-root",
+				"--ip", "0.0.0.0", "--notebook-dir", "test",
+				"--NotebookApp.password", "test", "--NotebookApp.token", "''",
+				"--port", "8888",
+			},
+		},
+		{
+			graph:    ir.Graph{},
+			dir:      "test",
+			expected: []string{},
+		},
+	}
+	for _, tc := range testcases {
+		actual := GenerateCommand(tc.graph, tc.dir)
+		if !equal(actual, tc.expected) {
+			t.Errorf("failed to generate the command: expected %v, got %v", tc.expected, actual)
+		}
+	}
+}
+
+// Equal tells whether a and b contain the same elements.
+// A nil argument is equivalent to an empty slice.
+func equal(a, b []string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i, v := range a {
+		if v != b[i] {
+			return false
+		}
+	}
+	return true
+}

pkg/home/manager_suite_test.go

@@ -0,0 +1,13 @@
+package home
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestManager(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Home Manager Suite")
+}

pkg/home/manager_test.go

@@ -0,0 +1,44 @@
+package home
+
+import (
+	"path/filepath"
+
+	"github.com/adrg/xdg"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"github.com/tensorchord/envd/pkg/util/fileutil"
+)
+
+var _ = Describe("home manager", func() {
+	BeforeEach(func() {
+		// Cleanup the home cache.
+		Expect(Initialize()).NotTo(HaveOccurred())
+		m := GetManager()
+		Expect(fileutil.RemoveAll(m.(*generalManager).cacheStatusFile)).NotTo(HaveOccurred())
+	})
+	AfterEach(func() {
+		// Cleanup the home cache.
+		Expect(Initialize()).NotTo(HaveOccurred())
+		m := GetManager()
+		Expect(fileutil.RemoveAll(m.(*generalManager).cacheStatusFile)).NotTo(HaveOccurred())
+	})
+	When("initialized", func() {
+		It("should initialized successfully", func() {
+			Expect(Initialize()).NotTo(HaveOccurred())
+			m := GetManager()
+			Expect(m.CacheDir()).To(Equal(filepath.Join(xdg.CacheHome, "envd")))
+			Expect(m.ConfigFile()).To(Equal(filepath.Join(xdg.ConfigHome, "envd/config.envd")))
+		})
+		It("should return the cache status", func() {
+			Expect(Initialize()).NotTo(HaveOccurred())
+			m := GetManager()
+			Expect(m.Cached("test")).To(BeFalse())
+			Expect(m.MarkCache("test", true)).To(Succeed())
+			Expect(m.Cached("test")).To(BeTrue())
+			// Restart the init process, the cache should be persistent.
+			Expect(Initialize()).NotTo(HaveOccurred())
+			m = GetManager()
+			Expect(m.Cached("test")).To(BeTrue())
+		})
+	})
+})

pkg/lang/ir/graph.go

@@ -95,7 +95,7 @@ func (g Graph) Compile() (llb.State, error) {
 	diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell"))
 	pypiStage := llb.Diff(builtinSystemStage, g.compilePyPIPackages(builtinSystemStage), llb.WithCustomName("install PyPI packages"))
 	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage), llb.WithCustomName("install system packages"))
-	sshStage := g.copyenvdSSHServer()
+	sshStage := g.copyEnvdSSHServer()
 
 	vscodeStage, err := g.compileVSCode()
 	if err != nil {
@@ -233,7 +233,7 @@ func (g Graph) compileSystemPackages(root llb.State) llb.State {
 	return run.Root()
 }
 
-func (g Graph) copyenvdSSHServer() llb.State {
+func (g Graph) copyEnvdSSHServer() llb.State {
 	// TODO(gaocegege): Remove global var ssh image.
 	run := llb.Image(viper.GetString(flag.FlagSSHImage)).
 		File(llb.Copy(llb.Image(viper.GetString(flag.FlagSSHImage)),

pkg/shell/zsh.go

@@ -25,6 +25,10 @@ import (
 	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
+const (
+	cacheKey = "oh-my-zsh"
+)
+
 //go:embed install.sh
 var installScript string
 
@@ -46,7 +50,7 @@ func (m generalManager) InstallScript() string {
 }
 
 func (m generalManager) DownloadOrCache() (bool, error) {
-	if home.GetManager().Cached("oh-my-zsh") {
+	if home.GetManager().Cached(cacheKey) {
 		logrus.WithFields(logrus.Fields{
 			"cache-dir": m.OHMyZSHDir(),
 		}).Debug("oh-my-zsh already exists in cache")
@@ -71,7 +75,7 @@ func (m generalManager) DownloadOrCache() (bool, error) {
 		return false, err
 	}
 
-	if err := home.GetManager().MarkCache("oh-my-zsh", true); err != nil {
+	if err := home.GetManager().MarkCache(cacheKey, true); err != nil {
 		return false, errors.Wrap(err, "failed to update cache status")
 	}
 	l.Debug("oh-my-zsh is downloaded")

pkg/shell/zsh_suite_test.go

@@ -0,0 +1,13 @@
+package shell
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestZSH(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "ZSH Suite")
+}

pkg/shell/zsh_test.go

@@ -0,0 +1,44 @@
+package shell
+
+import (
+	"path/filepath"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"github.com/tensorchord/envd/pkg/home"
+	"github.com/tensorchord/envd/pkg/util/fileutil"
+)
+
+var _ = Describe("zsh manager", func() {
+	zshManager := NewManager()
+	BeforeEach(func() {
+		Expect(home.Initialize()).NotTo(HaveOccurred())
+	})
+	AfterEach(func() {
+		// Cleanup the home cache.
+		Expect(home.Initialize()).NotTo(HaveOccurred())
+		m := home.GetManager()
+		Expect(fileutil.RemoveAll(m.CacheDir())).NotTo(HaveOccurred())
+	})
+	When("cached", func() {
+		It("should skip", func() {
+			err := home.GetManager().MarkCache(cacheKey, true)
+			Expect(err).NotTo(HaveOccurred())
+			cached, err := zshManager.DownloadOrCache()
+			Expect(cached).To(BeTrue())
+			Expect(err).NotTo(HaveOccurred())
+		})
+	})
+	When("not cached", func() {
+		It("should download", func() {
+			err := home.GetManager().MarkCache(cacheKey, false)
+			Expect(err).NotTo(HaveOccurred())
+			cached, err := zshManager.DownloadOrCache()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(cached).To(BeFalse())
+			exists, err := fileutil.DirExists(filepath.Join(home.GetManager().CacheDir(), "oh-my-zsh"))
+			Expect(err).NotTo(HaveOccurred())
+			Expect(exists).To(BeTrue())
+		})
+	})
+})