github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/testing/fake_listener.go (about)

     1  package testing
     2  
     3  import (
     4  	"errors"
     5  	"net"
     6  	"sync"
     7  )
     8  
     9  // FakeListener implements net.Listener and has an extra method (Provide) that allows
    10  //   you to queue a connection to be consumed via Accept.
    11  type FakeListener struct {
    12  	ch     chan net.Conn
    13  	closed bool
    14  	mutex  *sync.Mutex
    15  }
    16  
    17  func NewFakeListener() *FakeListener {
    18  	return &FakeListener{
    19  		ch:    make(chan net.Conn),
    20  		mutex: &sync.Mutex{},
    21  	}
    22  }
    23  
    24  func (fl *FakeListener) Provide(conn net.Conn) error {
    25  	fl.mutex.Lock()
    26  	defer fl.mutex.Unlock()
    27  
    28  	if fl.closed {
    29  		return errors.New("connection closed")
    30  	}
    31  
    32  	fl.ch <- conn
    33  	return nil
    34  }
    35  
    36  func (fl *FakeListener) Accept() (net.Conn, error) {
    37  	conn, more := <-fl.ch
    38  	if more {
    39  		return conn, nil
    40  	}
    41  
    42  	return nil, errors.New("connection closed")
    43  }
    44  
    45  func (fl *FakeListener) Close() error {
    46  	fl.mutex.Lock()
    47  	defer fl.mutex.Unlock()
    48  
    49  	if fl.closed {
    50  		return errors.New("connection closed")
    51  	}
    52  
    53  	fl.closed = true
    54  	close(fl.ch)
    55  	return nil
    56  }
    57  
    58  func (*FakeListener) Addr() net.Addr {
    59  	a, _ := net.ResolveTCPAddr("tcp", "fake-listener:1111")
    60  	return a
    61  }