github.com/simpleiot/simpleiot@v0.18.3/test/io-sim.go (about)

     1  package test
     2  
     3  import (
     4  	"bytes"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  // IoSim is used to simulate an io channel -- provides both sides so you can easily
    10  // test code that uses an io.ReadWriter interface, etc
    11  type IoSim struct {
    12  	out  *bytes.Buffer
    13  	in   *bytes.Buffer
    14  	m    *sync.Mutex
    15  	stop chan struct{}
    16  }
    17  
    18  // NewIoSim creates a new IO sim and returns the A and B side of an IO simulator
    19  // that implements a io.ReadWriteCloser. Typically the A side would be used by your
    20  // test code to read/write simulated data, and the B side would be passed to your code
    21  // that uses an io.ReadWriter.
    22  func NewIoSim() (*IoSim, *IoSim) {
    23  	var a2b bytes.Buffer
    24  	var b2a bytes.Buffer
    25  	var m sync.Mutex
    26  
    27  	a := IoSim{&a2b, &b2a, &m, make(chan struct{})}
    28  	b := IoSim{&b2a, &a2b, &m, make(chan struct{})}
    29  
    30  	return &a, &b
    31  }
    32  
    33  func (ios *IoSim) Write(d []byte) (int, error) {
    34  	ios.m.Lock()
    35  	defer ios.m.Unlock()
    36  	return ios.in.Write(d)
    37  }
    38  
    39  // Read blocks until there is some data in the out buffer or the ioSim is closed.
    40  func (ios *IoSim) Read(d []byte) (int, error) {
    41  	ret := make(chan struct{})
    42  
    43  	go func() {
    44  		for {
    45  			ios.m.Lock()
    46  			if ios.out.Len() > 0 {
    47  				close(ret)
    48  				ios.m.Unlock()
    49  				return
    50  			}
    51  			ios.m.Unlock()
    52  			select {
    53  			case <-time.After(time.Millisecond):
    54  				// continue
    55  			case <-ios.stop:
    56  				close(ret)
    57  				return
    58  			}
    59  		}
    60  	}()
    61  
    62  	// block until we have data
    63  	<-ret
    64  	ios.m.Lock()
    65  	defer ios.m.Unlock()
    66  	return ios.out.Read(d)
    67  }
    68  
    69  // Close simulator
    70  func (ios *IoSim) Close() error {
    71  	close(ios.stop)
    72  	return nil
    73  }