github.com/jingruilea/kubeedge@v1.2.0-beta.0.0.20200410162146-4bb8902b3879/edge/test/cloudhub/stub.go (about)

     1  package test
     2  
     3  import (
     4  	"encoding/json"
     5  	"io"
     6  	"io/ioutil"
     7  	"net/http"
     8  
     9  	"github.com/gorilla/mux"
    10  	"github.com/gorilla/websocket"
    11  	"k8s.io/api/core/v1"
    12  	"k8s.io/klog"
    13  
    14  	"github.com/kubeedge/beehive/pkg/core"
    15  	beehiveContext "github.com/kubeedge/beehive/pkg/core/context"
    16  	"github.com/kubeedge/beehive/pkg/core/model"
    17  	"github.com/kubeedge/kubeedge/edge/pkg/common/modules"
    18  )
    19  
    20  func init() {
    21  	core.Register(&stubCloudHub{enable: true})
    22  }
    23  
    24  type Attributes struct {
    25  	RoleName  string `json:"iam_role"`
    26  	ProjectID string `json:"project_id"`
    27  }
    28  
    29  type record struct {
    30  	Data         string `json:"data"`
    31  	PartitionKey string `json:"partition_key"`
    32  }
    33  
    34  type stubCloudHub struct {
    35  	wsConn *websocket.Conn
    36  	enable bool
    37  }
    38  
    39  func (*stubCloudHub) Name() string {
    40  	return "stubCloudHub"
    41  }
    42  
    43  func (*stubCloudHub) Group() string {
    44  	//return core.MetaGroup
    45  	return modules.MetaGroup
    46  }
    47  
    48  func (tm *stubCloudHub) Enable() bool {
    49  	return tm.enable
    50  }
    51  
    52  func (tm *stubCloudHub) eventReadLoop(conn *websocket.Conn, stop chan bool) {
    53  	for {
    54  		var event interface{}
    55  		err := conn.ReadJSON(&event)
    56  		if err != nil {
    57  			klog.Errorf("read error, connection will be closed: %v", err)
    58  			stop <- true
    59  			return
    60  		}
    61  		klog.Infof("cloud hub receive message %+v", event)
    62  	}
    63  }
    64  
    65  func (tm *stubCloudHub) serveEvent(w http.ResponseWriter, r *http.Request) {
    66  	upgrader := websocket.Upgrader{}
    67  	conn, err := upgrader.Upgrade(w, r, nil)
    68  	if err != nil {
    69  		klog.Errorf("fail to build websocket connection: %v", err)
    70  		http.Error(w, "fail to upgrade to websocket protocol", http.StatusInternalServerError)
    71  		return
    72  	}
    73  	tm.wsConn = conn
    74  	stop := make(chan bool, 1)
    75  	klog.Info("edge connected")
    76  	go tm.eventReadLoop(conn, stop)
    77  	<-stop
    78  	tm.wsConn = nil
    79  	klog.Info("edge disconnected")
    80  }
    81  
    82  func (tm *stubCloudHub) podHandler(w http.ResponseWriter, req *http.Request) {
    83  	if req.Body != nil {
    84  		body, err := ioutil.ReadAll(req.Body)
    85  		if err != nil {
    86  			klog.Errorf("read body error %v", err)
    87  			w.Write([]byte("read request body error"))
    88  			return
    89  		}
    90  		klog.Infof("request body is %s\n", string(body))
    91  
    92  		var pod v1.Pod
    93  		if err = json.Unmarshal(body, &pod); err != nil {
    94  			klog.Errorf("unmarshal request body error %v", err)
    95  			w.Write([]byte("unmarshal request body error"))
    96  			return
    97  		}
    98  		var msgReq *model.Message
    99  		switch req.Method {
   100  		case "POST":
   101  			msgReq = model.NewMessage("").BuildRouter("edgecontroller", "resource",
   102  				"node/fake_node_id/pod/"+string(pod.UID), model.InsertOperation).FillBody(pod)
   103  		case "DELETE":
   104  			msgReq = model.NewMessage("").BuildRouter("edgecontroller", "resource",
   105  				"node/fake_node_id/pod/"+string(pod.UID), model.DeleteOperation).FillBody(pod)
   106  		}
   107  
   108  		if tm.wsConn != nil {
   109  			tm.wsConn.WriteJSON(*msgReq)
   110  			klog.Infof("send message to edgehub is %+v\n", *msgReq)
   111  		}
   112  
   113  		io.WriteString(w, "OK\n")
   114  	}
   115  }
   116  
   117  func (tm *stubCloudHub) Start() {
   118  	defer tm.Cleanup()
   119  
   120  	router := mux.NewRouter()
   121  	router.HandleFunc("/{group_id}/events", tm.serveEvent) // for edge-hub
   122  	router.HandleFunc("/pod", tm.podHandler)               // for pod test
   123  
   124  	s := http.Server{
   125  		Addr:    "127.0.0.1:20000",
   126  		Handler: router,
   127  	}
   128  	klog.Info("Start cloud hub service")
   129  	err := s.ListenAndServe()
   130  	if err != nil {
   131  		klog.Errorf("ListenAndServe: %v", err)
   132  	}
   133  
   134  }
   135  
   136  func (tm *stubCloudHub) Cleanup() {
   137  	beehiveContext.Cleanup(tm.Name())
   138  }