b2: implement --b2-test-mode to set X-Bz-Test-Mode header #420

b2/b2.go

@@ -39,6 +39,7 @@ const (
 	maxSleep        = 2 * time.Second
 	decayConstant   = 2 // bigger for slower decay, exponential
 	maxParts        = 10000
+	testModeHeader  = "X-Bz-Test-Mode"
 )
 
 // Globals
@@ -49,6 +50,7 @@ var (
 	errorAuthTokenExpired       = errors.New("b2 auth token expired")
 	errorUploadTokenExpired     = errors.New("b2 upload token expired")
 	errorUploadPartTokenExpired = errors.New("b2 upload part token expired")
+	b2TestMode                  = pflag.StringP("b2-test-mode", "", "", "A flag string for X-Bz-Test-Mode header.")
 )
 
 // Register with Fs
@@ -224,6 +226,12 @@ func NewFs(name, root string) (fs.Fs, error) {
 		pacer:        pacer.New().SetMinSleep(minSleep).SetMaxSleep(maxSleep).SetDecayConstant(decayConstant),
 		uploadTokens: make(chan struct{}, fs.Config.Transfers),
 	}
+	// Set the test flag if required
+	if *b2TestMode != "" {
+		testMode := strings.TrimSpace(*b2TestMode)
+		f.srv.SetHeader(testModeHeader, testMode)
+		fs.Debug(f, "Setting test header \"%s: %s\"", testModeHeader, testMode)
+	}
 	// Fill up the upload tokens
 	for i := 0; i < fs.Config.Transfers; i++ {
 		f.returnUploadToken()