From 5060704a768308ece761c9dc88ccd064d92ebe59 Mon Sep 17 00:00:00 2001
From: Matthew Wong <matthewwong.wong@alum.utoronto.ca>
Date: Fri, 19 Jul 2019 19:54:35 +0000
Subject: [PATCH] Add CSIDriver creation factory to e2e framework

---
 test/e2e/framework/BUILD     |  1 +
 test/e2e/framework/create.go | 25 +++++++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/test/e2e/framework/BUILD b/test/e2e/framework/BUILD
index cfbfb88ece5..c8e7df37c53 100644
--- a/test/e2e/framework/BUILD
+++ b/test/e2e/framework/BUILD
@@ -57,6 +57,7 @@ go_library(
         "//staging/src/k8s.io/api/policy/v1beta1:go_default_library",
         "//staging/src/k8s.io/api/rbac/v1:go_default_library",
         "//staging/src/k8s.io/api/storage/v1:go_default_library",
+        "//staging/src/k8s.io/api/storage/v1beta1:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/api/errors:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/api/resource:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/apis/meta/v1:go_default_library",
diff --git a/test/e2e/framework/create.go b/test/e2e/framework/create.go
index b700da99c31..bf5362b5aab 100644
--- a/test/e2e/framework/create.go
+++ b/test/e2e/framework/create.go
@@ -27,6 +27,7 @@ import (
 	"k8s.io/api/core/v1"
 	rbacv1 "k8s.io/api/rbac/v1"
 	storagev1 "k8s.io/api/storage/v1"
+	storagev1beta1 "k8s.io/api/storage/v1beta1"
 	apierrs "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -278,6 +279,7 @@ var errorItemNotSupported = errors.New("not supported")
 var factories = map[What]ItemFactory{
 	{"ClusterRole"}:        &clusterRoleFactory{},
 	{"ClusterRoleBinding"}: &clusterRoleBindingFactory{},
+	{"CSIDriver"}:          &csiDriverFactory{},
 	{"DaemonSet"}:          &daemonSetFactory{},
 	{"Role"}:               &roleFactory{},
 	{"RoleBinding"}:        &roleBindingFactory{},
@@ -327,6 +329,8 @@ func (f *Framework) patchItemRecursively(item interface{}) error {
 		f.PatchName(&item.Name)
 	case *storagev1.StorageClass:
 		f.PatchName(&item.Name)
+	case *storagev1beta1.CSIDriver:
+		f.PatchName(&item.Name)
 	case *v1.ServiceAccount:
 		f.PatchNamespace(&item.ObjectMeta.Namespace)
 	case *v1.Secret:
@@ -557,6 +561,27 @@ func (*storageClassFactory) Create(f *Framework, i interface{}) (func() error, e
 	}, nil
 }
 
+type csiDriverFactory struct{}
+
+func (f *csiDriverFactory) New() runtime.Object {
+	return &storagev1beta1.CSIDriver{}
+}
+
+func (*csiDriverFactory) Create(f *Framework, i interface{}) (func() error, error) {
+	item, ok := i.(*storagev1beta1.CSIDriver)
+	if !ok {
+		return nil, errorItemNotSupported
+	}
+
+	client := f.ClientSet.StorageV1beta1().CSIDrivers()
+	if _, err := client.Create(item); err != nil {
+		return nil, errors.Wrap(err, "create CSIDriver")
+	}
+	return func() error {
+		return client.Delete(item.GetName(), &metav1.DeleteOptions{})
+	}, nil
+}
+
 type secretFactory struct{}
 
 func (f *secretFactory) New() runtime.Object {