diff --git a/pkg/expect/expect_test.go b/pkg/expect/expect_test.go
index d5a22848d..65746851c 100644
--- a/pkg/expect/expect_test.go
+++ b/pkg/expect/expect_test.go
@@ -75,7 +75,8 @@ func TestEcho(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	l, eerr := ep.Expect("world")
+	ctx := context.Background()
+	l, eerr := ep.ExpectWithContext(ctx, "world")
 	if eerr != nil {
 		t.Fatal(eerr)
 	}
@@ -86,7 +87,7 @@ func TestEcho(t *testing.T) {
 	if cerr := ep.Close(); cerr != nil {
 		t.Fatal(cerr)
 	}
-	if _, eerr = ep.Expect("..."); eerr == nil {
+	if _, eerr = ep.ExpectWithContext(ctx, "..."); eerr == nil {
 		t.Fatalf("expected error on closed expect process")
 	}
 }
@@ -97,7 +98,7 @@ func TestLineCount(t *testing.T) {
 		t.Fatal(err)
 	}
 	wstr := "3"
-	l, eerr := ep.Expect(wstr)
+	l, eerr := ep.ExpectWithContext(context.Background(), wstr)
 	if eerr != nil {
 		t.Fatal(eerr)
 	}
@@ -120,7 +121,7 @@ func TestSend(t *testing.T) {
 	if err := ep.Send("a\r"); err != nil {
 		t.Fatal(err)
 	}
-	if _, err := ep.Expect("b"); err != nil {
+	if _, err := ep.ExpectWithContext(context.Background(), "b"); err != nil {
 		t.Fatal(err)
 	}
 	if err := ep.Stop(); err != nil {