From 136b8f200db5e6b534d769ec5b7452ff4e1f2213 Mon Sep 17 00:00:00 2001
From: Andy Goldstein <andy.goldstein@gmail.com>
Date: Wed, 13 Sep 2017 16:04:46 -0400
Subject: [PATCH] Support restore logs when backup name has -

Signed-off-by: Andy Goldstein <andy.goldstein@gmail.com>
---
 pkg/cloudprovider/backup_service.go      |  6 +-
 pkg/cloudprovider/backup_service_test.go | 70 ++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/pkg/cloudprovider/backup_service.go b/pkg/cloudprovider/backup_service.go
index 6543043d4..521e32786 100644
--- a/pkg/cloudprovider/backup_service.go
+++ b/pkg/cloudprovider/backup_service.go
@@ -211,7 +211,11 @@ func (br *backupService) CreateSignedURL(target api.DownloadTarget, bucket strin
 		return br.objectStorage.CreateSignedURL(bucket, getBackupLogKey(target.Name), ttl)
 	case api.DownloadTargetKindRestoreLog:
 		// restore name is formatted as <backup name>-<timestamp>
-		backup := strings.Split(target.Name, "-")[0]
+		i := strings.LastIndex(target.Name, "-")
+		if i < 0 {
+			i = len(target.Name)
+		}
+		backup := target.Name[0:i]
 		return br.objectStorage.CreateSignedURL(bucket, getRestoreLogKey(backup, target.Name), ttl)
 	default:
 		return "", fmt.Errorf("unsupported download target kind %q", target.Kind)
diff --git a/pkg/cloudprovider/backup_service_test.go b/pkg/cloudprovider/backup_service_test.go
index 455636393..afb705d92 100644
--- a/pkg/cloudprovider/backup_service_test.go
+++ b/pkg/cloudprovider/backup_service_test.go
@@ -24,6 +24,7 @@ import (
 	"io/ioutil"
 	"strings"
 	"testing"
+	"time"
 
 	testutil "github.com/heptio/ark/pkg/util/test"
 	"github.com/stretchr/testify/assert"
@@ -241,6 +242,75 @@ func TestGetAllBackups(t *testing.T) {
 	}
 }
 
+func TestCreateSignedURL(t *testing.T) {
+	tests := []struct {
+		name        string
+		targetKind  api.DownloadTargetKind
+		targetName  string
+		expectedKey string
+	}{
+		{
+			name:        "backup contents",
+			targetKind:  api.DownloadTargetKindBackupContents,
+			targetName:  "my-backup",
+			expectedKey: "my-backup/my-backup.tar.gz",
+		},
+		{
+			name:        "backup log",
+			targetKind:  api.DownloadTargetKindBackupLog,
+			targetName:  "my-backup",
+			expectedKey: "my-backup/my-backup-logs.gz",
+		},
+		{
+			name:        "scheduled backup contents",
+			targetKind:  api.DownloadTargetKindBackupContents,
+			targetName:  "my-backup-20170913154901",
+			expectedKey: "my-backup-20170913154901/my-backup-20170913154901.tar.gz",
+		},
+		{
+			name:        "scheduled backup log",
+			targetKind:  api.DownloadTargetKindBackupLog,
+			targetName:  "my-backup-20170913154901",
+			expectedKey: "my-backup-20170913154901/my-backup-20170913154901-logs.gz",
+		},
+		{
+			name:        "restore log - backup has no dash",
+			targetKind:  api.DownloadTargetKindRestoreLog,
+			targetName:  "b-20170913154901",
+			expectedKey: "b/restore-b-20170913154901-logs.gz",
+		},
+		{
+			name:        "restore log - backup has 1 dash",
+			targetKind:  api.DownloadTargetKindRestoreLog,
+			targetName:  "b-cool-20170913154901",
+			expectedKey: "b-cool/restore-b-cool-20170913154901-logs.gz",
+		},
+		{
+			name:        "restore log - backup has multiple dashes (e.g. restore of scheduled backup)",
+			targetKind:  api.DownloadTargetKindRestoreLog,
+			targetName:  "b-cool-20170913154901-20170913154902",
+			expectedKey: "b-cool-20170913154901/restore-b-cool-20170913154901-20170913154902-logs.gz",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			objectStorage := &testutil.ObjectStorageAdapter{}
+			backupService := NewBackupService(objectStorage)
+
+			target := api.DownloadTarget{
+				Kind: test.targetKind,
+				Name: test.targetName,
+			}
+			objectStorage.On("CreateSignedURL", "bucket", test.expectedKey, time.Duration(0)).Return("url", nil)
+			url, err := backupService.CreateSignedURL(target, "bucket", 0)
+			require.NoError(t, err)
+			assert.Equal(t, "url", url)
+			objectStorage.AssertExpectations(t)
+		})
+	}
+}
+
 func jsonMarshal(obj interface{}) []byte {
 	res, err := json.Marshal(obj)
 	if err != nil {