Skip to content

Commit

Permalink
Improved the unit test for the EnvFromSecret function
Browse files Browse the repository at this point in the history
  • Loading branch information
gemblerz committed Jan 22, 2024
1 parent 00fe77c commit 442f6aa
Show file tree
Hide file tree
Showing 4 changed files with 105 additions and 70 deletions.
4 changes: 2 additions & 2 deletions cmd/runplugin/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ func main() {
os.Exit(1)
}

resourceManager, err := nodescheduler.NewK3SResourceManager(false, kubeconfig, "runplugin", false)
resourceManager, err := nodescheduler.NewK3SResourceManager(false, kubeconfig, "runplugin")
if err != nil {
log.Fatalf("nodescheduler.NewK3SResourceManager: %s", err.Error())
}
Expand Down Expand Up @@ -164,7 +164,7 @@ func runPlugin(resourceManager *nodescheduler.ResourceManager, plugin *datatype.
return nil
}

func updateDeployment(clientset *kubernetes.Clientset, deployment *v1.Deployment) error {
func updateDeployment(clientset kubernetes.Interface, deployment *v1.Deployment) error {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()

Expand Down
99 changes: 57 additions & 42 deletions pkg/nodescheduler/resourcemanager.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,10 @@ import (
"k8s.io/apimachinery/pkg/api/resource"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/labels"
"k8s.io/apimachinery/pkg/runtime"
"k8s.io/apimachinery/pkg/watch"
"k8s.io/client-go/kubernetes"
"k8s.io/client-go/kubernetes/fake"
"k8s.io/client-go/rest"
"k8s.io/client-go/tools/clientcmd"
"k8s.io/client-go/util/homedir"
Expand All @@ -56,7 +58,7 @@ var (
// ResourceManager structs a resource manager talking to a local computing cluster to schedule plugins
type ResourceManager struct {
Namespace string
Clientset *kubernetes.Clientset
Clientset kubernetes.Interface
MetricsClient *metrics.Clientset
RMQManagement *RMQManagement
Notifier *interfacing.Notifier
Expand All @@ -66,16 +68,7 @@ type ResourceManager struct {
}

// NewResourceManager returns an instance of ResourceManager
func NewK3SResourceManager(incluster bool, kubeconfig string, runner string, simulate bool) (rm *ResourceManager, err error) {
if simulate {
return &ResourceManager{
Namespace: namespace,
Clientset: nil,
MetricsClient: nil,
Simulate: simulate,
runner: runner,
}, nil
}
func NewK3SResourceManager(incluster bool, kubeconfig string, runner string) (rm *ResourceManager, err error) {
k3sClient, err := GetK3SClient(incluster, kubeconfig)
if err != nil {
return
Expand All @@ -88,11 +81,21 @@ func NewK3SResourceManager(incluster bool, kubeconfig string, runner string, sim
Namespace: namespace,
Clientset: k3sClient,
MetricsClient: metricsClient,
Simulate: simulate,
runner: runner,
}, nil
}

// NewFakeK3SResourceManager creates a ResourceManager object with the fake Kubernetes Clientset
// which holds given objects.
func NewFakeK3SResourceManager(objects []runtime.Object) (rm *ResourceManager) {
return &ResourceManager{
Namespace: namespace,
Clientset: fake.NewSimpleClientset(objects...),
MetricsClient: nil,
runner: "fake",
}
}

func generatePassword() string {
b := make([]byte, 16)
if _, err := rand.Read(b); err != nil {
Expand Down Expand Up @@ -178,6 +181,39 @@ func securityContextForConfig(pluginSpec *datatype.PluginSpec) *apiv1.SecurityCo
return nil
}

func (rm *ResourceManager) parseEnv(rawEnv map[string]string) (parsedEnv []apiv1.EnvVar, err error) {
for k, v := range rawEnv {
// If the environment variable value refers to a Kubernetes Secret,
// then we check and load the value from the Kubernetes Secret.
// Else, it must be just a value.
if sp := pluginEnvFromSecretRegexFormat.FindStringSubmatch(v); len(sp) == 3 {
secretName := sp[1]
secret, _err := rm.GetSecret(secretName)
if _err != nil {
err = _err
return
}

keyName := sp[2]
if secretV, found := secret.Data[keyName]; found {
parsedEnv = append(parsedEnv, apiv1.EnvVar{
Name: k,
Value: string(secretV),
})
} else {
err = fmt.Errorf("Secret %s does not have the variable %s. Please check.", secretName, keyName)
return
}
} else {
parsedEnv = append(parsedEnv, apiv1.EnvVar{
Name: k,
Value: v,
})
}
}
return
}

func (rm *ResourceManager) ConfigureKubernetes(inCluster bool, kubeconfig string) error {
k3sClient, err := GetK3SClient(inCluster, kubeconfig)
if err != nil {
Expand Down Expand Up @@ -393,7 +429,14 @@ func (rm *ResourceManager) WatchPod(name string, namespace string, retry int) (w
}

func (rm *ResourceManager) createPodTemplateSpecForPlugin(pr *datatype.PluginRuntime) (v1.PodTemplateSpec, error) {
envs := []apiv1.EnvVar{
// We put user environmental variables first, so that they don't
// override our environmental variagbles.
envs, err := rm.parseEnv(pr.Plugin.PluginSpec.Env)
if err != nil {
return v1.PodTemplateSpec{}, err
}

envs = append(envs, []apiv1.EnvVar{
{
Name: "PULSE_SERVER",
Value: "tcp:wes-audio-server.default.svc.cluster.local:4713",
Expand Down Expand Up @@ -474,35 +517,7 @@ func (rm *ResourceManager) createPodTemplateSpecForPlugin(pr *datatype.PluginRun
// },
// },
// },
}

for k, v := range pr.Plugin.PluginSpec.Env {
// If the environment variable value refers to a Kubernetes Secret,
// then we check and load the value from the Kubernetes Secret.
// Else, it must be just a value.
if sp := pluginEnvFromSecretRegexFormat.FindStringSubmatch(v); len(sp) == 3 {
secretName := sp[1]
secret, err := rm.GetSecret(secretName)
if err != nil {
return v1.PodTemplateSpec{}, err
}

keyName := sp[2]
if secretV, found := secret.Data[keyName]; found {
envs = append(envs, apiv1.EnvVar{
Name: k,
Value: string(secretV),
})
} else {
return v1.PodTemplateSpec{}, fmt.Errorf("Secret %s does not have the variable %s. Please check.", secretName, keyName)
}
} else {
envs = append(envs, apiv1.EnvVar{
Name: k,
Value: v,
})
}
}
}...)

tag, err := pr.Plugin.PluginSpec.GetImageTag()
if err != nil {
Expand Down
70 changes: 45 additions & 25 deletions pkg/nodescheduler/resourcemanager_test.go
Original file line number Diff line number Diff line change
@@ -1,38 +1,58 @@
package nodescheduler

import (
"reflect"
"regexp"
"testing"

"github.com/waggle-sensor/edge-scheduler/pkg/datatype"
"gopkg.in/yaml.v2"
corev1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/runtime"
)

func TestEnvFromSecret(t *testing.T) {
pluginEnvFromSecretRegexFormat = regexp.MustCompile(`^{secret\.([a-z0-9-]+).([a-zA-Z0-9]+)}$`)

kv := map[string]string{
"ENV1": "123",
"{secret.mysecret.myenv}": "helloworld",
mySecretWord := "helloworld"
// A secret which we will retrieve the data from.
objects := []runtime.Object{
&corev1.Secret{
ObjectMeta: metav1.ObjectMeta{
Name: "mysecret",
Namespace: "ses", // ResourceManager uses ses namespace by default
},
Data: map[string][]byte{
"greeting": []byte(mySecretWord),
},
},
}
expected := map[string]string{
"ENV1": "123",
"myenv": "helloworld",
fake_rm := NewFakeK3SResourceManager(objects)

pr := datatype.NewPluginRuntimeWithScienceRule(
datatype.Plugin{
Name: "test-plguin",
PluginSpec: &datatype.PluginSpec{
Image: "myimage:latest",
Env: map[string]string{
"ENV1": "123",
"MYENV": "{secret.mysecret.greeting}",
},
},
}, datatype.ScienceRule{},
)
pluginSpec, err := fake_rm.createPodTemplateSpecForPlugin(pr)
if err != nil {
t.Error(err)
}
parsed := make(map[string]string)
for k, v := range kv {
// if the environment variable name refers to a Kubernetes Secret,
// then we check and load the Secret from Kubernetes.
// Else, the it may be just a variable name
// FYI, len(sp) must be 3; [{secret.mysecret.myenv} mysecret myenv]
if sp := pluginEnvFromSecretRegexFormat.FindStringSubmatch(k); len(sp) == 3 {
// secretName := sp[1]
keyName := sp[2]
// fmt.Printf("%v and name is %q, env name is %q", sp, secretName, keyName)
parsed[keyName] = v
} else {
parsed[k] = v
pluginContainer := pluginSpec.Spec.Containers[0]
found := false
for _, e := range pluginContainer.Env {
if e.Name == "MYENV" && e.Value == mySecretWord {
found = true
break
}
}
if !reflect.DeepEqual(parsed, expected) {
t.Errorf("%+v but expected %+v", parsed, expected)
if !found {
t.Errorf("%s was expected in %v, but not found", mySecretWord, pluginContainer.Env)
}
out, _ := yaml.Marshal(pluginSpec)
t.Logf("%s", out)
}
2 changes: 1 addition & 1 deletion pkg/pluginctl/pluginctl.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ type Deployment struct {
}

func NewPluginCtl(kubeconfig string) (*PluginCtl, error) {
resourceManager, err := nodescheduler.NewK3SResourceManager(false, kubeconfig, "pluginctl", false)
resourceManager, err := nodescheduler.NewK3SResourceManager(false, kubeconfig, "pluginctl")
if err != nil {
return nil, err
}
Expand Down

0 comments on commit 442f6aa

Please sign in to comment.