discovery: try to fix timeout
Try to fix test timeout failure.
(Not sure yet whether this would fix it.)
Change-Id: Ib9457d4bde8f17e134ac0160c768dc1660dc6302
diff --git a/lib/discovery/plugins/testutil/util.go b/lib/discovery/plugins/testutil/util.go
index abf743b..5178542 100644
--- a/lib/discovery/plugins/testutil/util.go
+++ b/lib/discovery/plugins/testutil/util.go
@@ -35,7 +35,12 @@
func Scan(ctx *context.T, p idiscovery.Plugin, interfaceName string) (<-chan *idiscovery.AdInfo, func(), error) {
ctx, cancel := context.WithCancel(ctx)
scanCh := make(chan *idiscovery.AdInfo)
- callback := func(ad *idiscovery.AdInfo) { scanCh <- ad }
+ callback := func(ad *idiscovery.AdInfo) {
+ select {
+ case scanCh <- ad:
+ case <-ctx.Done():
+ }
+ }
var wg sync.WaitGroup
wg.Add(1)
if err := p.Scan(ctx, interfaceName, callback, wg.Done); err != nil {
diff --git a/lib/discovery/test/directory_test.go b/lib/discovery/test/directory_test.go
index 914923b..8b2422b 100644
--- a/lib/discovery/test/directory_test.go
+++ b/lib/discovery/test/directory_test.go
@@ -153,7 +153,13 @@
}
updateCh := make(chan *idiscovery.AdInfo)
- if err = mockPlugin.Scan(ctx, "", func(ad *idiscovery.AdInfo) { updateCh <- ad }, func() {}); err != nil {
+ callback := func(ad *idiscovery.AdInfo) {
+ select {
+ case updateCh <- ad:
+ case <-ctx.Done():
+ }
+ }
+ if err = mockPlugin.Scan(ctx, "", callback, func() {}); err != nil {
t.Fatal(err)
}
<-updateCh
