github.com/streamdal/segmentio-kafka-go@v0.4.47-streamdal/testing/conn.go (about)

     1  package testing
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"sync"
     7  )
     8  
     9  type ConnWaitGroup struct {
    10  	DialFunc func(context.Context, string, string) (net.Conn, error)
    11  	sync.WaitGroup
    12  }
    13  
    14  func (g *ConnWaitGroup) Dial(ctx context.Context, network, address string) (net.Conn, error) {
    15  	c, err := g.DialFunc(ctx, network, address)
    16  	if err != nil {
    17  		return nil, err
    18  	}
    19  	g.Add(1)
    20  	return &groupConn{Conn: c, group: g}, nil
    21  }
    22  
    23  type groupConn struct {
    24  	net.Conn
    25  	group *ConnWaitGroup
    26  	once  sync.Once
    27  }
    28  
    29  func (c *groupConn) Close() error {
    30  	defer c.once.Do(c.group.Done)
    31  	return c.Conn.Close()
    32  }