Merge pull request #10 from ncode/juliano/moar_tests

tests for updateServiceTags

README.md

@@ -53,5 +53,4 @@ sequenceDiagram
 
 ## Todo
 
-- [ ] Adds support for multiple services (currently only supports one service)
 - [ ] Adds a systemd unit file generator

cmd/run.go

@@ -16,6 +16,7 @@ limitations under the License.
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"time"
@@ -39,6 +40,7 @@ example: tagit run -s my-super-service -x '/tmp/tag-role.sh'
 		script := cmd.PersistentFlags().Lookup("script").Value.String()
 		tagPrefix := cmd.PersistentFlags().Lookup("tag-prefix").Value.String()
 		interval := cmd.PersistentFlags().Lookup("interval").Value.String()
+		ctx := context.Background()
 
 		if serviceID == "" {
 			fmt.Println("service-id is required")
@@ -71,7 +73,7 @@ example: tagit run -s my-super-service -x '/tmp/tag-role.sh'
 		}
 
 		t := tagit.New(tagit.NewConsulAPIWrapper(consulClient), &tagit.CmdExecutor{}, serviceID, script, validInterval, tagPrefix)
-		t.Run()
+		t.Run(ctx)
 	},
 }
 

pkg/tagit/tagit.go

@@ -1,6 +1,7 @@
 package tagit
 
 import (
+	"context"
 	"fmt"
 	"os/exec"
 	"strings"
@@ -77,16 +78,22 @@ func New(consulClient ConsulClient, commandExecutor CommandExecutor, serviceID s
 }
 
 // Run will run the tagit flow and tag consul services based on the script output
-func (t *TagIt) Run() {
+func (t *TagIt) Run(ctx context.Context) {
+	ticker := time.NewTicker(t.Interval)
 	for {
-		err := t.updateServiceTags()
-		if err != nil {
-			log.WithFields(log.Fields{
-				"service": t.ServiceID,
-				"error":   err,
-			}).Error("error updating service tags")
+		select {
+		case <-ctx.Done():
+			ticker.Stop()
+			return
+		case <-ticker.C:
+			err := t.updateServiceTags()
+			if err != nil {
+				log.WithFields(log.Fields{
+					"service": t.ServiceID,
+					"error":   err,
+				}).Error("error updating service tags")
+			}
 		}
-		time.Sleep(t.Interval)
 	}
 }
 

pkg/tagit/tagit_test.go

@@ -562,3 +562,64 @@ func TestCmdExecutor_Execute(t *testing.T) {
 		})
 	}
 }
+
+func TestUpdateServiceTags(t *testing.T) {
+	tests := []struct {
+		name             string
+		mockScriptOutput string
+		mockScriptError  error
+		existingTags     []string
+		newTags          []string
+		mockRegisterErr  error
+		expectError      bool
+	}{
+		{
+			name:             "Successful Update",
+			mockScriptOutput: "new-tag1 new-tag2",
+			existingTags:     []string{"old-tag"},
+			newTags:          []string{"tag-new-tag1", "tag-new-tag2"},
+			expectError:      false,
+		},
+		{
+			name:            "Script Error",
+			mockScriptError: fmt.Errorf("script error"),
+			expectError:     true,
+		},
+		{
+			name:             "Consul Register Error",
+			mockScriptOutput: "new-tag1 new-tag2",
+			mockRegisterErr:  fmt.Errorf("consul error"),
+			expectError:      true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockExecutor := &MockCommandExecutor{
+				MockOutput: []byte(tt.mockScriptOutput),
+				MockError:  tt.mockScriptError,
+			}
+			mockConsulClient := &MockConsulClient{
+				MockAgent: &MockAgent{
+					ServicesFunc: func() (map[string]*api.AgentService, error) {
+						return map[string]*api.AgentService{
+							"test-service": {
+								ID:   "test-service",
+								Tags: tt.existingTags,
+							},
+						}, nil
+					},
+					ServiceRegisterFunc: func(reg *api.AgentServiceRegistration) error {
+						return tt.mockRegisterErr
+					},
+				},
+			}
+			tagit := New(mockConsulClient, mockExecutor, "test-service", "echo test", 30*time.Second, "tag")
+
+			err := tagit.updateServiceTags()
+			if (err != nil) != tt.expectError {
+				t.Errorf("updateServiceTags() error = %v, wantErr %v", err, tt.expectError)
+			}
+		})
+	}
+}