Skip to content
This repository has been archived by the owner on Oct 9, 2023. It is now read-only.

Commit

Permalink
implemented unit tests for ReleaseReservation API call
Browse files Browse the repository at this point in the history 
Signed-off-by: Daniel Rammer <[email protected]>
  • Loading branch information
@hamersaw
hamersaw committed Aug 26, 2021
1 parent b3e16e0 commit 1627131
Show file tree
Hide file tree
Showing 2 changed files with 149 additions and 52 deletions.
170 changes: 118 additions & 52 deletions pkg/manager/impl/reservation_manager_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
"time"

errors2 "github.com/flyteorg/datacatalog/pkg/errors"
errors3 "github.com/flyteorg/datacatalog/pkg/repositories/errors"
"github.com/flyteorg/datacatalog/pkg/repositories/mocks"
"github.com/flyteorg/datacatalog/pkg/repositories/models"
"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog"
Expand Down Expand Up @@ -78,13 +79,6 @@ func TestGetOrReserveArtifact_ArtifactExists(t *testing.T) {
assert.Equal(t, expectedArtifact.ArtifactKey.ArtifactID, artifact.Id)
}

func getDatacatalogRepo() mocks.DataCatalogRepo {
return mocks.DataCatalogRepo{
MockReservationRepo: &mocks.ReservationRepo{},
MockTagRepo: &mocks.TagRepo{},
}
}

