From 3b509ba6c6ffe40ef5df979e02a17d4958fd34ed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BD=99=E6=B3=93=E9=93=AE?= <yuhongzheng@bytedance.com>
Date: Thu, 20 Feb 2025 21:39:15 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E4=BC=98=E5=8C=96driver=20teardown?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/app.go         | 2 +-
 server/ext/context.go | 2 +-
 server/main.go        | 6 +++---
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/server/app.go b/server/app.go
index 9c98b9a1..ac45cef5 100644
--- a/server/app.go
+++ b/server/app.go
@@ -59,7 +59,7 @@ func (r *Router) clearAppHandler(c *gin.Context) {
 	if err != nil {
 		return
 	}
-	err = driver.GetIDriver().(*uixt.ADBDriver).AppClear(appClearReq.PackageName)
+	err = driver.GetIDriver().AppClear(appClearReq.PackageName)
 	if err != nil {
 		RenderError(c, err)
 		return
diff --git a/server/ext/context.go b/server/ext/context.go
index 3677984e..73bb7771 100644
--- a/server/ext/context.go
+++ b/server/ext/context.go
@@ -37,8 +37,8 @@ func (p RouterBaseMethodExt) GetDriver(c *gin.Context) (driverExt uixt.IXTDriver
 		server.RenderErrorInitDriver(c, err)
 		return
 	}
-	c.Set("driver", driver)
 	driverExt = driver_ext.NewXTDriver(driver,
 		ai.WithCVService(ai.CVServiceTypeVEDEM))
+	c.Set("driver", driverExt)
 	return driverExt, nil
 }
diff --git a/server/main.go b/server/main.go
index 2fabe924..9fef05a0 100644
--- a/server/main.go
+++ b/server/main.go
@@ -32,7 +32,7 @@ type IRouterBaseMethod interface {
 }
 
 func (r *Router) Init() {
-	r.Engine.Use(teardown())
+	r.Engine.Use(r.teardown())
 	r.Engine.GET("/ping", r.pingHandler)
 	r.Engine.GET("/", r.pingHandler)
 	r.Engine.POST("/", r.pingHandler)
@@ -91,7 +91,7 @@ func (r *Router) pingHandler(c *gin.Context) {
 	RenderSuccess(c, true)
 }
 
-func teardown() gin.HandlerFunc {
+func (r *Router) teardown() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		logID := c.Request.Header.Get("x-tt-logid")
 		startTime := time.Now()
@@ -107,7 +107,7 @@ func teardown() gin.HandlerFunc {
 
 		driverObj, exists := c.Get("driver")
 		if exists {
-			if driver, ok := driverObj.(*uixt.XTDriver); ok {
+			if driver, ok := driverObj.(uixt.IXTDriver); ok {
 				_ = driver.TearDown()
 			}
 		}