diff --git a/pkg/probe/probe_test.go b/pkg/probe/probe_test.go
index 37e8013..9edc561 100644
--- a/pkg/probe/probe_test.go
+++ b/pkg/probe/probe_test.go
@@ -18,12 +18,14 @@
 import (
 	"context"
 	"encoding/json"
-	"github.com/opencord/voltha-lib-go/v3/pkg/log"
-	"github.com/stretchr/testify/assert"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
+
+	"github.com/opencord/voltha-lib-go/v3/pkg/log"
+	"github.com/stretchr/testify/assert"
 )
 
 func init() {
@@ -343,7 +345,8 @@
 }
 
 func TestGetProbeFromContextMssing(t *testing.T) {
-	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
 	pc := GetProbeFromContext(ctx)
 	assert.Nil(t, pc, "Context had a non-nil probe when it should have been nil")
 }
@@ -375,7 +378,8 @@
 func TestUpdateStatusFromContextWithoutProbe(t *testing.T) {
 	p := &Probe{}
 	p.RegisterService("one")
-	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
 	UpdateStatusFromContext(ctx, "one", ServiceStatusRunning)
 
 	assert.Equal(t, 1, len(p.status), "wrong number of services")