From ac86aca7c8e29dbd435f45677a11fdabdef889d0 Mon Sep 17 00:00:00 2001
From: wanyaoqi <d3lx.yq@gmail.com>
Date: Mon, 15 Apr 2024 21:45:04 +0800
Subject: [PATCH] feat(host): NewPciDevice add executor

---
 pkg/baremetal/tasks/baseprepare.go |  2 +-
 pkg/hostman/isolated_device/gpu.go | 56 +++++++++++++++++++++++-------
 pkg/util/ssh/ssh.go                |  4 +++
 3 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/pkg/baremetal/tasks/baseprepare.go b/pkg/baremetal/tasks/baseprepare.go
index 59f0afd3e5..9e4f0abf86 100644
--- a/pkg/baremetal/tasks/baseprepare.go
+++ b/pkg/baremetal/tasks/baseprepare.go
@@ -747,7 +747,7 @@ func getIsolatedDevicesInfo(cli *ssh.Client, ip net.IP) ([]*isolated_device.PCID
 	devs := []*isolated_device.PCIDevice{}
 	for _, line := range lines {
 		if len(line) > 0 {
-			dev := isolated_device.NewPCIDevice2(line)
+			dev := isolated_device.NewPCIDevice2(line, cli)
 			if len(dev.Addr) > 0 && utils.IsInArray(dev.ClassCode, isolated_device.GpuClassCodes) && !isBootVga(cli, dev, bootVgaPath) {
 				devs = append(devs, dev)
 			}
diff --git a/pkg/hostman/isolated_device/gpu.go b/pkg/hostman/isolated_device/gpu.go
index e41d21670d..2cb0475090 100644
--- a/pkg/hostman/isolated_device/gpu.go
+++ b/pkg/hostman/isolated_device/gpu.go
@@ -154,6 +154,22 @@ func getGPUPCIStr() ([]string, error) {
 	return GetPCIStrByAddr("")
 }
 
+type IExecutor interface {
+	RunCmd(cmd string) ([]string, error)
+}
+
+var defaultExecutor IExecutor = new(SDefaultExecutor)
+
+func GetDefaultExecutor() IExecutor {
+	return defaultExecutor
+}
+
+type SDefaultExecutor struct{}
+
+func (*SDefaultExecutor) RunCmd(cmd string) ([]string, error) {
+	return bashOutput(cmd)
+}
+
 type PCIDevice struct {
 	Addr          string `json:"bus_id"`
 	ClassName     string `json:"class_name"`
@@ -172,11 +188,23 @@ type PCIDevice struct {
 	PCIEInfo           *api.IsolatedDevicePCIEInfo `json:"pcie_info"`
 }
 
-func NewPCIDevice(line string) (*PCIDevice, error) {
-	if len(line) == 0 {
+func NewPCIDevice(addr string, executors ...IExecutor) (*PCIDevice, error) {
+	if len(addr) == 0 {
 		return nil, errors.Errorf("input line is empty")
 	}
-	dev := NewPCIDevice2(line)
+
+	var executor IExecutor
+	if len(executors) == 0 {
+		executor = GetDefaultExecutor()
+	} else {
+		executor = executors[0]
+	}
+	ret, err := executor.RunCmd(fmt.Sprintf("lspci -nnmm -s %s", addr))
+	if err != nil {
+		return nil, errors.Wrapf(err, "run lspci -nnmm -s %s", addr)
+	}
+
+	dev := NewPCIDevice2(strings.Join(ret, ""))
 	if err := dev.checkSameIOMMUGroupDevice(); err != nil {
 		return nil, err
 	}
@@ -186,9 +214,16 @@ func NewPCIDevice(line string) (*PCIDevice, error) {
 	return dev, nil
 }
 
-func NewPCIDevice2(line string) *PCIDevice {
+func NewPCIDevice2(line string, executors ...IExecutor) *PCIDevice {
+	var executor IExecutor
+	if len(executors) == 0 {
+		executor = GetDefaultExecutor()
+	} else {
+		executor = executors[0]
+	}
+
 	dev := parseLspci(line)
-	if err := dev.fillPCIEInfo(); err != nil {
+	if err := dev.fillPCIEInfo(executor); err != nil {
 		log.Warningf("fillPCIEInfo for line: %q, device: %s, error: %v", line, dev.String(), err)
 	}
 	return dev
@@ -416,12 +451,13 @@ func (d *PCIDevice) bindDriver() error {
 	)
 }
 
-func (d *PCIDevice) fillPCIEInfo() error {
+func (d *PCIDevice) fillPCIEInfo(executor IExecutor) error {
 	if d.Addr == "" {
 		return errors.Errorf("device address is empty: %s", d.String())
 	}
+
 	cmd := fmt.Sprintf("lspci -vvv -s %s", d.Addr)
-	lines, err := bashOutput(cmd)
+	lines, err := executor.RunCmd(cmd)
 	if err != nil {
 		return errors.Wrapf(err, "execute cmd: %s", cmd)
 	}
@@ -561,11 +597,7 @@ func (g *IOMMUGroup) String() string {
 }
 
 func detectPCIDevByAddr(addr string) (*PCIDevice, error) {
-	ret, err := bashOutput(fmt.Sprintf("lspci -nnmm -s %s", addr))
-	if err != nil {
-		return nil, err
-	}
-	return NewPCIDevice(strings.Join(ret, ""))
+	return NewPCIDevice(addr)
 }
 
 func detectPCIDevByAddrWithoutIOMMUGroup(addr string) (*PCIDevice, error) {
diff --git a/pkg/util/ssh/ssh.go b/pkg/util/ssh/ssh.go
index 20f3deff29..2bb653707e 100644
--- a/pkg/util/ssh/ssh.go
+++ b/pkg/util/ssh/ssh.go
@@ -147,6 +147,10 @@ func (s *Client) RawRun(cmds ...string) ([]string, error) {
 	return s.run(false, cmds, nil, false)
 }
 
+func (s *Client) RunCmd(cmd string) ([]string, error) {
+	return s.Run(cmd)
+}
+
 func (s *Client) Run(cmds ...string) ([]string, error) {
 	return s.run(true, cmds, nil, false)
 }