add ssh support

.gitignore

@@ -0,0 +1 @@
+.DS_Store

pkg/util/ssh/ssh.go

@@ -0,0 +1,24 @@
+package ssh
+
+import (
+	"context"
+	"fmt"
+	"os/exec"
+)
+
+// Exec executes the cmd on the remote node.
+// Here we assume we can run with `ssh node cmd` directly
+// TODO: add a SSH config?
+func Exec(ctx context.Context, node string, cmd string) error {
+	return exec.CommandContext(ctx, "ssh", node, cmd).Run()
+}
+
+// Upload uploads files from local path to remote node path.
+func Upload(ctx context.Context, localPath string, node string, remotePath string) error {
+	return exec.CommandContext(ctx, "scp", "-r", localPath, fmt.Sprintf("%s:%s", node, remotePath)).Run()
+}
+
+// Download downloads files from remote node path to local path.
+func Download(ctx context.Context, localPath string, node string, remotePath string) error {
+	return exec.CommandContext(ctx, "scp", "-r", fmt.Sprintf("%s:%s", node, remotePath), localPath).Run()
+}

pkg/util/ssh/ssh_test.go

@@ -0,0 +1,56 @@
+package ssh
+
+import (
+	"context"
+	"io/ioutil"
+	"os"
+	"path"
+	"testing"
+)
+
+func TestExec(t *testing.T) {
+	ctx := context.Background()
+	if err := Exec(ctx, "n1", "sleep 1"); err != nil {
+		t.Fatalf("ssh failed %v", err)
+	}
+
+	if err := Exec(ctx, "n1", "cat non_exist_file"); err == nil {
+		t.Fatal("ssh must fail")
+	}
+
+	if err := Exec(ctx, "n6", "sleep 1"); err == nil {
+		t.Fatal("ssh must fail")
+	}
+}
+
+func TestScp(t *testing.T) {
+	dir, err := ioutil.TempDir(".", "var")
+	if err != nil {
+		t.Fatal(err)
+	}
+	//defer os.RemoveAll(dir)
+
+	f, err := os.Create(path.Join(dir, "a.log"))
+	if err != nil {
+		t.Fatalf("create file failed %v", err)
+	}
+	f.WriteString("Hello world")
+	f.Close()
+
+	ctx := context.Background()
+	if err = Upload(ctx, path.Join(dir, "a.log"), "n1", "/tmp/b.log"); err != nil {
+		t.Fatalf("upload file failed %v", err)
+	}
+
+	if err = Download(ctx, path.Join(dir, "b.log"), "n1", "/tmp/b.log"); err != nil {
+		t.Fatalf("download file failed %v", err)
+	}
+
+	if err = Upload(ctx, dir, "n1", "/tmp/"); err != nil {
+		t.Fatalf("upload folder failed %v", err)
+	}
+
+	if err = Download(ctx, path.Join(dir, "sub"), "n1", path.Join("/tmp", path.Base(dir))); err != nil {
+		t.Fatalf("download foler failed %v", err)
+	}
+}

pkg/util/util_test.go

@@ -8,6 +8,7 @@ import (
 	"path"
 	"strconv"
 	"testing"
+	"time"
 )
 
 func TestWget(t *testing.T) {
@@ -93,6 +94,8 @@ func TestDaemon(t *testing.T) {
 		t.Fatalf("stop daemon failed %v", err)
 	}
 
+	time.Sleep(time.Second)
+
 	if IsProcessExist(pid) {
 		t.Fatalf("pid %d must not exist", pid)
 	}