diff --git a/core/block/object/objectgraph/graph_test.go b/core/block/object/objectgraph/graph_test.go
index 4f58addba..2c0c8eaab 100644
--- a/core/block/object/objectgraph/graph_test.go
+++ b/core/block/object/objectgraph/graph_test.go
@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/mock"
 
 	"github.com/anyproto/anytype-heart/core/relationutils"
+	"github.com/anyproto/anytype-heart/core/subscription"
 	"github.com/anyproto/anytype-heart/core/subscription/mock_subscription"
 	"github.com/anyproto/anytype-heart/pb"
 	"github.com/anyproto/anytype-heart/pkg/lib/bundle"
@@ -51,7 +52,7 @@ func Test(t *testing.T) {
 			{Relation: bundle.MustGetRelation(bundle.RelationKeyAuthor)},
 			{Relation: bundle.MustGetRelation(bundle.RelationKeyAttachments)},
 		}, nil)
-		fixture.subscriptionServiceMock.EXPECT().Search(mock.Anything).Return(&pb.RpcObjectSearchSubscribeResponse{
+		fixture.subscriptionServiceMock.EXPECT().Search(mock.Anything).Return(&subscription.SubscribeResponse{
 			Records: []*types.Struct{},
 		}, nil)
 		fixture.subscriptionServiceMock.EXPECT().Unsubscribe(mock.Anything).Return(nil)
@@ -73,7 +74,7 @@ func Test(t *testing.T) {
 			{Relation: bundle.MustGetRelation(bundle.RelationKeyAssignee)},
 			{Relation: bundle.MustGetRelation(bundle.RelationKeyAttachments)},
 		}, nil)
-		fixture.subscriptionServiceMock.EXPECT().Search(mock.Anything).Return(&pb.RpcObjectSearchSubscribeResponse{
+		fixture.subscriptionServiceMock.EXPECT().Search(mock.Anything).Return(&subscription.SubscribeResponse{
 			Records: []*types.Struct{
 				{Fields: map[string]*types.Value{
 					bundle.RelationKeyId.String():       pbtypes.String("id1"),
diff --git a/tests/integration/util_test.go b/tests/integration/util_test.go
index e5463b67e..67e6523e0 100644
--- a/tests/integration/util_test.go
+++ b/tests/integration/util_test.go
@@ -115,7 +115,7 @@ func newTestSubscription(t *testing.T, app *testApplication, keys []domain.Relat
 	}
 	subscriptionId := bson.NewObjectId().Hex()
 	subscriptionService := getService[subscription.Service](app)
-	_, err := subscriptionService.Search(pb.RpcObjectSearchSubscribeRequest{
+	_, err := subscriptionService.Search(subscription.SubscribeRequest{
 		SubId:   subscriptionId,
 		Keys:    keysConverted,
 		Filters: filters,