support set k8s by env variable

Issue-ID: MSB-256

Change-Id: Ie4b8cbd8bd1041e4a3fdcc977a0383d2ee4adb01
Signed-off-by: Lvbo163 <lv.bo163@zte.com.cn>

3 files changed, 83 insertions, 2 deletions

diff --git a/msb2pilot/src/msb2pilot/models/config.go b/msb2pilot/src/msb2pilot/models/config.go
index 7f34514..8dd2d3c 100644
--- a/msb2pilot/src/msb2pilot/models/config.go
+++ b/msb2pilot/src/msb2pilot/models/config.go
@@ -13,4 +13,5 @@ package models
 
 const (
 	EnvConsulAddress = "ConsulAddress" //http://localhost:8500
+	EnvK8sAddress    = "K8sAddress"
 )
diff --git a/msb2pilot/src/msb2pilot/pilot/controller.go b/msb2pilot/src/msb2pilot/pilot/controller.go
index d931605..dd45e7f 100644
--- a/msb2pilot/src/msb2pilot/pilot/controller.go
+++ b/msb2pilot/src/msb2pilot/pilot/controller.go
@@ -14,6 +14,10 @@ package pilot
 import (
 	"errors"
 	"msb2pilot/log"
+	"msb2pilot/models"
+	"msb2pilot/util"
+	"os"
+	"path/filepath"
 
 	"istio.io/istio/pilot/pkg/config/kube/crd"
 	"istio.io/istio/pilot/pkg/model"
@@ -22,12 +26,15 @@ import (
 type Operation string
 
 var (
-	client *crd.Client
+	client     *crd.Client
+	configPath = filepath.Join(util.GetCfgPath(), "k8s.yml")
 )
 
 func init() {
+	updateK8sAddress(configPath)
+
 	var err error
-	client, err = crd.NewClient("k8s.yml", model.ConfigDescriptor{
+	client, err = crd.NewClient(configPath, model.ConfigDescriptor{
 		model.RouteRule,
 		model.DestinationPolicy,
 		model.DestinationRule,
@@ -38,6 +45,39 @@ func init() {
 	}
 }
 
+func updateK8sAddress(path string) (string, error) {
+	addr := os.Getenv(models.EnvK8sAddress)
+	log.Log.Informational("k8s cfg address from env: ", addr)
+	if addr == "" {
+		return "", nil
+	}
+
+	// load cfg file
+	cfgstr, err := util.Read(path)
+	if err != nil {
+		log.Log.Error("file to load k8s config file", err)
+		return "", err
+	}
+
+	// update address
+	cfg := make(map[string]interface{})
+	util.UnmarshalYaml(cfgstr, &cfg)
+	if clusters, exist := cfg["clusters"]; exist {
+		clusterItem := clusters.([]interface{})[0]
+		cluster, _ := clusterItem.(map[interface{}]interface{})["cluster"]
+		cluster.(map[interface{}]interface{})["server"] = addr
+	}
+
+	updatedCfgstr, _ := util.MarshalYaml(cfg)
+
+	err = util.Write(path, updatedCfgstr, 0644)
+	if err != nil {
+		log.Log.Error("fail to write k8s cfg info to file", err)
+	}
+
+	return addr, err
+}
+
 func Get(typ, namespace, name string) (*model.Config, bool) {
 	proto, err := protoSchema(typ)
 	if err != nil {
diff --git a/msb2pilot/src/msb2pilot/pilot/controller_test.go b/msb2pilot/src/msb2pilot/pilot/controller_test.go
index 6dcea5b..d7fe7d9 100644
--- a/msb2pilot/src/msb2pilot/pilot/controller_test.go
+++ b/msb2pilot/src/msb2pilot/pilot/controller_test.go
@@ -13,6 +13,9 @@ package pilot
 
 import (
 	"fmt"
+	"msb2pilot/models"
+	"os"
+	"reflect"
 	"testing"
 )
 
@@ -24,3 +27,40 @@ func TestList(t *testing.T) {
 		fmt.Print(res)
 	}
 }
+
+func TestUpdateK8sAddress(t *testing.T) {
+	cases := []struct {
+		path, addr, want, err string
+	}{
+		{
+			path: "k8s.yml222",
+			addr: "filenoteexisttest",
+			want: "",
+			err:  "*os.PathError",
+		},
+		{
+			path: configPath,
+			addr: "",
+			want: "",
+			err:  "",
+		},
+		{
+			path: configPath,
+			addr: "k8stest",
+			want: "k8stest",
+			err:  "",
+		},
+	}
+
+	oldEnv := os.Getenv(models.EnvK8sAddress)
+	for _, cas := range cases {
+		os.Unsetenv(models.EnvK8sAddress)
+		os.Setenv(models.EnvK8sAddress, cas.addr)
+
+		got, err := updateK8sAddress(cas.path)
+		if got != cas.want || (err != nil && reflect.TypeOf(err).String() != cas.err) {
+			t.Errorf("updateK8sAddress(%s, %s) => got %s %v, want %s", cas.path, cas.addr, got, reflect.TypeOf(err), cas.want)
+		}
+	}
+	os.Setenv(models.EnvK8sAddress, oldEnv)
+}