From d2f92923ee9242ecef0ebca55b38218fe2fc2d4a Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Wed, 26 Mar 2025 19:06:51 +0800
Subject: [PATCH] feat: add PullImages for ADBDriver

---
 internal/version/VERSION   |  2 +-
 uixt/android_driver_adb.go | 69 ++++++++++++++++++++++++++++++++++++++
 uixt/android_test.go       |  3 ++
 uixt/browser_driver.go     |  4 +++
 uixt/driver.go             |  1 +
 uixt/harmony_driver_hdc.go |  5 +++
 uixt/ios_driver_wda.go     |  5 +++
 7 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/internal/version/VERSION b/internal/version/VERSION
index 3211bfbf..883634e4 100644
--- a/internal/version/VERSION
+++ b/internal/version/VERSION
@@ -1 +1 @@
-v5.0.0-beta-2503251433
+v5.0.0-beta-2503261906
diff --git a/uixt/android_driver_adb.go b/uixt/android_driver_adb.go
index ed5c7ab6..893621b9 100644
--- a/uixt/android_driver_adb.go
+++ b/uixt/android_driver_adb.go
@@ -961,6 +961,75 @@ func (ad *ADBDriver) ClearImages() error {
 	return nil
 }
 
+// PullImages pulls all images from device's DCIM/Camera directory to local directory
+func (ad *ADBDriver) PullImages(localDir string) error {
+	log.Info().Str("localDir", localDir).Msg("ADBDriver.PullImages")
+	remoteDir := "/sdcard/DCIM/Camera/"
+
+	// create local directory if not exists
+	if err := os.MkdirAll(localDir, 0o755); err != nil {
+		return fmt.Errorf("failed to create local directory: %w", err)
+	}
+
+	files, err := ad.Device.List(remoteDir)
+	if err != nil {
+		return fmt.Errorf("failed to list directory %s: %w", remoteDir, err)
+	}
+
+	for _, file := range files {
+		// filter image files by extension
+		ext := strings.ToLower(path.Ext(file.Name))
+		if !isImageFile(ext) {
+			continue
+		}
+
+		remotePath := path.Join(remoteDir, file.Name)
+		localPath := path.Join(localDir, file.Name)
+
+		// check if file already exists
+		if _, err := os.Stat(localPath); err == nil {
+			log.Debug().Str("localPath", localPath).Msg("file already exists, skipping")
+			continue
+		}
+
+		// create local file
+		f, err := os.Create(localPath)
+		if err != nil {
+			log.Error().Err(err).Str("localPath", localPath).Msg("failed to create local file")
+			continue
+		}
+		defer f.Close()
+
+		// pull image file
+		if err := ad.Device.Pull(remotePath, f); err != nil {
+			log.Error().Err(err).
+				Str("remotePath", remotePath).
+				Str("localPath", localPath).
+				Msg("failed to pull image")
+			continue // continue with next file
+		}
+		log.Info().
+			Str("remotePath", remotePath).
+			Str("localPath", localPath).
+			Msg("image pulled successfully")
+	}
+	return nil
+}
+
+// isImageFile checks if the file extension is an image format
+func isImageFile(ext string) bool {
+	imageExts := map[string]bool{
+		".jpg":  true,
+		".jpeg": true,
+		".png":  true,
+		".gif":  true,
+		".bmp":  true,
+		".webp": true,
+		".heic": true,
+	}
+	return imageExts[ext]
+}
+
 type ExportPoint struct {
 	Start     int         `json:"start" yaml:"start"`
 	End       int         `json:"end" yaml:"end"`
diff --git a/uixt/android_test.go b/uixt/android_test.go
index cdcdcd88..6e8076ad 100644
--- a/uixt/android_test.go
+++ b/uixt/android_test.go
@@ -260,6 +260,9 @@ func TestDriver_ADB_PushImage(t *testing.T) {
 	err = driver.PushImage(path)
 	assert.Nil(t, err)
 
+	err = driver.PullImages("./test")
+	assert.Nil(t, err)
+
 	err = driver.ClearImages()
 	assert.Nil(t, err)
 }
diff --git a/uixt/browser_driver.go b/uixt/browser_driver.go
index 7a455529..75bd9fc1 100644
--- a/uixt/browser_driver.go
+++ b/uixt/browser_driver.go
@@ -452,6 +452,10 @@ func (wd *BrowserDriver) PushImage(localPath string) error {
 	return errors.New("not support")
 }
 
+func (wd *BrowserDriver) PullImages(localDir string) error {
+	return errors.New("not support")
+}
+
 func (wd *BrowserDriver) Orientation() (orientation types.Orientation, err error) {
 	log.Warn().Msg("Orientation not implemented in ADBDriver")
 	return
diff --git a/uixt/driver.go b/uixt/driver.go
index 70accc51..897bbc76 100644
--- a/uixt/driver.go
+++ b/uixt/driver.go
@@ -68,6 +68,7 @@ type IDriver interface {
 
 	// image related
 	PushImage(localPath string) error
+	PullImages(localDir string) error
 	ClearImages() error
 
 	// triggers the log capture and returns the log entries
diff --git a/uixt/harmony_driver_hdc.go b/uixt/harmony_driver_hdc.go
index 1541245f..bec2ba83 100644
--- a/uixt/harmony_driver_hdc.go
+++ b/uixt/harmony_driver_hdc.go
@@ -288,6 +288,11 @@ func (hd *HDCDriver) PushImage(localPath string) error {
 	return nil
 }
 
+func (hd *HDCDriver) PullImages(localDir string) error {
+	log.Warn().Msg("PullImages not implemented in HDCDriver")
+	return nil
+}
+
 func (hd *HDCDriver) ClearImages() error {
 	log.Warn().Msg("ClearImages not implemented in HDCDriver")
 	return nil
diff --git a/uixt/ios_driver_wda.go b/uixt/ios_driver_wda.go
index 0c56aa66..8f5e9897 100644
--- a/uixt/ios_driver_wda.go
+++ b/uixt/ios_driver_wda.go
@@ -989,6 +989,11 @@ func (wd *WDADriver) PushImage(localPath string) error {
 	return err
 }
 
+func (wd *WDADriver) PullImages(localDir string) error {
+	log.Warn().Msg("PullImages not implemented in WDADriver")
+	return nil
+}
+
 func (wd *WDADriver) ClearImages() error {
 	log.Info().Msg("WDADriver.ClearImages")
 	data := map[string]interface{}{}