diff --git a/pkg/backend/proxy_test.go b/pkg/backend/proxy_test.go
index 7a201f1..55b35d8 100644
--- a/pkg/backend/proxy_test.go
+++ b/pkg/backend/proxy_test.go
@@ -85,7 +85,7 @@ func TestProxyBackendRun(t *testing.T) {
 
 	checkCh := make(chan struct{})
 	go func() {
-		<-ctx.Done()
+		pb.Wait()
 		close(checkCh)
 	}()
 
diff --git a/pkg/backend/transparent_test.go b/pkg/backend/transparent_test.go
index db4f415..dc7d041 100644
--- a/pkg/backend/transparent_test.go
+++ b/pkg/backend/transparent_test.go
@@ -114,7 +114,7 @@ func TestTransparentBackendRun(t *testing.T) {
 
 	checkCh := make(chan struct{})
 	go func() {
-		<-ctx.Done()
+		be.Wait()
 		close(checkCh)
 	}()