3pp-mirror/pomerium
1
0
Fork 0
mirror of https://github.com/pomerium/pomerium.git synced 2025-08-01 07:50:26 +02:00
Code Issues Projects Releases Packages Wiki Activity

databroker: implement leases (#2172)

Browse source 
* databroker: implement leases

* return error

* handle gRPC errors
This commit is contained in:
Caleb Doxsey 2021-05-10 13:30:25 -06:00 • committed by GitHub
parent a54d43b937
commit 94aa0b1a48
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 2135 additions and 149 deletions
Download patch file Download diff file Expand all files Collapse all files

138
pkg/grpc/databroker/leaser_test.go Normal file
View file

@ -0,0 +1,138 @@
package databroker_test
import (
"context"
"errors"
"testing"
"time"
"github.com/golang/mock/gomock"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/types/known/durationpb"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/grpc/databroker/mock_databroker"
)
func TestLeaser(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
exitErr := errors.New("EXIT")
t.Run("acquires lease", func(t *testing.T) {
client := mock_databroker.NewMockDataBrokerServiceClient(ctrl)
client.EXPECT().
AcquireLease(gomock.Any(), &databroker.AcquireLeaseRequest{
Name: "TEST",
Duration: durationpb.New(time.Second * 30),
}).
Return(&databroker.AcquireLeaseResponse{
Id: "lease1",
}, nil).
Times(1)
client.EXPECT().
ReleaseLease(gomock.Any(), &databroker.ReleaseLeaseRequest{
Name: "TEST",
Id: "lease1",
}).
Times(1)
handler := mock_databroker.NewMockLeaserHandler(ctrl)
handler.EXPECT().
GetDataBrokerServiceClient().
Return(client).
AnyTimes()
handler.EXPECT().
RunLeased(gomock.Any()).
Return(exitErr).
Times(1)
leaser := databroker.NewLeaser("TEST", time.Second*30, handler)
err := leaser.Run(context.Background())
assert.Equal(t, exitErr, err)
})
t.Run("retries acquire", func(t *testing.T) {
client := mock_databroker.NewMockDataBrokerServiceClient(ctrl)
client.EXPECT().
AcquireLease(gomock.Any(), &databroker.AcquireLeaseRequest{
Name: "TEST",
Duration: durationpb.New(time.Second * 30),
}).
Return(nil, status.Error(codes.Unavailable, "UNAVAILABLE")).
Times(2)
client.EXPECT().
AcquireLease(gomock.Any(), &databroker.AcquireLeaseRequest{
Name: "TEST",
Duration: durationpb.New(time.Second * 30),
}).
Return(&databroker.AcquireLeaseResponse{
Id: "lease1",
}, nil).
Times(1)
client.EXPECT().
ReleaseLease(gomock.Any(), &databroker.ReleaseLeaseRequest{
Name: "TEST",
Id: "lease1",
}).
Times(1)
handler := mock_databroker.NewMockLeaserHandler(ctrl)
handler.EXPECT().
GetDataBrokerServiceClient().
Return(client).
AnyTimes()
handler.EXPECT().
RunLeased(gomock.Any()).
Return(exitErr).
Times(1)
leaser := databroker.NewLeaser("TEST", time.Second*30, handler)
err := leaser.Run(context.Background())
assert.Equal(t, exitErr, err)
})
t.Run("renews", func(t *testing.T) {
client := mock_databroker.NewMockDataBrokerServiceClient(ctrl)
client.EXPECT().
AcquireLease(gomock.Any(), &databroker.AcquireLeaseRequest{
Name: "TEST",
Duration: durationpb.New(time.Millisecond),
}).
Return(&databroker.AcquireLeaseResponse{
Id: "lease1",
}, nil).
Times(1)
client.EXPECT().
RenewLease(gomock.Any(), &databroker.RenewLeaseRequest{
Name: "TEST",
Id: "lease1",
Duration: durationpb.New(time.Millisecond),
}).
MinTimes(1)
client.EXPECT().
ReleaseLease(gomock.Any(), &databroker.ReleaseLeaseRequest{
Name: "TEST",
Id: "lease1",
}).
Times(1)
handler := mock_databroker.NewMockLeaserHandler(ctrl)
handler.EXPECT().
GetDataBrokerServiceClient().
Return(client).
AnyTimes()
handler.EXPECT().
RunLeased(gomock.Any()).
DoAndReturn(func(ctx context.Context) error {
time.Sleep(time.Millisecond * 20)
return exitErr
}).
Times(1)
leaser := databroker.NewLeaser("TEST", time.Millisecond, handler)
err := leaser.Run(context.Background())
assert.Equal(t, exitErr, err)
})
}