Golang safe.Buffer类代码示例

func execWithTimeout(timeout time.Duration, cmd string, args ...string) (output []byte, err error) {
	var buf safe.Buffer
	ch := make(chan []byte, 1)
	errCh := make(chan error, 1)
	command := osexec.Command(cmd, args...)
	command.Stdout = &Writer{&buf}
	command.Stderr = &Writer{&buf}
	if err = command.Start(); err != nil {
		return nil, err
	}
	go func() {
		if err := command.Wait(); err == nil {
			ch <- buf.Bytes()
		} else {
			errCh <- err
			ch <- buf.Bytes()
		}
	}()
	select {
	case output = <-ch:
		select {
		case err = <-errCh:
		case <-time.After(1e9):
		}
	case err = <-errCh:
		output = <-ch
	case <-time.After(timeout):
		argsStr := strings.Join(args, " ")
		err = fmt.Errorf("%q ran for more than %s.", cmd+" "+argsStr, timeout)
		command.Process.Kill()
	}
	return output, err
}

func TestImportImage(t *testing.T) {
	server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("importing from 1"))
	}))
	defer server1.Close()
	server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("importing from 2"))
	}))
	defer server2.Close()
	cluster, err := New(nil, &MapStorage{},
		Node{Address: server1.URL},
		Node{Address: server2.URL},
	)
	if err != nil {
		t.Fatal(err)
	}
	var buf safe.Buffer
	opts := docker.ImportImageOptions{
		Repository:   "tsuru/python",
		Source:       "http://url.to/tar",
		OutputStream: &buf,
	}
	err = cluster.ImportImage(opts)
	if err != nil {
		t.Error(err)
	}
	re := regexp.MustCompile(`^importing from \d`)
	if !re.MatchString(buf.String()) {
		t.Errorf("Wrong output: Want %q. Got %q.", "importing from [12]", buf.String())
	}
}

func TestPullImageSpecifyNode(t *testing.T) {
	server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/images/tsuru/python/json" {
			w.Write([]byte(`{"Id": "id1"}`))
		} else {
			w.Write([]byte("Pulling from 1!"))
		}
	}))
	defer server1.Close()
	server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/images/tsuru/python/json" {
			w.Write([]byte(`{"Id": "id1"}`))
		} else {
			w.Write([]byte("Pulling from 2!"))
		}
	}))
	defer server2.Close()
	var buf safe.Buffer
	cluster, err := New(nil, &MapStorage{},
		Node{Address: server1.URL},
		Node{Address: server2.URL},
	)
	if err != nil {
		t.Fatal(err)
	}
	opts := docker.PullImageOptions{Repository: "tsuru/python", OutputStream: &buf}
	err = cluster.PullImage(opts, docker.AuthConfiguration{}, server2.URL)
	if err != nil {
		t.Error(err)
	}
	expected := "Pulling from 2!"
	if r := buf.String(); r != expected {
		t.Errorf("Wrong output: Want %q. Got %q.", expected, r)
	}
}

func TestPushImage(t *testing.T) {
	server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Pushing to server 1!"))
	}))
	defer server1.Close()
	server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Pushing to server 2!"))
	}))
	defer server2.Close()
	var buf safe.Buffer
	stor := &MapStorage{}
	err := stor.StoreImage("tsuru/ruby", "id1", server1.URL)
	if err != nil {
		t.Fatal(err)
	}
	cluster, err := New(nil, stor,
		Node{Address: server1.URL},
		Node{Address: server2.URL},
	)
	if err != nil {
		t.Fatal(err)
	}
	var auth docker.AuthConfiguration
	err = cluster.PushImage(docker.PushImageOptions{Name: "tsuru/ruby", OutputStream: &buf}, auth)
	if err != nil {
		t.Fatal(err)
	}
	re := regexp.MustCompile(`^Pushing to server \d`)
	if !re.MatchString(buf.String()) {
		t.Errorf("Wrong output: Want %q. Got %q.", "Pushing to server [12]", buf.String())
	}
}

