github.com/rkt/rkt@v1.30.1-0.20200224141603-171c416fac02/tests/testutils/goroutineassistant.go (about)

     1  // Copyright 2015 The rkt Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package testutils
    16  
    17  import (
    18  	"fmt"
    19  	"runtime"
    20  	"sync"
    21  	"testing"
    22  
    23  	"github.com/coreos/gexpect"
    24  )
    25  
    26  type GoroutineAssistant struct {
    27  	s  chan error
    28  	wg sync.WaitGroup
    29  	t  *testing.T
    30  }
    31  
    32  func NewGoroutineAssistant(t *testing.T) *GoroutineAssistant {
    33  	return &GoroutineAssistant{
    34  		s: make(chan error),
    35  		t: t,
    36  	}
    37  }
    38  
    39  func (a *GoroutineAssistant) Fatalf(s string, args ...interface{}) {
    40  	a.s <- fmt.Errorf(s, args...)
    41  	runtime.Goexit()
    42  }
    43  
    44  func (a *GoroutineAssistant) Add(n int) {
    45  	a.wg.Add(n)
    46  }
    47  
    48  func (a *GoroutineAssistant) Done() {
    49  	a.wg.Done()
    50  }
    51  
    52  func (a *GoroutineAssistant) Wait() {
    53  	go func() {
    54  		a.wg.Wait()
    55  		a.s <- nil
    56  	}()
    57  	err := <-a.s
    58  	if err == nil {
    59  		// success
    60  		return
    61  	}
    62  	// Let's mimic testing.Fatalf()'s behavior and fatal with the first received error.
    63  	// All other goroutines will be killed.
    64  	a.t.Logf("Error encountered - shutting down")
    65  	a.t.Fatal(err)
    66  }
    67  
    68  func (a *GoroutineAssistant) WaitOrFail(child *gexpect.ExpectSubprocess) {
    69  	err := child.Wait()
    70  	if err != nil {
    71  		a.Fatalf("rkt didn't terminate correctly: %v", err)
    72  	}
    73  }
    74  
    75  func (a *GoroutineAssistant) SpawnOrFail(cmd string) *gexpect.ExpectSubprocess {
    76  	a.t.Logf("Command: %v", cmd)
    77  	child, err := gexpect.Spawn(cmd)
    78  	if err != nil {
    79  		a.Fatalf("Cannot exec rkt: %v", err)
    80  	}
    81  	return child
    82  }