From a5fe983735c6f396b8f618e1c5c3ae368497966b Mon Sep 17 00:00:00 2001
From: Catalin Toda <ctoda@lyft.com>
Date: Thu, 1 Oct 2020 14:12:14 -0700
Subject: [PATCH] Add GetProjectLabels() to TaskExecutionMetadata

---
 .../pluginmachinery/core/exec_metadata.go     |  1 +
 .../core/mocks/task_execution_metadata.go     | 34 ++++++++++
 .../k8s/mocks/plugin_context.go               | 68 +++++++++++++++++++
 3 files changed, 103 insertions(+)

diff --git a/go/tasks/pluginmachinery/core/exec_metadata.go b/go/tasks/pluginmachinery/core/exec_metadata.go
index 1f697b2e1..044b03b81 100644
--- a/go/tasks/pluginmachinery/core/exec_metadata.go
+++ b/go/tasks/pluginmachinery/core/exec_metadata.go
@@ -33,4 +33,5 @@ type TaskExecutionMetadata interface {
 	GetAnnotations() map[string]string
 	GetK8sServiceAccount() string
 	IsInterruptible() bool
+        GetProjectLabels() map[string]string
 }
diff --git a/go/tasks/pluginmachinery/core/mocks/task_execution_metadata.go b/go/tasks/pluginmachinery/core/mocks/task_execution_metadata.go
index be337961c..0365dda9e 100644
--- a/go/tasks/pluginmachinery/core/mocks/task_execution_metadata.go
+++ b/go/tasks/pluginmachinery/core/mocks/task_execution_metadata.go
@@ -246,6 +246,40 @@ func (_m *TaskExecutionMetadata) GetOwnerReference() v1.OwnerReference {
 	return r0
 }
 
+type TaskExecutionMetadata_GetProjectLabels struct {
+	*mock.Call
+}
+
+func (_m TaskExecutionMetadata_GetProjectLabels) Return(_a0 map[string]string) *TaskExecutionMetadata_GetProjectLabels {
+	return &TaskExecutionMetadata_GetProjectLabels{Call: _m.Call.Return(_a0)}
+}
+
+func (_m *TaskExecutionMetadata) OnGetProjectLabels() *TaskExecutionMetadata_GetProjectLabels {
+	c := _m.On("GetProjectLabels")
+	return &TaskExecutionMetadata_GetProjectLabels{Call: c}
+}
+
+func (_m *TaskExecutionMetadata) OnGetProjectLabelsMatch(matchers ...interface{}) *TaskExecutionMetadata_GetProjectLabels {
+	c := _m.On("GetProjectLabels", matchers...)
+	return &TaskExecutionMetadata_GetProjectLabels{Call: c}
+}
+
+// GetProjectLabels provides a mock function with given fields:
+func (_m *TaskExecutionMetadata) GetProjectLabels() map[string]string {
+	ret := _m.Called()
+
+	var r0 map[string]string
+	if rf, ok := ret.Get(0).(func() map[string]string); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(map[string]string)
+		}
+	}
+
+	return r0
+}
+
 type TaskExecutionMetadata_GetTaskExecutionID struct {
 	*mock.Call
 }
diff --git a/go/tasks/pluginmachinery/k8s/mocks/plugin_context.go b/go/tasks/pluginmachinery/k8s/mocks/plugin_context.go
index a9011ddb0..abb8b3f84 100644
--- a/go/tasks/pluginmachinery/k8s/mocks/plugin_context.go
+++ b/go/tasks/pluginmachinery/k8s/mocks/plugin_context.go
@@ -7,6 +7,8 @@ import (
 	io "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io"
 
 	mock "github.com/stretchr/testify/mock"
+
+	storage "github.com/lyft/flytestdlib/storage"
 )
 
 // PluginContext is an autogenerated mock type for the PluginContext type
@@ -14,6 +16,40 @@ type PluginContext struct {
 	mock.Mock
 }
 
+type PluginContext_DataStore struct {
+	*mock.Call
+}
+
+func (_m PluginContext_DataStore) Return(_a0 *storage.DataStore) *PluginContext_DataStore {
+	return &PluginContext_DataStore{Call: _m.Call.Return(_a0)}
+}
+
+func (_m *PluginContext) OnDataStore() *PluginContext_DataStore {
+	c := _m.On("DataStore")
+	return &PluginContext_DataStore{Call: c}
+}
+
+func (_m *PluginContext) OnDataStoreMatch(matchers ...interface{}) *PluginContext_DataStore {
+	c := _m.On("DataStore", matchers...)
+	return &PluginContext_DataStore{Call: c}
+}
+
+// DataStore provides a mock function with given fields:
+func (_m *PluginContext) DataStore() *storage.DataStore {
+	ret := _m.Called()
+
+	var r0 *storage.DataStore
+	if rf, ok := ret.Get(0).(func() *storage.DataStore); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*storage.DataStore)
+		}
+	}
+
+	return r0
+}
+
 type PluginContext_InputReader struct {
 	*mock.Call
 }
@@ -48,6 +84,38 @@ func (_m *PluginContext) InputReader() io.InputReader {
 	return r0
 }
 
+type PluginContext_MaxDatasetSizeBytes struct {
+	*mock.Call
+}
+
+func (_m PluginContext_MaxDatasetSizeBytes) Return(_a0 int64) *PluginContext_MaxDatasetSizeBytes {
+	return &PluginContext_MaxDatasetSizeBytes{Call: _m.Call.Return(_a0)}
+}
+
+func (_m *PluginContext) OnMaxDatasetSizeBytes() *PluginContext_MaxDatasetSizeBytes {
+	c := _m.On("MaxDatasetSizeBytes")
+	return &PluginContext_MaxDatasetSizeBytes{Call: c}
+}
+
+func (_m *PluginContext) OnMaxDatasetSizeBytesMatch(matchers ...interface{}) *PluginContext_MaxDatasetSizeBytes {
+	c := _m.On("MaxDatasetSizeBytes", matchers...)
+	return &PluginContext_MaxDatasetSizeBytes{Call: c}
+}
+
+// MaxDatasetSizeBytes provides a mock function with given fields:
+func (_m *PluginContext) MaxDatasetSizeBytes() int64 {
+	ret := _m.Called()
+
+	var r0 int64
+	if rf, ok := ret.Get(0).(func() int64); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(int64)
+	}
+
+	return r0
+}
+
 type PluginContext_OutputWriter struct {
 	*mock.Call
 }