util: added SelectWithTimeout && WaitWithTImeout

This commit is contained in:
CPunch 2023-12-01 20:13:42 -06:00
parent 557117f093
commit 3445b852fd
2 changed files with 28 additions and 25 deletions

View File

@ -8,10 +8,10 @@ import (
"os" "os"
"sync" "sync"
"testing" "testing"
"time"
"github.com/CPunch/gopenfusion/cnet" "github.com/CPunch/gopenfusion/cnet"
"github.com/CPunch/gopenfusion/cnet/protocol" "github.com/CPunch/gopenfusion/cnet/protocol"
"github.com/CPunch/gopenfusion/util"
"github.com/matryer/is" "github.com/matryer/is"
) )
@ -24,25 +24,6 @@ const (
maxDummyPeers = 5 maxDummyPeers = 5
) )
func selectWithTimeout(ch <-chan struct{}, seconds int) bool {
select {
case <-ch:
return true
case <-time.After(time.Duration(seconds) * time.Second):
return false
}
}
func waitWithTimeout(wg *sync.WaitGroup, seconds int) bool {
done := make(chan struct{})
go func() {
defer close(done)
wg.Wait()
}()
return selectWithTimeout(done, seconds)
}
func TestMain(m *testing.M) { func TestMain(m *testing.M) {
var err error var err error
srvcPort, err = cnet.RandomPort() srvcPort, err = cnet.RandomPort()
@ -62,7 +43,7 @@ func TestService(t *testing.T) {
// shutdown service when test is done // shutdown service when test is done
defer func() { defer func() {
cancel() cancel()
is.True(selectWithTimeout(srvc.Stopped(), timeout)) // wait for service to stop with timeout is.True(util.SelectWithTimeout(srvc.Stopped(), timeout)) // wait for service to stop with timeout
}() }()
// our dummy packet handler // our dummy packet handler
@ -86,7 +67,7 @@ func TestService(t *testing.T) {
// run service // run service
go func() { is.NoErr(srvc.Start()) }() // srvc.Start error go func() { is.NoErr(srvc.Start()) }() // srvc.Start error
is.True(selectWithTimeout(srvc.Started(), timeout)) // wait for service to start with timeout is.True(util.SelectWithTimeout(srvc.Started(), timeout)) // wait for service to start with timeout
wg.Add(maxDummyPeers * 2) // 2 wg.Done() per peer for receiving packets wg.Add(maxDummyPeers * 2) // 2 wg.Done() per peer for receiving packets
for i := 0; i < maxDummyPeers; i++ { for i := 0; i < maxDummyPeers; i++ {
@ -111,5 +92,5 @@ func TestService(t *testing.T) {
}() }()
} }
is.True(waitWithTimeout(&wg, timeout)) // wait for all dummy peers to be done with timeout is.True(util.WaitWithTimeout(&wg, timeout)) // wait for all dummy peers to be done with timeout
} }

View File

@ -1,7 +1,29 @@
package util package util
import "time" import (
"sync"
"time"
)
func GetTime() uint64 { func GetTime() uint64 {
return uint64(time.Now().UnixMilli()) return uint64(time.Now().UnixMilli())
} }
func SelectWithTimeout(ch <-chan struct{}, seconds int) bool {
select {
case <-ch:
return true
case <-time.After(time.Duration(seconds) * time.Second):
return false
}
}
func WaitWithTimeout(wg *sync.WaitGroup, seconds int) bool {
done := make(chan struct{})
go func() {
defer close(done)
wg.Wait()
}()
return SelectWithTimeout(done, seconds)
}