func TestGetOrReserveArtifact_CreateReservation(t *testing.T) {
dcRepo := getDatacatalogRepo()

Expand Down Expand Up @@ -132,13 +126,13 @@ func TestGetOrReserveArtifact_CreateReservation(t *testing.T) {
assert.Equal(t, datacatalog.ReservationStatus_ACQUIRED, resp.GetReservationStatus().State)
}

func TestGetOrReserveArtifact_TakeOverReservation(t *testing.T) {
func TestGetOrReserveArtifact_ExtendReservation(t *testing.T) {
dcRepo := getDatacatalogRepo()

setUpTagRepoGetNotFound(&dcRepo)

now := time.Now()
prevExpiresAt := now.Add(time.Second * 10 * time.Duration(-1))
prevExpiresAt := now.Add(time.Second * 10)

setUpReservationRepoGet(&dcRepo, prevExpiresAt)

Expand All @@ -150,7 +144,7 @@ func TestGetOrReserveArtifact_TakeOverReservation(t *testing.T) {
reservation.DatasetName == datasetID.Name &&
reservation.DatasetVersion == datasetID.Version &&
reservation.TagName == tagName &&
reservation.OwnerID == currentOwner &&
reservation.OwnerID == prevOwner &&
reservation.ExpiresAt == now.Add(heartbeatInterval*heartbeatGracePeriodMultiplier)
}),
mock.MatchedBy(func(now time.Time) bool { return true }),
Expand All @@ -162,58 +156,23 @@ func TestGetOrReserveArtifact_TakeOverReservation(t *testing.T) {

req := datacatalog.GetOrReserveArtifactRequest{
ReservationId: &reservationID,
OwnerId: currentOwner,
OwnerId: prevOwner,
}

resp, err := reservationManager.GetOrReserveArtifact(context.Background(), &req)

assert.Nil(t, err)
assert.Equal(t, currentOwner, resp.GetReservationStatus().OwnerId)
assert.Equal(t, prevOwner, resp.GetReservationStatus().OwnerId)
assert.Equal(t, datacatalog.ReservationStatus_ACQUIRED, resp.GetReservationStatus().State)
}

func setUpReservationRepoGet(dcRepo *mocks.DataCatalogRepo, prevExpiresAt time.Time) {
dcRepo.MockReservationRepo.On("Get",
mock.MatchedBy(func(ctx context.Context) bool { return true }),
mock.MatchedBy(func(key models.ReservationKey) bool {
return key.DatasetProject == datasetID.Project &&
key.DatasetDomain == datasetID.Domain &&
key.DatasetVersion == datasetID.Version &&
key.DatasetName == datasetID.Name &&
key.TagName == tagName
})).Return(
models.Reservation{
ReservationKey: getReservationKey(),
OwnerID: prevOwner,
ExpiresAt: prevExpiresAt,
}, nil,
)
}

func setUpTagRepoGetNotFound(dcRepo *mocks.DataCatalogRepo) {
dcRepo.MockTagRepo.On("Get",
mock.Anything,
mock.Anything,
).Return(models.Tag{}, errors2.NewDataCatalogErrorf(codes.NotFound, "entry not found"))
}

func getReservationKey() models.ReservationKey {
return models.ReservationKey{
DatasetProject: project,
DatasetName: name,
DatasetDomain: domain,
DatasetVersion: version,
TagName: tagName,
}
}

func TestGetOrReserveArtifact_ExtendReservation(t *testing.T) {
func TestGetOrReserveArtifact_TakeOverReservation(t *testing.T) {
dcRepo := getDatacatalogRepo()

setUpTagRepoGetNotFound(&dcRepo)

now := time.Now()
prevExpiresAt := now.Add(time.Second * 10)
prevExpiresAt := now.Add(time.Second * 10 * time.Duration(-1))

setUpReservationRepoGet(&dcRepo, prevExpiresAt)

Expand All @@ -225,7 +184,7 @@ func TestGetOrReserveArtifact_ExtendReservation(t *testing.T) {
reservation.DatasetName == datasetID.Name &&
reservation.DatasetVersion == datasetID.Version &&
reservation.TagName == tagName &&
reservation.OwnerID == prevOwner &&
reservation.OwnerID == currentOwner &&
reservation.ExpiresAt == now.Add(heartbeatInterval*heartbeatGracePeriodMultiplier)
}),
mock.MatchedBy(func(now time.Time) bool { return true }),
Expand All @@ -237,13 +196,13 @@ func TestGetOrReserveArtifact_ExtendReservation(t *testing.T) {

req := datacatalog.GetOrReserveArtifactRequest{
ReservationId: &reservationID,
OwnerId: prevOwner,
OwnerId: currentOwner,
}

resp, err := reservationManager.GetOrReserveArtifact(context.Background(), &req)

assert.Nil(t, err)
assert.Equal(t, prevOwner, resp.GetReservationStatus().OwnerId)
assert.Equal(t, currentOwner, resp.GetReservationStatus().OwnerId)
assert.Equal(t, datacatalog.ReservationStatus_ACQUIRED, resp.GetReservationStatus().State)
}

Expand Down Expand Up @@ -272,3 +231,110 @@ func TestGetOrReserveArtifact_AlreadyInProgress(t *testing.T) {
assert.Equal(t, prevOwner, resp.GetReservationStatus().OwnerId)
assert.Equal(t, datacatalog.ReservationStatus_ALREADY_IN_PROGRESS, resp.GetReservationStatus().State)
}

func TestReleaseReservation(t *testing.T) {
dcRepo := getDatacatalogRepo()

now := time.Now()

dcRepo.MockReservationRepo.On("Delete",
mock.MatchedBy(func(ctx context.Context) bool { return true }),
mock.MatchedBy(func(reservationKey models.ReservationKey) bool {
return reservationKey.DatasetProject == datasetID.Project &&
reservationKey.DatasetDomain == datasetID.Domain &&
reservationKey.DatasetName == datasetID.Name &&
reservationKey.DatasetVersion == datasetID.Version &&
reservationKey.TagName == tagName
}),
).Return(nil)

reservationManager := NewReservationManager(&dcRepo,
heartbeatGracePeriodMultiplier, heartbeatInterval,
func() time.Time { return now }, mockScope.NewTestScope())

req := datacatalog.ReleaseReservationRequest{
ReservationId: &reservationID,
OwnerId: currentOwner,
}

_, err := reservationManager.ReleaseReservation(context.Background(), &req)

assert.Nil(t, err)
}

func TestReleaseReservation_Failure(t *testing.T) {
dcRepo := getDatacatalogRepo()

now := time.Now()
reservationErr := errors3.GetMissingEntityError("Reservation",
&datacatalog.ReservationID{
DatasetId: &datasetID,
TagName: tagName,
})

dcRepo.MockReservationRepo.On("Delete",
mock.MatchedBy(func(ctx context.Context) bool { return true }),
mock.MatchedBy(func(reservationKey models.ReservationKey) bool {
return reservationKey.DatasetProject == datasetID.Project &&
reservationKey.DatasetDomain == datasetID.Domain &&
reservationKey.DatasetName == datasetID.Name &&
reservationKey.DatasetVersion == datasetID.Version &&
reservationKey.TagName == tagName
}),
).Return(reservationErr)

reservationManager := NewReservationManager(&dcRepo,
heartbeatGracePeriodMultiplier, heartbeatInterval,
func() time.Time { return now }, mockScope.NewTestScope())

req := datacatalog.ReleaseReservationRequest{
ReservationId: &reservationID,
OwnerId: currentOwner,
}

_, err := reservationManager.ReleaseReservation(context.Background(), &req)

assert.Equal(t, err, reservationErr)
}

func getDatacatalogRepo() mocks.DataCatalogRepo {
return mocks.DataCatalogRepo{
MockReservationRepo: &mocks.ReservationRepo{},
MockTagRepo: &mocks.TagRepo{},
}
}

func setUpReservationRepoGet(dcRepo *mocks.DataCatalogRepo, prevExpiresAt time.Time) {
dcRepo.MockReservationRepo.On("Get",
mock.MatchedBy(func(ctx context.Context) bool { return true }),
mock.MatchedBy(func(key models.ReservationKey) bool {
return key.DatasetProject == datasetID.Project &&
key.DatasetDomain == datasetID.Domain &&
key.DatasetVersion == datasetID.Version &&
key.DatasetName == datasetID.Name &&
key.TagName == tagName
})).Return(
models.Reservation{
ReservationKey: getReservationKey(),
OwnerID: prevOwner,
ExpiresAt: prevExpiresAt,
}, nil,
)
}

func setUpTagRepoGetNotFound(dcRepo *mocks.DataCatalogRepo) {
dcRepo.MockTagRepo.On("Get",
mock.Anything,
mock.Anything,
).Return(models.Tag{}, errors2.NewDataCatalogErrorf(codes.NotFound, "entry not found"))
}

func getReservationKey() models.ReservationKey {
return models.ReservationKey{
DatasetProject: project,
DatasetName: name,
DatasetDomain: domain,
DatasetVersion: version,
TagName: tagName,
}
}
31 changes: 31 additions & 0 deletions pkg/repositories/gormimpl/reservation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,37 @@ func TestCreate(t *testing.T) {
assert.NoError(t, err)
}

func TestDelete(t *testing.T) {
GlobalMock := mocket.Catcher.Reset()
GlobalMock.Logging = true
expectedReservation := GetReservation()

GlobalMock.NewMock().WithQuery(
`DELETE FROM "reservations" WHERE "reservations"."dataset_project" = $1 AND "reservations"."dataset_name" = $2 AND "reservations"."dataset_domain" = $3 AND "reservations"."dataset_version" = $4 AND "reservations"."tag_name" = $5`,
).WithRowsNum(1)

reservationRepo := getReservationRepo(t)

err := reservationRepo.Delete(context.Background(), expectedReservation.ReservationKey)
assert.NoError(t, err)
}

func TestDeleteFailure(t *testing.T) {
GlobalMock := mocket.Catcher.Reset()
GlobalMock.Logging = true
expectedReservation := GetReservation()

GlobalMock.NewMock().WithQuery(
`DELETE FROM "reservations" WHERE "reservations"."dataset_project" = $1 AND "reservations"."dataset_name" = $2 AND "reservations"."dataset_domain" = $3 AND "reservations"."dataset_version" = $4 AND "reservations"."tag_name" = $5`,
).WithRowsNum(0)

reservationRepo := getReservationRepo(t)

err := reservationRepo.Delete(context.Background(), expectedReservation.ReservationKey)
assert.Error(t, err)
assert.Equal(t, "missing entity of type Reservation with identifier dataset_id:<project:\"testProject\" name:\"testDataset\" domain:\"testDomain\" version:\"testVersion\" > tag_name:\"testTag\" ", err.Error())
}

func TestGet(t *testing.T) {
expectedReservation := GetReservation()

Expand Down

0 comments on commit 1627131

Please sign in to comment.