github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/utils/progress_test.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package utils
     4  
     5  import (
     6  	"context"
     7  	"time"
     8  
     9  	. "github.com/pingcap/check"
    10  )
    11  
    12  type testProgressSuite struct{}
    13  
    14  var _ = Suite(&testProgressSuite{})
    15  
    16  type testWriter struct {
    17  	fn func(string)
    18  }
    19  
    20  func (t *testWriter) Write(p []byte) (int, error) {
    21  	t.fn(string(p))
    22  	return len(p), nil
    23  }
    24  
    25  func (r *testProgressSuite) TestProgress(c *C) {
    26  	ctx, cancel := context.WithCancel(context.Background())
    27  
    28  	var p string
    29  	pCh2 := make(chan string, 2)
    30  	progress2 := NewProgressPrinter("test", 2, false)
    31  	progress2.goPrintProgress(ctx, nil, &testWriter{
    32  		fn: func(p string) { pCh2 <- p },
    33  	})
    34  	progress2.Inc()
    35  	time.Sleep(2 * time.Second)
    36  	p = <-pCh2
    37  	c.Assert(p, Matches, `.*"P":"50\.00%".*`)
    38  	progress2.Inc()
    39  	time.Sleep(2 * time.Second)
    40  	p = <-pCh2
    41  	c.Assert(p, Matches, `.*"P":"100\.00%".*`)
    42  	progress2.Inc()
    43  	time.Sleep(2 * time.Second)
    44  	p = <-pCh2
    45  	c.Assert(p, Matches, `.*"P":"100\.00%".*`)
    46  
    47  	pCh4 := make(chan string, 4)
    48  	progress4 := NewProgressPrinter("test", 4, false)
    49  	progress4.goPrintProgress(ctx, nil, &testWriter{
    50  		fn: func(p string) { pCh4 <- p },
    51  	})
    52  	progress4.Inc()
    53  	time.Sleep(2 * time.Second)
    54  	p = <-pCh4
    55  	c.Assert(p, Matches, `.*"P":"25\.00%".*`)
    56  	progress4.Inc()
    57  	progress4.Close()
    58  	time.Sleep(2 * time.Second)
    59  	p = <-pCh4
    60  	c.Assert(p, Matches, `.*"P":"100\.00%".*`)
    61  
    62  	pCh8 := make(chan string, 8)
    63  	progress8 := NewProgressPrinter("test", 8, false)
    64  	progress8.goPrintProgress(ctx, nil, &testWriter{
    65  		fn: func(p string) { pCh8 <- p },
    66  	})
    67  	progress8.Inc()
    68  	progress8.Inc()
    69  	time.Sleep(2 * time.Second)
    70  	p = <-pCh8
    71  	c.Assert(p, Matches, `.*"P":"25\.00%".*`)
    72  
    73  	// Cancel should stop progress at the current position.
    74  	cancel()
    75  	p = <-pCh8
    76  	c.Assert(p, Matches, `.*"P":"25\.00%".*`)
    77  }