Merge branch 'push_image' into 'master'

feat: 支持添加图片/视频文件

See merge request iesqa/httprunner!153

uixt/driver_session.go

@@ -265,8 +265,10 @@ func (s *DriverSession) Request(method string, urlStr string, rawBody []byte, op
 			logger = log.Debug().Bool("success", true)
 		}
 
-		logger = logger.Str("logid", logid).Str("request_method", method).Str("request_url", rawURL).
-			Str("request_body", string(rawBody))
+		logger = logger.Str("logid", logid).Str("request_method", method).Str("request_url", rawURL)
+		if len(rawBody) < 1024 {
+			logger = logger.Str("request_body", string(rawBody))
+		}
 		if !driverResult.RequestTime.IsZero() {
 			logger = logger.Int64("request_time", driverResult.RequestTime.UnixMilli())
 		}

uixt/image_utils.go

@@ -0,0 +1,116 @@
+package uixt
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/rs/zerolog/log"
+)
+
+// DetectAndRenameImageFile examines the file content to determine its media type
+// and renames the file with the appropriate extension (.jpg, .png, .mp4, etc.)
+func DetectAndRenameMediaFile(filePath string) (string, error) {
+	// Open the file
+	file, err := os.Open(filePath)
+	if err != nil {
+		return "", fmt.Errorf("failed to open file for type detection: %v", err)
+	}
+	defer file.Close()
+
+	// Read the first 512 bytes to detect content type
+	buffer := make([]byte, 512)
+	_, err = file.Read(buffer)
+	if err != nil && err != io.EOF {
+		return "", fmt.Errorf("failed to read file for type detection: %v", err)
+	}
+
+	// Reset file pointer
+	_, err = file.Seek(0, 0)
+	if err != nil {
+		return "", fmt.Errorf("failed to reset file pointer: %v", err)
+	}
+
+	// Detect content type
+	contentType := http.DetectContentType(buffer)
+	log.Info().Str("filePath", filePath).Str("contentType", contentType).Msg("Detected content type")
+
+	// Determine file extension based on content type
+	var extension string
+	switch {
+	// Image types
+	case strings.Contains(contentType, "image/jpeg"):
+		extension = ".jpg"
+	case strings.Contains(contentType, "image/png"):
+		extension = ".png"
+	case strings.Contains(contentType, "image/gif"):
+		extension = ".gif"
+	case strings.Contains(contentType, "image/webp"):
+		extension = ".webp"
+	case strings.Contains(contentType, "image/bmp"):
+		extension = ".bmp"
+	case strings.Contains(contentType, "image/tiff"):
+		extension = ".tiff"
+	case strings.Contains(contentType, "image/svg+xml"):
+		extension = ".svg"
+
+	// Video types
+	case strings.Contains(contentType, "video/mp4"):
+		extension = ".mp4"
+	case strings.Contains(contentType, "video/quicktime"):
+		extension = ".mov"
+	case strings.Contains(contentType, "video/x-msvideo"):
+		extension = ".avi"
+	case strings.Contains(contentType, "video/x-ms-wmv"):
+		extension = ".wmv"
+	case strings.Contains(contentType, "video/x-flv"):
+		extension = ".flv"
+	case strings.Contains(contentType, "video/webm"):
+		extension = ".webm"
+	case strings.Contains(contentType, "video/x-matroska"):
+		extension = ".mkv"
+
+	default:
+		// Check for general image or video types
+		if strings.Contains(contentType, "image/") {
+			extension = ".jpg" // Default for unknown image types
+		} else if strings.Contains(contentType, "video/") {
+			extension = ".mp4" // Default for unknown video types
+		} else {
+			// Try to determine from original file extension
+			origExt := strings.ToLower(filepath.Ext(filePath))
+			if origExt == ".mp4" || origExt == ".mov" || origExt == ".avi" ||
+				origExt == ".wmv" || origExt == ".flv" || origExt == ".webm" || origExt == ".mkv" {
+				extension = origExt
+			} else if origExt == ".jpg" || origExt == ".jpeg" || origExt == ".png" ||
+				origExt == ".gif" || origExt == ".webp" || origExt == ".bmp" ||
+				origExt == ".tiff" || origExt == ".svg" {
+				extension = origExt
+			} else {
+				return filePath, fmt.Errorf("not a recognized media type: %s", contentType)
+			}
+		}
+	}
+
+	// Create new file path with extension
+	dir := filepath.Dir(filePath)
+	base := filepath.Base(filePath)
+	newFilePath := filepath.Join(dir, base+extension)
+
+	// If the file already has the correct extension, just return it
+	if filePath == newFilePath {
+		return filePath, nil
+	}
+
+	// Rename the file
+	err = os.Rename(filePath, newFilePath)
+	if err != nil {
+		return "", fmt.Errorf("failed to rename file: %v", err)
+	}
+
+	log.Info().Str("oldPath", filePath).Str("newPath", newFilePath).Msg("Renamed image file with proper extension")
+	return newFilePath, nil
+}

