1 /* 2 * 3 * Copyright 2022 gRPC authors. 4 * 5 * Licensed under the Apache License, Version 2.0 (the "License"); 6 * you may not use this file except in compliance with the License. 7 * You may obtain a copy of the License at 8 * 9 * http://www.apache.org/licenses/LICENSE-2.0 10 * 11 * Unless required by applicable law or agreed to in writing, software 12 * distributed under the License is distributed on an "AS IS" BASIS, 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 * See the License for the specific language governing permissions and 15 * limitations under the License. 16 * 17 */ 18 19 package testutils 20 21 import ( 22 "net" 23 "testing" 24 ) 25 26 // ConnWrapper wraps a net.Conn and pushes on a channel when closed. 27 type ConnWrapper struct { 28 net.Conn 29 CloseCh *Channel 30 } 31 32 // Close closes the connection and sends a value on the close channel. 33 func (cw *ConnWrapper) Close() error { 34 err := cw.Conn.Close() 35 cw.CloseCh.Replace(nil) 36 return err 37 } 38 39 // ListenerWrapper wraps a net.Listener and the returned net.Conn. 40 // 41 // It pushes on a channel whenever it accepts a new connection. 42 type ListenerWrapper struct { 43 net.Listener 44 NewConnCh *Channel 45 } 46 47 // Accept wraps the Listener Accept and sends the accepted connection on a 48 // channel. 49 func (l *ListenerWrapper) Accept() (net.Conn, error) { 50 c, err := l.Listener.Accept() 51 if err != nil { 52 return nil, err 53 } 54 closeCh := NewChannel() 55 conn := &ConnWrapper{Conn: c, CloseCh: closeCh} 56 l.NewConnCh.Send(conn) 57 return conn, nil 58 } 59 60 // NewListenerWrapper returns a ListenerWrapper. 61 func NewListenerWrapper(t *testing.T, lis net.Listener) *ListenerWrapper { 62 if lis == nil { 63 var err error 64 lis, err = LocalTCPListener() 65 if err != nil { 66 t.Fatal(err) 67 } 68 } 69 70 return &ListenerWrapper{ 71 Listener: lis, 72 NewConnCh: NewChannel(), 73 } 74 } 75