From bc59486896e156a1bf69e311c25223108e38b794 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=80=97=E5=AD=90?= <i@haozi.net>
Date: Mon, 9 Sep 2024 12:30:15 +0800
Subject: [PATCH] feat: optimize code

---
 contracts/session/manager.go                     |  2 ++
 session/manager.go                               | 10 +++++++---
 session/manager_test.go                          |  7 ++++++-
 session/middleware/start_session.go              |  3 +--
 session/session.go                               | 16 ++++++++--------
 support/docker/{postgresql.go => postgres.go}    |  0
 .../{postgresql_test.go => postgres_test.go}     |  0
 translation/file_loader_test.go                  |  2 +-
 8 files changed, 25 insertions(+), 15 deletions(-)
 rename support/docker/{postgresql.go => postgres.go} (100%)
 rename support/docker/{postgresql_test.go => postgres_test.go} (100%)

diff --git a/contracts/session/manager.go b/contracts/session/manager.go
index 8c3d621b2..bf2dd9bfb 100644
--- a/contracts/session/manager.go
+++ b/contracts/session/manager.go
@@ -7,4 +7,6 @@ type Manager interface {
 	Driver(name ...string) (Driver, error)
 	// Extend extends the session manager with a custom driver.
 	Extend(driver string, handler func() Driver) error
+	// ReleaseSession releases the session back to the pool.
+	ReleaseSession(session Session)
 }
diff --git a/session/manager.go b/session/manager.go
index c820ce3f5..93be429d3 100644
--- a/session/manager.go
+++ b/session/manager.go
@@ -36,6 +36,9 @@ func NewManager(config config.Config, json foundation.Json) *Manager {
 }
 
 func (m *Manager) BuildSession(handler sessioncontract.Driver, sessionID ...string) sessioncontract.Session {
+	if handler == nil {
+		panic("session driver cannot be nil")
+	}
 	session := m.AcquireSession()
 	session.setDriver(handler)
 	session.setJson(m.json)
@@ -82,9 +85,10 @@ func (m *Manager) AcquireSession() *Session {
 	return session
 }
 
-func (m *Manager) ReleaseSession(session *Session) {
-	session.reset()
-	m.sessionPool.Put(session)
+func (m *Manager) ReleaseSession(session sessioncontract.Session) {
+	s := session.(*Session)
+	s.reset()
+	m.sessionPool.Put(s)
 }
 
 func (m *Manager) getDefaultDriver() string {
diff --git a/session/manager_test.go b/session/manager_test.go
index 1295f6d41..a45a5ffc1 100644
--- a/session/manager_test.go
+++ b/session/manager_test.go
@@ -88,8 +88,13 @@ func (s *ManagerTestSuite) TestExtend() {
 }
 
 func (s *ManagerTestSuite) TestBuildSession() {
+	driver, err := s.manager.Driver("file")
+	s.Nil(err)
+	s.NotNil(driver)
+	s.Equal("*driver.File", fmt.Sprintf("%T", driver))
+
 	s.mockConfig.On("GetString", "session.cookie").Return("test_cookie").Once()
-	session := s.manager.BuildSession(nil)
+	session := s.manager.BuildSession(driver)
 	s.NotNil(session)
 	s.Equal("test_cookie", session.GetName())
 }
diff --git a/session/middleware/start_session.go b/session/middleware/start_session.go
index 41d867a5a..9df9d46a0 100644
--- a/session/middleware/start_session.go
+++ b/session/middleware/start_session.go
@@ -55,7 +55,6 @@ func StartSession() http.Middleware {
 		}
 
 		// Release session
-		// TODO - any better way to release the session?
-		session.SessionFacade.(*session.Manager).ReleaseSession(s.(*session.Session))
+		session.SessionFacade.ReleaseSession(s)
 	}
 }
diff --git a/session/session.go b/session/session.go
index f1ad06287..7c796e09f 100644
--- a/session/session.go
+++ b/session/session.go
@@ -180,14 +180,6 @@ func (s *Session) SetName(name string) sessioncontract.Session {
 	return s
 }
 
-func (s *Session) setDriver(driver sessioncontract.Driver) {
-	s.driver = driver
-}
-
-func (s *Session) setJson(json foundation.Json) {
-	s.json = json
-}
-
 func (s *Session) Start() bool {
 	s.loadSession()
 
@@ -288,6 +280,14 @@ func (s *Session) reset() {
 	s.started = false
 }
 
+func (s *Session) setDriver(driver sessioncontract.Driver) {
+	s.driver = driver
+}
+
+func (s *Session) setJson(json foundation.Json) {
+	s.json = json
+}
+
 // toStringSlice converts an interface slice to a string slice.
 func toStringSlice(anySlice []any) []string {
 	strSlice := make([]string, len(anySlice))
diff --git a/support/docker/postgresql.go b/support/docker/postgres.go
similarity index 100%
rename from support/docker/postgresql.go
rename to support/docker/postgres.go
diff --git a/support/docker/postgresql_test.go b/support/docker/postgres_test.go
similarity index 100%
rename from support/docker/postgresql_test.go
rename to support/docker/postgres_test.go
diff --git a/translation/file_loader_test.go b/translation/file_loader_test.go
index dd9177f7b..12f01a595 100644
--- a/translation/file_loader_test.go
+++ b/translation/file_loader_test.go
@@ -34,7 +34,7 @@ func (f *FileLoaderTestSuite) SetupSuite() {
 }
 
 func (f *FileLoaderTestSuite) TearDownSuite() {
-	assert.Nil(f.T(), file.Remove("lang"))
+	f.Nil(file.Remove("lang"))
 }
 
 func (f *FileLoaderTestSuite) SetupTest() {