You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
198 lines
3.7 KiB
198 lines
3.7 KiB
package graceful
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
var b = make([]byte, 0)
|
|
|
|
func connify(c net.Conn) *conn {
|
|
switch c.(type) {
|
|
case (*conn):
|
|
return c.(*conn)
|
|
case (*sendfile):
|
|
return &c.(*sendfile).conn
|
|
default:
|
|
panic("IDK")
|
|
}
|
|
}
|
|
|
|
func assertState(t *testing.T, n net.Conn, st connstate) {
|
|
c := connify(n)
|
|
c.m.Lock()
|
|
defer c.m.Unlock()
|
|
if c.state != st {
|
|
t.Fatalf("conn was %v, but expected %v", c.state, st)
|
|
}
|
|
}
|
|
|
|
// Not super happy about making the tests dependent on the passing of time, but
|
|
// I'm not really sure what else to do.
|
|
|
|
func expectCall(t *testing.T, ch <-chan struct{}, name string) {
|
|
select {
|
|
case <-ch:
|
|
case <-time.After(5 * time.Millisecond):
|
|
t.Fatalf("Expected call to %s", name)
|
|
}
|
|
}
|
|
|
|
func TestCounting(t *testing.T) {
|
|
kill = make(chan struct{})
|
|
c := WrapConn(fakeConn{})
|
|
ch := make(chan struct{})
|
|
|
|
go func() {
|
|
wg.Wait()
|
|
ch <- struct{}{}
|
|
}()
|
|
|
|
select {
|
|
case <-ch:
|
|
t.Fatal("Expected connection to keep us from quitting")
|
|
case <-time.After(5 * time.Millisecond):
|
|
}
|
|
|
|
c.Close()
|
|
expectCall(t, ch, "wg.Wait()")
|
|
}
|
|
|
|
func TestStateTransitions1(t *testing.T) {
|
|
kill = make(chan struct{})
|
|
ch := make(chan struct{})
|
|
|
|
onclose := make(chan struct{})
|
|
read := make(chan struct{})
|
|
deadline := make(chan struct{})
|
|
c := WrapConn(fakeConn{
|
|
onClose: func() {
|
|
onclose <- struct{}{}
|
|
},
|
|
onRead: func() {
|
|
read <- struct{}{}
|
|
},
|
|
onSetReadDeadline: func() {
|
|
deadline <- struct{}{}
|
|
},
|
|
})
|
|
|
|
go func() {
|
|
wg.Wait()
|
|
ch <- struct{}{}
|
|
}()
|
|
|
|
assertState(t, c, csWaiting)
|
|
|
|
// Waiting + Read() = Working
|
|
go c.Read(b)
|
|
expectCall(t, read, "c.Read()")
|
|
assertState(t, c, csWorking)
|
|
|
|
// Working + SetReadDeadline() = Waiting
|
|
go c.SetReadDeadline(time.Now())
|
|
expectCall(t, deadline, "c.SetReadDeadline()")
|
|
assertState(t, c, csWaiting)
|
|
|
|
// Waiting + kill = Dead
|
|
close(kill)
|
|
expectCall(t, onclose, "c.Close()")
|
|
assertState(t, c, csDead)
|
|
|
|
expectCall(t, ch, "wg.Wait()")
|
|
}
|
|
|
|
func TestStateTransitions2(t *testing.T) {
|
|
kill = make(chan struct{})
|
|
ch := make(chan struct{})
|
|
onclose := make(chan struct{})
|
|
read := make(chan struct{})
|
|
deadline := make(chan struct{})
|
|
c := WrapConn(fakeConn{
|
|
onClose: func() {
|
|
onclose <- struct{}{}
|
|
},
|
|
onRead: func() {
|
|
read <- struct{}{}
|
|
},
|
|
onSetReadDeadline: func() {
|
|
deadline <- struct{}{}
|
|
},
|
|
})
|
|
|
|
go func() {
|
|
wg.Wait()
|
|
ch <- struct{}{}
|
|
}()
|
|
|
|
assertState(t, c, csWaiting)
|
|
|
|
// Waiting + Read() = Working
|
|
go c.Read(b)
|
|
expectCall(t, read, "c.Read()")
|
|
assertState(t, c, csWorking)
|
|
|
|
// Working + Read() = Working
|
|
go c.Read(b)
|
|
expectCall(t, read, "c.Read()")
|
|
assertState(t, c, csWorking)
|
|
|
|
// Working + kill = Dying
|
|
close(kill)
|
|
time.Sleep(5 * time.Millisecond)
|
|
assertState(t, c, csDying)
|
|
|
|
// Dying + Read() = Dying
|
|
go c.Read(b)
|
|
expectCall(t, read, "c.Read()")
|
|
assertState(t, c, csDying)
|
|
|
|
// Dying + SetReadDeadline() = Dead
|
|
go c.SetReadDeadline(time.Now())
|
|
expectCall(t, deadline, "c.SetReadDeadline()")
|
|
assertState(t, c, csDead)
|
|
|
|
expectCall(t, ch, "wg.Wait()")
|
|
}
|
|
|
|
func TestHijack(t *testing.T) {
|
|
kill = make(chan struct{})
|
|
fake := fakeConn{}
|
|
c := WrapConn(fake)
|
|
ch := make(chan struct{})
|
|
|
|
go func() {
|
|
wg.Wait()
|
|
ch <- struct{}{}
|
|
}()
|
|
|
|
cc := connify(c)
|
|
if _, ok := cc.hijack().(fakeConn); !ok {
|
|
t.Error("Expected original connection back out")
|
|
}
|
|
assertState(t, c, csDead)
|
|
expectCall(t, ch, "wg.Wait()")
|
|
}
|
|
|
|
type fakeSendfile struct {
|
|
fakeConn
|
|
}
|
|
|
|
func (f fakeSendfile) ReadFrom(r io.Reader) (int64, error) {
|
|
return 0, nil
|
|
}
|
|
|
|
func TestReadFrom(t *testing.T) {
|
|
kill = make(chan struct{})
|
|
c := WrapConn(fakeSendfile{})
|
|
r := strings.NewReader("Hello world")
|
|
|
|
if rf, ok := c.(io.ReaderFrom); ok {
|
|
rf.ReadFrom(r)
|
|
} else {
|
|
t.Fatal("Expected a ReaderFrom in return")
|
|
}
|
|
}
|