From 6d5530eedc13bdf5a5da3917e2c7c4dabed3bebc Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Wed, 25 Aug 2021 17:29:01 -0500 Subject: [PATCH] added unit tests for reservation transformer Signed-off-by: Daniel Rammer --- pkg/manager/impl/reservation_manager.go | 6 +- pkg/manager/impl/reservation_manager_test.go | 25 ++++---- pkg/repositories/transformers/reservation.go | 26 ++++---- .../transformers/reservation_test.go | 60 +++++++++++++++---- 4 files changed, 75 insertions(+), 42 deletions(-) diff --git a/pkg/manager/impl/reservation_manager.go b/pkg/manager/impl/reservation_manager.go index 03c801ae..f6f94543 100644 --- a/pkg/manager/impl/reservation_manager.go +++ b/pkg/manager/impl/reservation_manager.go @@ -163,7 +163,7 @@ func (r *reservationManager) tryAcquireReservation(ctx context.Context, reservat logger.Debugf(ctx, "Reservation: %+v is held by %s", reservationKey, reservation.OwnerID) reservationStatus, err := transformers.CreateReservationStatus( - reservation, r.heartbeatInterval, + &reservation, r.heartbeatInterval, datacatalog.ReservationStatus_ALREADY_IN_PROGRESS) if err != nil { return reservationStatus, err @@ -183,7 +183,7 @@ func (r *reservationManager) tryAcquireReservation(ctx context.Context, reservat } reservationStatus, err := transformers.CreateReservationStatus( - rsv1, r.heartbeatInterval, + &rsv1, r.heartbeatInterval, datacatalog.ReservationStatus_ALREADY_IN_PROGRESS) if err != nil { return reservationStatus, err @@ -197,7 +197,7 @@ func (r *reservationManager) tryAcquireReservation(ctx context.Context, reservat } // Reservation has been acquired or extended without error - reservationStatus, err := transformers.CreateReservationStatus(newReservation, + reservationStatus, err := transformers.CreateReservationStatus(&newReservation, r.heartbeatInterval, datacatalog.ReservationStatus_ACQUIRED) if err != nil { return reservationStatus, err diff --git a/pkg/manager/impl/reservation_manager_test.go b/pkg/manager/impl/reservation_manager_test.go index c4ca69aa..d14181fc 100644 --- a/pkg/manager/impl/reservation_manager_test.go +++ b/pkg/manager/impl/reservation_manager_test.go @@ -29,6 +29,10 @@ var datasetID = datacatalog.DatasetID{ Domain: domain, Version: version, } +var reservationID = datacatalog.ReservationID{ + DatasetId: &datasetID, + TagName: tagName, +} var heartbeatInterval = time.Second * 5 var heartbeatGracePeriodMultiplier = time.Second * 3 var prevOwner = "prevOwner" @@ -63,8 +67,7 @@ func TestGetOrReserveArtifact_ArtifactExists(t *testing.T) { heartbeatInterval, time.Now, mockScope.NewTestScope()) req := datacatalog.GetOrReserveArtifactRequest{ - DatasetId: &datasetID, - TagName: tagName, + ReservationId: &reservationID, OwnerId: currentOwner, } @@ -118,9 +121,8 @@ func TestGetOrReserveArtifact_CreateReservation(t *testing.T) { func() time.Time { return now }, mockScope.NewTestScope()) req := datacatalog.GetOrReserveArtifactRequest{ - DatasetId: &datasetID, - TagName: tagName, - OwnerId: currentOwner, + ReservationId: &reservationID, + OwnerId: currentOwner, } resp, err := reservationManager.GetOrReserveArtifact(context.Background(), &req) @@ -159,9 +161,8 @@ func TestGetOrReserveArtifact_TakeOverReservation(t *testing.T) { func() time.Time { return now }, mockScope.NewTestScope()) req := datacatalog.GetOrReserveArtifactRequest{ - DatasetId: &datasetID, - TagName: tagName, - OwnerId: currentOwner, + ReservationId: &reservationID, + OwnerId: currentOwner, } resp, err := reservationManager.GetOrReserveArtifact(context.Background(), &req) @@ -235,8 +236,7 @@ func TestGetOrReserveArtifact_ExtendReservation(t *testing.T) { func() time.Time { return now }, mockScope.NewTestScope()) req := datacatalog.GetOrReserveArtifactRequest{ - DatasetId: &datasetID, - TagName: tagName, + ReservationId: &reservationID, OwnerId: prevOwner, } @@ -262,9 +262,8 @@ func TestGetOrReserveArtifact_AlreadyInProgress(t *testing.T) { func() time.Time { return now }, mockScope.NewTestScope()) req := datacatalog.GetOrReserveArtifactRequest{ - DatasetId: &datasetID, - TagName: tagName, - OwnerId: currentOwner, + ReservationId: &reservationID, + OwnerId: currentOwner, } resp, err := reservationManager.GetOrReserveArtifact(context.Background(), &req) diff --git a/pkg/repositories/transformers/reservation.go b/pkg/repositories/transformers/reservation.go index d27c9b45..71a471a8 100644 --- a/pkg/repositories/transformers/reservation.go +++ b/pkg/repositories/transformers/reservation.go @@ -12,17 +12,6 @@ import ( "google.golang.org/grpc/codes" ) -// TODO - remove -/*func ToReservationKey(datasetID datacatalog.DatasetID, tagName string) models.ReservationKey { - return models.ReservationKey{ - DatasetProject: datasetID.Project, - DatasetName: datasetID.Name, - DatasetDomain: datasetID.Domain, - DatasetVersion: datasetID.Version, - TagName: tagName, - } -}*/ - func FromReservationID(reservationID *datacatalog.ReservationID) models.ReservationKey { datasetID := reservationID.DatasetId @@ -35,17 +24,26 @@ func FromReservationID(reservationID *datacatalog.ReservationID) models.Reservat } } -func CreateReservationStatus(reservation models.Reservation, heartbeatInterval time.Duration, state datacatalog.ReservationStatus_State) (datacatalog.ReservationStatus, error) { - expiresAtPb, err := ptypes.TimestampProto(reservation.CreatedAt) +func CreateReservationStatus(reservation *models.Reservation, heartbeatInterval time.Duration, state datacatalog.ReservationStatus_State) (datacatalog.ReservationStatus, error) { + expiresAtPb, err := ptypes.TimestampProto(reservation.ExpiresAt) if err != nil { return datacatalog.ReservationStatus{}, errors.NewDataCatalogErrorf(codes.Internal, "failed to serialize expires at time") } heartbeatIntervalPb := ptypes.DurationProto(heartbeatInterval) return datacatalog.ReservationStatus{ + ReservationId: &datacatalog.ReservationID { + DatasetId: &datacatalog.DatasetID { + Project: reservation.DatasetProject, + Domain: reservation.DatasetDomain, + Name: reservation.DatasetName, + Version: reservation.DatasetVersion, + }, + TagName: reservation.TagName, + }, + OwnerId: reservation.OwnerID, State: state, ExpiresAt: expiresAtPb, HeartbeatInterval: heartbeatIntervalPb, - OwnerId: reservation.OwnerID, }, nil } diff --git a/pkg/repositories/transformers/reservation_test.go b/pkg/repositories/transformers/reservation_test.go index 998c025b..496a1487 100644 --- a/pkg/repositories/transformers/reservation_test.go +++ b/pkg/repositories/transformers/reservation_test.go @@ -2,23 +2,59 @@ package transformers import ( "testing" + "time" + "github.com/flyteorg/datacatalog/pkg/repositories/models" "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog" "github.com/magiconair/properties/assert" ) -func TestToReservationKey(t *testing.T) { - datasetID := datacatalog.DatasetID{ - Project: "p", - Name: "n", - Domain: "d", - Version: "v", +func TestFromReservationID(t *testing.T) { + reservationID := datacatalog.ReservationID{ + DatasetId: &datacatalog.DatasetID{ + Project: "p", + Name: "n", + Domain: "d", + Version: "v", + }, + TagName: "t", } - reservationKey := ToReservationKey(datasetID, "t") - assert.Equal(t, datasetID.Project, reservationKey.DatasetProject) - assert.Equal(t, datasetID.Name, reservationKey.DatasetName) - assert.Equal(t, datasetID.Domain, reservationKey.DatasetDomain) - assert.Equal(t, datasetID.Version, reservationKey.DatasetVersion) - assert.Equal(t, "t", reservationKey.TagName) + reservationKey := FromReservationID(&reservationID) + assert.Equal(t, reservationKey.DatasetProject, reservationID.DatasetId.Project) + assert.Equal(t, reservationKey.DatasetName, reservationID.DatasetId.Name) + assert.Equal(t, reservationKey.DatasetDomain, reservationID.DatasetId.Domain) + assert.Equal(t, reservationKey.DatasetVersion, reservationID.DatasetId.Version) + assert.Equal(t, reservationKey.TagName, reservationID.TagName) +} + +func TestCreateReservationStatus(t *testing.T) { + now := time.Now() + heartbeatInterval := time.Duration(time.Second * 5) + reservation := models.Reservation { + ReservationKey: models.ReservationKey { + DatasetProject: "p", + DatasetName: "n", + DatasetDomain: "d", + DatasetVersion: "v", + TagName: "t", + }, + OwnerID: "o", + ExpiresAt: now, + } + + reservationStatus, err := CreateReservationStatus(&reservation, heartbeatInterval, datacatalog.ReservationStatus_ACQUIRED) + + assert.Equal(t, err, nil) + assert.Equal(t, reservationStatus.OwnerId, reservation.OwnerID) + assert.Equal(t, reservationStatus.ExpiresAt.AsTime(), reservation.ExpiresAt.UTC()) + + reservationID := reservationStatus.ReservationId + assert.Equal(t, reservationID.TagName, reservation.TagName) + + datasetID := reservationID.DatasetId + assert.Equal(t, datasetID.Project, reservation.DatasetProject) + assert.Equal(t, datasetID.Name, reservation.DatasetName) + assert.Equal(t, datasetID.Domain, reservation.DatasetDomain) + assert.Equal(t, datasetID.Version, reservation.DatasetVersion) }