From 5e9d49e1391a0905e38ab0567deadda66e8261fb Mon Sep 17 00:00:00 2001
From: WYike <77846369+WangYikewyk@users.noreply.github.com>
Date: Sun, 7 Feb 2021 19:31:42 +0800
Subject: [PATCH] fix vela traits not compatible with installed capability
 (#1033)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix issue 1029

* add test case

* fix check diff

Co-authored-by: 王易可 <wangyike@B-V1QHML7L-1909.local>
---
 pkg/serverlib/capability.go      | 24 +++++++++++++++++++++
 pkg/serverlib/capability_test.go | 37 ++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 pkg/serverlib/capability_test.go

diff --git a/pkg/serverlib/capability.go b/pkg/serverlib/capability.go
index cff295b75..2f75976e6 100644
--- a/pkg/serverlib/capability.go
+++ b/pkg/serverlib/capability.go
@@ -2,6 +2,7 @@ package serverlib
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -114,6 +115,10 @@ func InstallCapability(client client.Client, mapper discoverymapper.DiscoveryMap
 			APIVersion: gvk.GroupVersion().String(),
 			Kind:       gvk.Kind,
 		}
+		err = addSourceIntoExtension(wd.Spec.Extension, tp.Source)
+		if err != nil {
+			return err
+		}
 		if err = client.Create(context.Background(), &wd); err != nil && !apierrors.IsAlreadyExists(err) {
 			return err
 		}
@@ -145,6 +150,10 @@ func InstallCapability(client client.Client, mapper discoverymapper.DiscoveryMap
 			APIVersion: gvk.GroupVersion().String(),
 			Kind:       gvk.Kind,
 		}
+		err = addSourceIntoExtension(td.Spec.Extension, tp.Source)
+		if err != nil {
+			return err
+		}
 		if err = client.Create(context.Background(), &td); err != nil && !apierrors.IsAlreadyExists(err) {
 			return err
 		}
@@ -421,3 +430,18 @@ func checkInstallStatus(repoName string, tmp types.Capability) string {
 	}
 	return status
 }
+
+func addSourceIntoExtension(in *runtime.RawExtension, source *types.Source) error {
+	var extension map[string]interface{}
+	err := json.Unmarshal(in.Raw, &extension)
+	if err != nil {
+		return err
+	}
+	extension["source"] = source
+	data, err := json.Marshal(extension)
+	if err != nil {
+		return err
+	}
+	in.Raw = data
+	return nil
+}
diff --git a/pkg/serverlib/capability_test.go b/pkg/serverlib/capability_test.go
new file mode 100644
index 000000000..22080341e
--- /dev/null
+++ b/pkg/serverlib/capability_test.go
@@ -0,0 +1,37 @@
+package serverlib
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+
+	"k8s.io/apimachinery/pkg/runtime"
+
+	"github.com/oam-dev/kubevela/apis/types"
+)
+
+func TestAddSourceIntoDefinition(t *testing.T) {
+	caseJson := []byte(`{"template":""}`)
+	wantJson := []byte(`{"source":{"repoName":"foo"},"template":""}`)
+	source := types.Source{RepoName: "foo"}
+	testcase := runtime.RawExtension{Raw: caseJson}
+	err := addSourceIntoExtension(&testcase, &source)
+	if err != nil {
+		t.Error("meet an error ", err)
+		return
+	}
+	var result, want map[string]interface{}
+	err = json.Unmarshal(testcase.Raw, &result)
+	if err != nil {
+		t.Error("marshaling object meet an error ", err)
+		return
+	}
+	err = json.Unmarshal(wantJson, &want)
+	if err != nil {
+		t.Error("marshaling object meet an error ", err)
+		return
+	}
+	if !reflect.DeepEqual(result, want) {
+		t.Errorf("error result want %s, got %s", result, testcase)
+	}
+}