added unit tests for reservation transformer

Signed-off-by: Daniel Rammer <[email protected]>

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

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)

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
 }

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)
 }