From 899b95b4e6ade2299f88d3bbada27c306427c9ad Mon Sep 17 00:00:00 2001
From: CPunch <sethtstubbs@gmail.com>
Date: Fri, 1 Dec 2023 20:23:27 -0600
Subject: [PATCH] SelectWithTimeout && WaitWithTImeout now use time.Duration

---
 cnet/service_test.go |  3 ++-
 util/util.go         |  8 ++++----
 util/util_test.go    | 12 ++++++------
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/cnet/service_test.go b/cnet/service_test.go
index 705e555..8c73dcc 100644
--- a/cnet/service_test.go
+++ b/cnet/service_test.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/CPunch/gopenfusion/cnet"
 	"github.com/CPunch/gopenfusion/cnet/protocol"
@@ -20,7 +21,7 @@ var (
 )
 
 const (
-	timeout       = 2
+	timeout       = 2 * time.Second
 	maxDummyPeers = 5
 )
 
diff --git a/util/util.go b/util/util.go
index 4208723..60bb296 100644
--- a/util/util.go
+++ b/util/util.go
@@ -9,21 +9,21 @@ func GetTime() uint64 {
 	return uint64(time.Now().UnixMilli())
 }
 
-func SelectWithTimeout(ch <-chan struct{}, seconds int) bool {
+func SelectWithTimeout(ch <-chan struct{}, timeout time.Duration) bool {
 	select {
 	case <-ch:
 		return true
-	case <-time.After(time.Duration(seconds) * time.Second):
+	case <-time.After(timeout):
 		return false
 	}
 }
 
-func WaitWithTimeout(wg *sync.WaitGroup, seconds int) bool {
+func WaitWithTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
 		wg.Wait()
 	}()
 
-	return SelectWithTimeout(done, seconds)
+	return SelectWithTimeout(done, timeout)
 }
diff --git a/util/util_test.go b/util/util_test.go
index b92be3c..ea0f021 100644
--- a/util/util_test.go
+++ b/util/util_test.go
@@ -13,23 +13,23 @@ func TestWaitWithTimeout(t *testing.T) {
 	is := is.New(t)
 	wg := &sync.WaitGroup{}
 	go func() {
-		time.Sleep(2 * time.Second)
+		time.Sleep(1 * time.Second)
 		wg.Done()
 	}()
 
 	wg.Add(1)
-	is.True(!util.WaitWithTimeout(wg, 1)) // timeout should occur
-	is.True(util.WaitWithTimeout(wg, 2))  // timeout shouldn't occur
+	is.True(!util.WaitWithTimeout(wg, 500*time.Millisecond)) // timeout should occur
+	is.True(util.WaitWithTimeout(wg, 750*time.Millisecond))  // timeout shouldn't occur
 }
 
 func TestSelectWithTimeout(t *testing.T) {
 	is := is.New(t)
 	ch := make(chan struct{})
 	go func() {
-		time.Sleep(2 * time.Second)
+		time.Sleep(1 * time.Second)
 		close(ch)
 	}()
 
-	is.True(!util.SelectWithTimeout(ch, 1)) // timeout should occur
-	is.True(util.SelectWithTimeout(ch, 2))  // timeout shouldn't occur
+	is.True(!util.SelectWithTimeout(ch, 500*time.Millisecond)) // timeout should occur
+	is.True(util.SelectWithTimeout(ch, 750*time.Millisecond))  // timeout shouldn't occur
 }