diff --git a/go.mod b/go.mod
index 7b1517867c95c6030fe574bc5f4335c316c28db6..8046f082ba2a4ac426d4725c1463fe67283190f1 100644
--- a/go.mod
+++ b/go.mod
@@ -10,7 +10,7 @@ require (
 	github.com/stretchr/testify v1.7.0
 	github.com/thijsheijden/alice v0.1.18
 	golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4 // indirect
-	golang.org/x/sys v0.0.0-20210511113859-b0526f3d8744 // indirect
+	golang.org/x/sys v0.0.0-20210514084401-e8d321eab015 // indirect
 	google.golang.org/grpc v1.37.0
 	google.golang.org/protobuf v1.26.0
 )
diff --git a/go.sum b/go.sum
index 805abd263bd64e9e5ba50bde46d486fb634462bd..8685c26d1b1be88c1e1f46e6a0cd6a475bee8f85 100644
--- a/go.sum
+++ b/go.sum
@@ -132,8 +132,8 @@ golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210511113859-b0526f3d8744 h1:yhBbb4IRs2HS9PPlAg6DMC6mUOKexJBNsLf4Z+6En1Q=
-golang.org/x/sys v0.0.0-20210511113859-b0526f3d8744/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210514084401-e8d321eab015 h1:hZR0X1kPW+nwyJ9xRxqZk1vx5RUObAPBdKVvXPDUH/E=
+golang.org/x/sys v0.0.0-20210514084401-e8d321eab015/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k=
diff --git a/internal/usecases/consume/consume_test.go b/internal/usecases/consume/consume_test.go
index 92b9cea86c1db7848b9cdd39d8d6400ed251ee6d..0e750c82a906da96e5102adc50f76157998abe6c 100644
--- a/internal/usecases/consume/consume_test.go
+++ b/internal/usecases/consume/consume_test.go
@@ -1,14 +1,15 @@
 package consume
 
 import (
+	"context"
 	"encoding/json"
-	mockkeyvaluedriver "query-service/internal/drivers/keyvaluedriver/mock"
 	mockdatabaseinfo "query-service/internal/usecases/databaseinfo/mock"
 	"query-service/internal/usecases/produce"
 	mockrequest "query-service/internal/usecases/request/mock"
 	"testing"
 
 	"git.science.uu.nl/datastrophe/broker"
+	"git.science.uu.nl/datastrophe/keyvaluestore"
 	mockconvertquery "git.science.uu.nl/datastrophe/query-conversion/aql"
 
 	"github.com/stretchr/testify/assert"
@@ -18,7 +19,7 @@ func TestHandleCorrectMessage(t *testing.T) {
 	// Create a mock broker
 	mockBroker := broker.NewMockDriver().(*broker.MockDriver)
 	// Create mock key value store
-	keyValueStore := mockkeyvaluedriver.CreateKeyValueStore()
+	keyValueStore := keyvaluestore.NewMockDriver()
 	// Create new producer service
 	producerService := produce.NewService(mockBroker, keyValueStore)
 	producerService.Start()
@@ -37,7 +38,7 @@ func TestHandleCorrectMessage(t *testing.T) {
 	mockQueue := "mock-queue"
 
 	// Set the test-session sessionID queue to mock-queue in key value store
-	keyValueStore.Set(&mockSession, &mockQueue)
+	keyValueStore.Set(context.Background(), mockSession, mockQueue)
 
 	// Create headers containing a sessionID
 	headers := make(map[string]interface{})
@@ -75,7 +76,7 @@ func TestHandleMessageNoSessionID(t *testing.T) {
 	// Create a mock broker
 	mockBroker := broker.NewMockDriver().(*broker.MockDriver)
 	// Create mock key value store
-	keyValueStore := mockkeyvaluedriver.CreateKeyValueStore()
+	keyValueStore := keyvaluestore.NewMockDriver()
 	// Create new producer service
 	producerService := produce.NewService(mockBroker, keyValueStore)
 	// Create new convert query service
@@ -109,7 +110,7 @@ func TestFailToConvertQuery(t *testing.T) {
 	// Create a mock broker
 	mockBroker := broker.NewMockDriver().(*broker.MockDriver)
 	// Create mock key value store
-	keyValueStore := mockkeyvaluedriver.CreateKeyValueStore()
+	keyValueStore := keyvaluestore.NewMockDriver()
 	// Create new producer service
 	producerService := produce.NewService(mockBroker, keyValueStore)
 	producerService.Start()
@@ -128,7 +129,7 @@ func TestFailToConvertQuery(t *testing.T) {
 	mockQueue := "mock-queue"
 
 	// Set the test-session sessionID queue to mock-queue in key value store
-	keyValueStore.Set(&mockSession, &mockQueue)
+	keyValueStore.Set(context.Background(), mockSession, mockQueue)
 
 	// Create headers containing a sessionID
 	headers := make(map[string]interface{})
@@ -159,7 +160,7 @@ func TestArangoError(t *testing.T) {
 	// Create a mock broker
 	mockBroker := broker.NewMockDriver().(*broker.MockDriver)
 	// Create mock key value store
-	keyValueStore := mockkeyvaluedriver.CreateKeyValueStore()
+	keyValueStore := keyvaluestore.NewMockDriver()
 	// Create new producer service
 	producerService := produce.NewService(mockBroker, keyValueStore)
 	producerService.Start()
@@ -178,7 +179,7 @@ func TestArangoError(t *testing.T) {
 	mockQueue := "mock-queue"
 
 	// Set the test-session sessionID queue to mock-queue in key value store
-	keyValueStore.Set(&mockSession, &mockQueue)
+	keyValueStore.Set(context.Background(), mockSession, mockQueue)
 
 	// Create headers containing a sessionID
 	headers := make(map[string]interface{})