uixt/mcp_server.go

@@ -133,6 +133,10 @@ func (s *MCPServer4XTDriver) registerTools() {
 	s.registerTool(&ToolGetScreenSize{})
 	s.registerTool(&ToolGetSource{})
 
+	// Media Album Tools
+	s.registerTool(&ToolPushAlbums{})
+	s.registerTool(&ToolClearAlbums{})
+
 	// Utility Tools
 	s.registerTool(&ToolSleep{})
 	s.registerTool(&ToolSleepMS{})

uixt/mcp_tools_device.go

@@ -3,6 +3,9 @@ package uixt
 import (
 	"context"
 	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
 
 	"github.com/danielpaulus/go-ios/ios"
 	"github.com/mark3labs/mcp-go/mcp"
@@ -216,3 +219,198 @@ func (t *ToolScreenRecord) Implement() server.ToolHandlerFunc {
 func (t *ToolScreenRecord) ConvertActionToCallToolRequest(action option.MobileAction) (mcp.CallToolRequest, error) {
 	return BuildMCPCallToolRequest(t.Name(), map[string]any{}, action), nil
 }
+
+// ToolPushAlbums implements the push_albums tool call.
+type ToolPushAlbums struct {
+	// Return data fields - these define the structure of data returned by this tool
+	FilePath string `json:"filePath" desc:"Path of the file that was pushed"`
+	FileUrl  string `json:"fileUrl,omitempty" desc:"URL of the file that was downloaded and pushed (if applicable)"`
+	FileType string `json:"fileType" desc:"Type of the file that was pushed (image or video)"`
+	Cleared  bool   `json:"cleared,omitempty" desc:"Whether albums were cleared before pushing (if applicable)"`
+}
+
+func (t *ToolPushAlbums) Name() option.ActionName {
+	return option.ACTION_PushAlbums
+}
+
+func (t *ToolPushAlbums) Description() string {
+	return "Push a media file (image or video) to the device's gallery. For Android, this will push the file to the DCIM/Camera directory. For iOS, this will add the file to the photo album."
+}
+
+func (t *ToolPushAlbums) Options() []mcp.ToolOption {
+	return []mcp.ToolOption{
+		mcp.WithString("platform", mcp.Enum("android", "ios"), mcp.Description("The platform type of device to push media to")),
+		mcp.WithString("serial", mcp.Description("The device serial number or UDID")),
+		mcp.WithString("filePath", mcp.Description("Path to the local media file to push to the device")),
+		mcp.WithString("fileUrl", mcp.Description("URL of the media file to download and push to the device")),
+		mcp.WithBoolean("cleanup", mcp.Description("Whether to delete the downloaded file after pushing it to the device")),
+		mcp.WithBoolean("clearBefore", mcp.Description("Whether to clear albums before pushing (if applicable)")),
+	}
+}
+
+func (t *ToolPushAlbums) ConvertActionToCallToolRequest(action option.MobileAction) (mcp.CallToolRequest, error) {
+	arguments := map[string]any{}
+
+	// Handle string param as fileUrl
+	if fileUrl, ok := action.Params.(string); ok && fileUrl != "" {
+		arguments["fileUrl"] = fileUrl
+	}
+
+	// Handle map params with fileUrl or filePath
+	if params, ok := action.Params.(map[string]interface{}); ok {
+		if fileUrl, ok := params["fileUrl"].(string); ok && fileUrl != "" {
+			arguments["fileUrl"] = fileUrl
+		}
+		if filePath, ok := params["filePath"].(string); ok && filePath != "" {
+			arguments["filePath"] = filePath
+		}
+		if cleanup, ok := params["cleanup"].(bool); ok {
+			arguments["cleanup"] = cleanup
+		}
+		if clearBefore, ok := params["clearBefore"].(bool); ok {
+			arguments["clearBefore"] = clearBefore
+		}
+	}
+
+	return BuildMCPCallToolRequest(t.Name(), arguments, action), nil
+}
+
+func (t *ToolPushAlbums) Implement() server.ToolHandlerFunc {
+	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		driverExt, err := setupXTDriver(ctx, request.GetArguments())
+		if err != nil {
+			return nil, err
+		}
+
+		// Get file path or URL
+		filePath, hasPath := request.GetArguments()["filePath"].(string)
+		fileUrl, hasUrl := request.GetArguments()["fileUrl"].(string)
+		cleanup, _ := request.GetArguments()["cleanup"].(bool)
+		clearBefore, _ := request.GetArguments()["clearBefore"].(bool)
+
+		// Check if we have either path or URL
+		if (!hasPath || filePath == "") && (!hasUrl || fileUrl == "") {
+			return nil, fmt.Errorf("either filePath or fileUrl is required")
+		}
+
+		// If we have a URL, download it
+		downloadedFile := false
+		fileType := "image" // Default file type
+		if hasUrl && fileUrl != "" {
+			log.Info().Str("fileUrl", fileUrl).Msg("Downloading media file from URL")
+			downloadedPath, err := DownloadFileByUrl(fileUrl)
+			if err != nil {
+				return nil, fmt.Errorf("failed to download media file from URL: %v", err)
+			}
+
+			// Detect file type and rename with proper extension
+			renamedPath, err := DetectAndRenameMediaFile(downloadedPath)
+			if err != nil {
+				log.Warn().Err(err).Str("path", downloadedPath).Msg("Failed to detect file type or rename file, using original file")
+				filePath = downloadedPath
+			} else {
+				filePath = renamedPath
+				// Determine if it's a video based on extension
+				ext := strings.ToLower(filepath.Ext(renamedPath))
+				if ext == ".mp4" || ext == ".mov" || ext == ".avi" || ext == ".wmv" || ext == ".flv" || ext == ".webm" || ext == ".mkv" {
+					fileType = "video"
+				}
+			}
+			downloadedFile = true
+		}
+
+		// Clear albums before pushing if requested
+		cleared := false
+		if clearBefore {
+			log.Info().Msg("Clearing albums before pushing new media file")
+			err := driverExt.IDriver.ClearImages()
+			if err != nil {
+				log.Warn().Err(err).Msg("Failed to clear albums before pushing, continuing anyway")
+			} else {
+				cleared = true
+			}
+		}
+
+		// Push the file to the device
+		err = driverExt.IDriver.PushImage(filePath)
+		if err != nil {
+			// If we downloaded the file and failed to push it, clean up
+			if downloadedFile && cleanup {
+				_ = os.Remove(filePath)
+			}
+			return nil, err
+		}
+
+		// Clean up downloaded file if requested
+		if downloadedFile && cleanup {
+			log.Info().Str("filePath", filePath).Msg("Cleaning up downloaded media file")
+			_ = os.Remove(filePath)
+		}
+
+		message := fmt.Sprintf("Successfully pushed %s to device", fileType)
+		returnData := ToolPushAlbums{
+			FilePath: filePath,
+			FileType: fileType,
+			Cleared:  cleared,
+		}
+
+		// Include URL in response if it was used
+		if hasUrl && fileUrl != "" {
+			returnData.FileUrl = fileUrl
+			message = fmt.Sprintf("Successfully downloaded and pushed %s from %s to device", fileType, fileUrl)
+		}
+
+		// Add cleared info to message if applicable
+		if cleared {
+			message = fmt.Sprintf("%s (albums cleared before pushing)", message)
+		}
+
+		return NewMCPSuccessResponse(message, &returnData), nil
+	}
+}
+
+// Old ToolPushImage implementation has been removed as part of the refactoring to ToolPushAlbums
+
+// ToolClearAlbums implements the clear_albums tool call.
+type ToolClearAlbums struct {
+	// Return data fields - these define the structure of data returned by this tool
+	Cleared bool `json:"cleared" desc:"Whether albums were cleared successfully"`
+}
+
+func (t *ToolClearAlbums) Name() option.ActionName {
+	return option.ACTION_ClearAlbums
+}
+
+func (t *ToolClearAlbums) Description() string {
+	return "Clear media files (images and videos) from the device's gallery. For Android, this will clear media from the DCIM/Camera directory. For iOS, this will clear media from the device's photo album."
+}
+
+func (t *ToolClearAlbums) Options() []mcp.ToolOption {
+	return []mcp.ToolOption{
+		mcp.WithString("platform", mcp.Enum("android", "ios"), mcp.Description("The platform type of device to clear media from")),
+		mcp.WithString("serial", mcp.Description("The device serial number or UDID")),
+	}
+}
+
+func (t *ToolClearAlbums) Implement() server.ToolHandlerFunc {
+	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		driverExt, err := setupXTDriver(ctx, request.GetArguments())
+		if err != nil {
+			return nil, err
+		}
+
+		err = driverExt.IDriver.ClearImages()
+		if err != nil {
+			return nil, err
+		}
+
+		message := "Successfully cleared media files from device"
+		returnData := ToolClearAlbums{Cleared: true}
+
+		return NewMCPSuccessResponse(message, &returnData), nil
+	}
+}
+
+func (t *ToolClearAlbums) ConvertActionToCallToolRequest(action option.MobileAction) (mcp.CallToolRequest, error) {
+	return BuildMCPCallToolRequest(t.Name(), map[string]any{}, action), nil
+}

uixt/option/action.go

@@ -98,6 +98,10 @@ const (
 	ACTION_ListAvailableDevices ActionName = "list_available_devices"
 	ACTION_SelectDevice         ActionName = "select_device"
 
+	// album actions (images and videos)
+	ACTION_PushAlbums  ActionName = "push_albums"
+	ACTION_ClearAlbums ActionName = "clear_albums"
+
 	// custom actions
 	ACTION_SwipeToTapApp   ActionName = "swipe_to_tap_app"   // swipe left & right to find app and tap
 	ACTION_SwipeToTapText  ActionName = "swipe_to_tap_text"  // swipe up & down to find text and tap