Merge pull request #58 from andywaltlova/refactor-server

Refactor server.go file and add tests

src/runner.go

@@ -71,8 +71,8 @@ func verifyYamlFile(yamlData []byte) bool {
 // If signature is valid then extracts the bash script to temporary file,
 // sets env variables if present and then runs the script.
 // Return stdout of executed script or error message if the signature wasn't valid.
-func processSignedScript(yamlFileContet []byte) string {
-	signatureIsValid := verifyYamlFile(yamlFileContet)
+func processSignedScript(incomingContent []byte) string {
+	signatureIsValid := verifyYamlFile(incomingContent)
 	if !signatureIsValid {
 		errorMsg := "Signature of yaml file is invalid"
 		log.Errorln(errorMsg)
@@ -82,9 +82,10 @@ func processSignedScript(yamlFileContet []byte) string {
 
 	// Parse the YAML data into the yamlConfig struct
 	var yamlContent signedYamlFile
-	err := yaml.Unmarshal(yamlFileContet, &yamlContent)
+	err := yaml.Unmarshal(incomingContent, &yamlContent)
 	if err != nil {
 		log.Errorln(err)
+		return "Yaml couldn't be unmarshaled"
 	}
 
 	// Set env variables
@@ -123,8 +124,8 @@ func processSignedScript(yamlFileContet []byte) string {
 
 	// Execute the script
 	log.Infoln("Executing bash script")
-
 	out, err := exec.Command("/bin/sh", scriptFileName).Output()
+
 	if err != nil {
 		log.Errorln("Failed to execute script: ", err)
 		return ""

src/server.go

@@ -40,6 +40,40 @@ func createDataMessage(commandOutput string, metadata map[string]string, directi
 	return data
 }
 
+// Processes signed script and sends message back to dispatcher
+func processData(d *pb.Data) *pb.Data {
+	log.Infoln("Processing received yaml data")
+	commandOutput := processSignedScript(d.GetContent())
+
+	// Create a data message to send back to the dispatcher.
+	log.Infof("Creating payload for message %s", d.GetMessageId())
+	data := createDataMessage(commandOutput, d.GetMetadata(), d.GetDirective(), d.GetMessageId())
+	return data
+}
+
+// Sends data back to dispatcher
+func sendDataToDispatcher(data *pb.Data) *pb.Data {
+	// Dial the Dispatcher and call "Finish"
+	conn, err := grpc.Dial(yggdDispatchSocketAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		log.Error(err)
+	}
+	defer conn.Close()
+
+	// Create a client of the Dispatch service
+	client := pb.NewDispatcherClient(conn)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	// Call "Send"
+	log.Infof("Sending response message to %s", data.GetResponseTo())
+	log.Infoln("pb.Data message: ", data)
+	if _, err := client.Send(ctx, data); err != nil {
+		log.Error(err)
+	}
+	return data
+}
+
 // jobServer implements the Worker gRPC service as defined by the yggdrasil
 // gRPC protocol. It accepts Assignment messages, unmarshals the data into a
 // string, and echoes the content back to the Dispatch service by calling the
@@ -59,35 +93,12 @@ type jobServer struct {
 //  4. Creates a client of the Dispatcher service.
 //  5. Constructs a data message to send back to the dispatcher.
 //  6. Sends the data message using the "Send" method of the Dispatcher service.
-func (s *jobServer) Send(ctx context.Context, d *pb.Data) (*pb.Receipt, error) {
-	go func() {
-		log.Infoln("Processing received yaml data")
-		commandOutput := processSignedScript(d.GetContent())
-
-		// Dial the Dispatcher and call "Finish"
-		conn, err := grpc.Dial(
-			yggdDispatchSocketAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
-		if err != nil {
-			log.Error(err)
-		}
-		defer conn.Close()
+func (s *jobServer) Send(_ context.Context, d *pb.Data) (*pb.Receipt, error) {
 
-		// Create a client of the Dispatch service
-		c := pb.NewDispatcherClient(conn)
-		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-		defer cancel()
-
-		// Create a data message to send back to the dispatcher.
-		log.Infof("Creating payload for message %s", d.GetMessageId())
-		data := createDataMessage(
-			commandOutput, d.GetMetadata(), d.GetDirective(), d.GetMessageId())
-
-		// Call "Send"
-		log.Infof("Sending message to %s", d.GetMessageId())
-		log.Infoln("pb.Data message: ", data)
-		if _, err := c.Send(ctx, data); err != nil {
-			log.Error(err)
-		}
+	// Goroutine processing the data, cancels the context when processing is done
+	go func() {
+		data := processData(d)
+		sendDataToDispatcher(data)
 	}()
 
 	// Respond to the start request that the work was accepted.

src/server_test.go

@@ -3,6 +3,9 @@ package main
 import (
 	"strings"
 	"testing"
+
+	"github.com/google/uuid"
+	pb "github.com/redhatinsights/yggdrasil/protocol"
 )
 
 func TestCreateDataMessage(t *testing.T) {
@@ -47,3 +50,68 @@ func TestCreateDataMessage(t *testing.T) {
 		t.Errorf("Expected Content to be empty, but got %s", data.Content)
 	}
 }
+
+func TestProcessData(t *testing.T) {
+	// FIXME: this should ideally test that all correct functions are called
+	// Probably easiest would be to move them to interface and then make the function argument take mock interface
+	yggdDispatchSocketAddr = "mock-target"
+
+	shouldVerifyYaml := false
+	shouldDoInsightsCoreGPGCheck := false
+	temporaryWorkerDirectory := "test-dir"
+	config = &Config{
+		VerifyYAML:               &shouldVerifyYaml,
+		TemporaryWorkerDirectory: &temporaryWorkerDirectory,
+		InsightsCoreGPGCheck:     &shouldDoInsightsCoreGPGCheck,
+	}
+
+	yamlData := []byte(`
+vars:
+    _insights_signature: "invalid-signature"
+    _insights_signature_exclude: "/vars/insights_signature,/vars/content_vars"
+    content: |
+        #!/bin/sh
+        echo "$RHC_WORKER_FOO $RHC_WORKER_BAR!"
+    content_vars:
+        FOO: Hello
+        BAR: World`)
+
+	returnURL := "bar"
+	testData := &pb.Data{
+		Content: yamlData,
+		Metadata: map[string]string{
+			"return_content_type": "foo",
+			"return_url":          returnURL,
+			"correlation_id":      "000",
+		},
+		Directive: "Your directive",
+		MessageId: "Your message ID",
+	}
+
+	data := processData(testData)
+	expectedOutput := "Hello World!"
+
+	if !strings.Contains(string(data.GetContent()), expectedOutput) {
+		t.Errorf("Expected content to contain '%s', but it didn't", expectedOutput)
+	}
+
+	if data.GetDirective() != returnURL {
+		t.Errorf("Expected directive to contain '%s', but it didn't", returnURL)
+	}
+}
+
+func TestSendDataToDispatcher(t *testing.T) {
+	// Tests only that the function doesn't modify data sent to dispatcher
+
+	yggdDispatchSocketAddr = "mock-target"
+
+	testData := &pb.Data{
+		MessageId:  uuid.New().String(),
+		ResponseTo: "mock-id",
+	}
+
+	data := sendDataToDispatcher(testData)
+	if data != testData {
+		t.Errorf("Function should NOT change data before sent, but it did: %s", data)
+	}
+}