fix(home): Remove afterall to avoid flaky tests (#215)

* fix: Remove afterall Signed-off-by: Ce Gao <cegao@tensorchord.ai> * fix: Update Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · May 30, 2022 · cd2c0db · cd2c0db
1 parent b8a5516
commit cd2c0db
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 18 deletions.
diff --git a/pkg/home/manager.go b/pkg/home/manager.go
@@ -137,11 +137,6 @@ func (m *generalManager) init() error {
 		}
 	}
 
-	// Generate SSH keys when init
-	if err := sshconfig.GenerateKeys(); err != nil {
-		return errors.Wrap(err, "failed to generate ssh key")
-	}
-
 	file, err := os.Open(m.cacheStatusFile)
 	if err != nil {
 		return errors.Wrap(err, "failed to open cache status file")
@@ -152,6 +147,11 @@ func (m *generalManager) init() error {
 		return errors.Wrap(err, "failed to decode cache map")
 	}
 
+	// Generate SSH keys when init
+	if err := sshconfig.GenerateKeys(); err != nil {
+		return errors.Wrap(err, "failed to generate ssh key")
+	}
+
 	m.logger = logrus.WithFields(logrus.Fields{
 		"cache-dir":    m.cacheDir,
 		"config":       m.configFile,

diff --git a/pkg/home/manager_test.go b/pkg/home/manager_test.go
@@ -15,27 +15,15 @@
 package home
 
 import (
+	"os"
 	"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())
@@ -44,8 +32,10 @@ var _ = Describe("home manager", func() {
 			Expect(m.ConfigFile()).To(Equal(filepath.Join(xdg.ConfigHome, "envd/config.envd")))
 		})
 		It("should return the cache status", func() {
+			Expect(os.RemoveAll(filepath.Join(xdg.CacheHome, "envd/cache.status"))).NotTo(HaveOccurred())
 			Expect(Initialize()).NotTo(HaveOccurred())
 			m := GetManager()
+			m.(*generalManager).cacheMap = make(map[string]bool)
 			Expect(m.Cached("test")).To(BeFalse())
 			Expect(m.MarkCache("test", true)).To(Succeed())
 			Expect(m.Cached("test")).To(BeTrue())