From 9f0042a3bcda124473ad715921406d74d75b2b9d Mon Sep 17 00:00:00 2001
From: Clive Cox <cc@seldon.io>
Date: Mon, 14 Jun 2021 11:37:39 +0100
Subject: [PATCH] Allow MLServer envs to be overriden

---
 operator/controllers/mlserver.go      | 11 +++++++++--
 operator/controllers/mlserver_test.go |  3 +++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/operator/controllers/mlserver.go b/operator/controllers/mlserver.go
index 5a21a8097d..4689c04059 100644
--- a/operator/controllers/mlserver.go
+++ b/operator/controllers/mlserver.go
@@ -3,6 +3,7 @@ package controllers
 import (
 	"errors"
 	"fmt"
+	"github.com/seldonio/seldon-core/operator/utils"
 	"strconv"
 
 	machinelearningv1 "github.com/seldonio/seldon-core/operator/apis/machinelearning.seldon.io/v1"
@@ -41,8 +42,14 @@ func mergeMLServerContainer(existing *v1.Container, mlServer *v1.Container) *v1.
 		existing.Env = []v1.EnvVar{}
 	}
 
-	// TODO: Allow overriding some of the env vars
-	existing.Env = append(existing.Env, mlServer.Env...)
+	for _, envVar := range existing.Env {
+		if utils.HasEnvVar(mlServer.Env, envVar.Name) {
+			mlServer.Env = utils.SetEnvVar(mlServer.Env, envVar)
+		} else {
+			mlServer.Env = append(mlServer.Env, envVar)
+		}
+	}
+	existing.Env = mlServer.Env
 
 	if existing.ReadinessProbe == nil {
 		existing.ReadinessProbe = mlServer.ReadinessProbe
diff --git a/operator/controllers/mlserver_test.go b/operator/controllers/mlserver_test.go
index e4ef596aa7..affeaf769c 100644
--- a/operator/controllers/mlserver_test.go
+++ b/operator/controllers/mlserver_test.go
@@ -33,11 +33,13 @@ var _ = Describe("MLServer helpers", func() {
 	Describe("mergeMLServerContainer", func() {
 		var existing *v1.Container
 		var mlServer *v1.Container
+		customEnvValue := "{\"custom\":1}"
 
 		BeforeEach(func() {
 			existing = &v1.Container{
 				Env: []v1.EnvVar{
 					{Name: "FOO", Value: "BAR"},
+					{Name: MLServerTempoRuntimeEnv, Value: customEnvValue},
 				},
 			}
 
@@ -49,6 +51,7 @@ var _ = Describe("MLServer helpers", func() {
 
 			Expect(merged).ToNot(BeNil())
 			Expect(merged.Env).To(ContainElement(v1.EnvVar{Name: "FOO", Value: "BAR"}))
+			Expect(merged.Env).To(ContainElement(v1.EnvVar{Name: MLServerTempoRuntimeEnv, Value: customEnvValue}))
 			Expect(merged.Env).To(ContainElements(mlServer.Env))
 			Expect(merged.Image).To(Equal(mlServer.Image))
 		})