// PullImage pulls an image from a remote registry server, returning an error
// in case of failure.
//
// It will pull all images in parallel, so users need to make sure that the
// given buffer is safe.
func (c *Cluster) PullImage(opts docker.PullImageOptions, auth docker.AuthConfiguration, nodes ...string) error {
	var w safe.Buffer
	if opts.OutputStream != nil {
		mw := io.MultiWriter(&w, opts.OutputStream)
		opts.OutputStream = mw
	} else {
		opts.OutputStream = &w
	}
	key := imageKey(opts.Repository, opts.Tag)
	_, err := c.runOnNodes(func(n node) (interface{}, error) {
		n.setPersistentClient()
		err := n.PullImage(opts, auth)
		if err != nil {
			return nil, err
		}
		img, err := n.InspectImage(key)
		if err != nil {
			return nil, err
		}
		return nil, c.storage().StoreImage(key, img.ID, n.addr)
	}, docker.ErrNoSuchImage, true, nodes...)
	if err != nil {
		return err
	}
	digest, _ := fix.GetImageDigest(w.String())
	return c.storage().SetImageDigest(key, digest)
}

func TestPullImageSpecifyMultipleNodes(t *testing.T) {
	server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Pulling from 1!"))
	}))
	defer server1.Close()
	server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Pulling from 2!"))
	}))
	defer server2.Close()
	server3 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Pulling from 3!"))
	}))
	defer server3.Close()
	var buf safe.Buffer
	cluster, err := New(nil, &MapStorage{},
		Node{Address: server1.URL},
		Node{Address: server2.URL},
		Node{Address: server3.URL},
	)
	if err != nil {
		t.Fatal(err)
	}
	opts := docker.PullImageOptions{Repository: "tsuru/python", OutputStream: &buf}
	err = cluster.PullImage(opts, docker.AuthConfiguration{}, server2.URL, server3.URL)
	if err != nil {
		t.Error(err)
	}
	alternatives := []string{
		"Pulling from 2!Pulling from 3!",
		"Pulling from 3!Pulling from 2!",
	}
	if r := buf.String(); r != alternatives[0] && r != alternatives[1] {
		t.Errorf("Wrong output: Want %q. Got %q.", "Pulling from 2!Pulling from 3!", r)
	}
}

// pushImage sends the given image to the registry server defined in the
// configuration file.
func pushImage(name string) error {
	if _, err := config.GetString("docker:registry"); err == nil {
		var buf safe.Buffer
		pushOpts := docker.PushImageOptions{Name: name, OutputStream: &buf}
		err = dockerCluster().PushImage(pushOpts, docker.AuthConfiguration{})
		if err != nil {
			log.Errorf("[docker] Failed to push image %q (%s): %s", name, err, buf.String())
			return err
		}
	}
	return nil
}

// PushImage sends the given image to the registry server defined in the
// configuration file.
func (p *dockerProvisioner) PushImage(name, tag string) error {
	if _, err := config.GetString("docker:registry"); err == nil {
		var buf safe.Buffer
		pushOpts := docker.PushImageOptions{Name: name, Tag: tag, OutputStream: &buf}
		err = p.Cluster().PushImage(pushOpts, p.RegistryAuthConfig())
		if err != nil {
			log.Errorf("[docker] Failed to push image %q (%s): %s", name, err, buf.String())
			return err
		}
	}
	return nil
}

func TestPullImage(t *testing.T) {
	server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/images/tsuru/python/json" {
			w.Write([]byte(`{"Id": "id1"}`))
		} else {
			w.Write([]byte("Pulling from 1!"))
		}
	}))
	defer server1.Close()
	server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/images/tsuru/python/json" {
			w.Write([]byte(`{"Id": "id1"}`))
		} else {
			w.Write([]byte("Pulling from 2!"))
		}
	}))
	defer server2.Close()
	var buf safe.Buffer
	cluster, err := New(nil, &MapStorage{},
		Node{Address: server1.URL},
		Node{Address: server2.URL},
	)
	if err != nil {
		t.Fatal(err)
	}
	opts := docker.PullImageOptions{Repository: "tsuru/python", OutputStream: &buf}
	err = cluster.PullImage(opts, docker.AuthConfiguration{})
	if err != nil {
		t.Error(err)
	}
	alternatives := []string{
		"Pulling from 1!Pulling from 2!",
		"Pulling from 2!Pulling from 1!",
	}
	if r := buf.String(); r != alternatives[0] && r != alternatives[1] {
		t.Errorf("Wrong output: Want %q. Got %q.", "Pulling from 1!Pulling from 2!", buf.String())
	}
	img, err := cluster.storage().RetrieveImage("tsuru/python")
	if err != nil {
		t.Fatal(err)
	}
	expected := []ImageHistory{
		{Node: server1.URL, ImageId: "id1"},
		{Node: server2.URL, ImageId: "id1"},
	}
	if !reflect.DeepEqual(img.History, expected) {
		expected[0], expected[1] = expected[1], expected[0]
		if !reflect.DeepEqual(img.History, expected) {
			t.Errorf("Wrong output: Want %q. Got %q.", expected, img)
		}
	}
}

func TestPullImage(t *testing.T) {
	server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Pulling from 1!"))
	}))
	defer server1.Close()
	server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Pulling from 2!"))
	}))
	defer server2.Close()
	var buf safe.Buffer
	cluster, err := New(nil, &MapStorage{},
		Node{Address: server1.URL},
		Node{Address: server2.URL},
	)
	if err != nil {
		t.Fatal(err)
	}
	opts := docker.PullImageOptions{Repository: "tsuru/python", OutputStream: &buf}
	err = cluster.PullImage(opts, docker.AuthConfiguration{})
	if err != nil {
		t.Error(err)
	}
	alternatives := []string{
		"Pulling from 1!Pulling from 2!",
		"Pulling from 2!Pulling from 1!",
	}
	if r := buf.String(); r != alternatives[0] && r != alternatives[1] {
		t.Errorf("Wrong output: Want %q. Got %q.", "Pulling from 1!Pulling from 2!", buf.String())
	}
	nodes, err := cluster.storage().RetrieveImage("tsuru/python")
	if err != nil {
		t.Fatal(err)
	}
	expected := []string{server1.URL, server2.URL}
	sort.Strings(nodes)
	sort.Strings(expected)
	if !reflect.DeepEqual(nodes, expected) {
		t.Errorf("Wrong output: Want %q. Got %q.", expected, nodes)
	}
}

func (s *S) TestRecreateBsContainers(c *check.C) {
	_, err := nodecontainer.InitializeBS()
	c.Assert(err, check.IsNil)
	p, err := dockertest.StartMultipleServersCluster()
	c.Assert(err, check.IsNil)
	defer p.Destroy()
	var buf safe.Buffer
	err = recreateContainers(p, &buf)
	c.Assert(err, check.IsNil)
	nodes, err := p.Cluster().Nodes()
	c.Assert(err, check.IsNil)
	c.Assert(nodes, check.HasLen, 2)
	client, err := nodes[0].Client()
	c.Assert(err, check.IsNil)
	containers, err := client.ListContainers(docker.ListContainersOptions{All: true})
	c.Assert(err, check.IsNil)
	c.Assert(containers, check.HasLen, 1)
	container, err := client.InspectContainer(containers[0].ID)
	c.Assert(err, check.IsNil)
	c.Assert(container.Name, check.Equals, nodecontainer.BsDefaultName)
	client, err = nodes[1].Client()
	c.Assert(err, check.IsNil)
	containers, err = client.ListContainers(docker.ListContainersOptions{All: true})
	c.Assert(err, check.IsNil)
	c.Assert(containers, check.HasLen, 1)
	container, err = client.InspectContainer(containers[0].ID)
	c.Assert(err, check.IsNil)
	c.Assert(container.Name, check.Equals, nodecontainer.BsDefaultName)
	// It runs in parallel, so we check both ordering
	output1 := fmt.Sprintf(`relaunching node container "big-sibling" in the node %s []
relaunching node container "big-sibling" in the node %s []
`, nodes[0].Address, nodes[1].Address)
	output2 := fmt.Sprintf(`relaunching node container "big-sibling" in the node %s []
relaunching node container "big-sibling" in the node %s []
`, nodes[1].Address, nodes[0].Address)
	if got := buf.String(); got != output1 && got != output2 {
		c.Errorf("Wrong output:\n%s", got)
	}
}