Merge pull request #12 from ncode/juliano/systemd_and_cleanup

updates

Author: ncode

cmd/cleanup.go

@@ -0,0 +1,58 @@
+/*
+Copyright © 2024 Juliano Martinez <[email protected]>
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package cmd
+
+import (
+	"fmt"
+	"github.com/hashicorp/consul/api"
+	"github.com/ncode/tagit/pkg/tagit"
+	"github.com/spf13/cobra"
+	"os"
+)
+
+// cleanupCmd represents the cleanup command
+var cleanupCmd = &cobra.Command{
+	Use:   "cleanup",
+	Short: "cleanup removes all services with the tag prefix from a given consul service",
+	Run: func(cmd *cobra.Command, args []string) {
+		config := api.DefaultConfig()
+		config.Address = cmd.InheritedFlags().Lookup("consul-addr").Value.String()
+		config.Token = cmd.InheritedFlags().Lookup("token").Value.String()
+		consulClient, err := api.NewClient(config)
+		if err != nil {
+			fmt.Printf("error creating consul client: %s", err.Error())
+			os.Exit(1)
+		}
+
+		t := tagit.New(
+			tagit.NewConsulAPIWrapper(consulClient),
+			&tagit.CmdExecutor{},
+			cmd.InheritedFlags().Lookup("service-id").Value.String(),
+			cmd.InheritedFlags().Lookup("script").Value.String(),
+			0,
+			cmd.InheritedFlags().Lookup("tag-prefix").Value.String(),
+		)
+		err = t.CleanupTags()
+		if err != nil {
+			fmt.Printf("error cleaning up tags: %s", err.Error())
+			os.Exit(1)
+		}
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(cleanupCmd)
+}

cmd/systemd.go

@@ -24,13 +24,7 @@ import (
 // systemdCmd represents the systemd command
 var systemdCmd = &cobra.Command{
 	Use:   "systemd",
-	Short: "A brief description of your command",
-	Long: `A longer description that spans multiple lines and likely contains examples
-and usage of using your command. For example:
-
-Cobra is a CLI library for Go that empowers applications.
-This application is a tool to generate the needed files
-to quickly create a Cobra application.`,
+	Short: "systemd generate a systemd service, that you can use for the tagit service",
 	Run: func(cmd *cobra.Command, args []string) {
 		fields := &systemd.Fields{
 			User:       cmd.PersistentFlags().Lookup("user").Value.String(),

pkg/tagit/tagit.go

@@ -97,6 +97,21 @@ func (t *TagIt) Run(ctx context.Context) {
 	}
 }
 
+// CleanupTags removes all tags with the given prefix from the service.
+func (t *TagIt) CleanupTags() error {
+	service, err := t.getService()
+	if err != nil {
+		return fmt.Errorf("error getting service: %w", err)
+	}
+	registration := t.copyServiceToRegistration(service)
+	updatedTags, tagged := t.excludeTagged(registration.Tags)
+	if tagged {
+		registration.Tags = updatedTags
+		return t.client.Agent().ServiceRegister(registration)
+	}
+	return nil
+}
+
 // runScript runs a command and returns the output.
 func (t *TagIt) runScript() ([]byte, error) {
 	log.WithFields(log.Fields{

pkg/tagit/tagit_test.go

@@ -1,10 +1,12 @@
 package tagit
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"slices"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -623,3 +625,150 @@ func TestUpdateServiceTags(t *testing.T) {
 		})
 	}
 }
+
+func TestCleanupTags(t *testing.T) {
+	tests := []struct {
+		name            string
+		serviceID       string
+		mockServices    map[string]*api.AgentService
+		tagPrefix       string
+		mockRegisterErr error
+		expectError     bool
+		expectTags      []string
+	}{
+		{
+			name:      "Successful Tag Cleanup",
+			serviceID: "test-service",
+			mockServices: map[string]*api.AgentService{
+				"test-service": {
+					ID:   "test-service",
+					Tags: []string{"tag-prefix1", "tag-prefix2", "other-tag"},
+				},
+			},
+			tagPrefix:   "tag",
+			expectError: false,
+			expectTags:  []string{"other-tag"},
+		},
+		{
+			name:      "No Tag Cleanup needed",
+			serviceID: "test-service",
+			mockServices: map[string]*api.AgentService{
+				"test-service": {
+					ID:   "test-service",
+					Tags: []string{"prefix1", "prefix2", "other-tag"},
+				},
+			},
+			tagPrefix:   "tag",
+			expectError: false,
+			expectTags:  []string{"prefix1", "prefix2", "other-tag"},
+		},
+		{
+			name:      "Service Not Found",
+			serviceID: "non-existent-service",
+			mockServices: map[string]*api.AgentService{
+				"other-service": {
+					ID:   "other-service",
+					Tags: []string{"some-tag", "another-tag"},
+				},
+			},
+			tagPrefix:   "tag-prefix",
+			expectError: true,
+		},
+		{
+			name:      "Consul Register Error",
+			serviceID: "test-service",
+			mockServices: map[string]*api.AgentService{
+				"test-service": {
+					ID:   "test-service",
+					Tags: []string{"tag-prefix1", "other-tag"},
+				},
+			},
+			tagPrefix:       "tag",
+			mockRegisterErr: fmt.Errorf("consul register error"),
+			expectError:     true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockConsulClient := &MockConsulClient{
+				MockAgent: &MockAgent{
+					ServicesFunc: func() (map[string]*api.AgentService, error) {
+						return tt.mockServices, nil
+					},
+					ServiceRegisterFunc: func(reg *api.AgentServiceRegistration) error {
+						// Ensure the service exists in the mock data
+						if service, exists := tt.mockServices[reg.ID]; exists && tt.mockRegisterErr == nil {
+							// Update the tags of the service
+							service.Tags = reg.Tags
+							tt.mockServices[reg.ID] = service // Update the map with the modified service
+						}
+						return tt.mockRegisterErr
+					},
+				},
+			}
+			tagit := New(mockConsulClient, nil, tt.serviceID, "", 0, tt.tagPrefix)
+
+			err := tagit.CleanupTags()
+			if (err != nil) != tt.expectError {
+				t.Errorf("CleanupTags() error = %v, wantErr %v", err, tt.expectError)
+			}
+
+			if !tt.expectError {
+				updatedService := tt.mockServices[tt.serviceID]
+				if updatedService != nil && !reflect.DeepEqual(updatedService.Tags, tt.expectTags) {
+					t.Errorf("Expected tags after cleanup: %v, got: %v", tt.expectTags, updatedService.Tags)
+				}
+			}
+		})
+	}
+}
+
+func TestRun(t *testing.T) {
+	// Setup
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	updateServiceTagsCalled := atomic.Int32{}
+	mockExecutor := &MockCommandExecutor{
+		MockOutput: []byte("new-tag1 new-tag2"),
+		MockError:  nil,
+	}
+	mockConsulClient := &MockConsulClient{
+		MockAgent: &MockAgent{
+			ServicesFunc: func() (map[string]*api.AgentService, error) {
+				updateServiceTagsCalled.Add(1)
+				if updateServiceTagsCalled.Load() == 2 {
+					return nil, fmt.Errorf("enter error")
+				}
+				return map[string]*api.AgentService{
+					"test-service": {
+						ID:   "test-service",
+						Tags: []string{"old-tag"},
+					},
+				}, nil
+			},
+			ServiceRegisterFunc: func(reg *api.AgentServiceRegistration) error {
+
+				return nil
+			},
+		},
+	}
+
+	tagit := New(mockConsulClient, mockExecutor, "test-service", "echo test", 100*time.Millisecond, "tag")
+
+	// Start Run in a goroutine
+	go tagit.Run(ctx)
+
+	// Allow some time to pass and then cancel the context
+	time.Sleep(350 * time.Millisecond) // Adjust this duration as needed
+	cancel()
+
+	// Allow some time for the goroutine to react to the context cancellation
+	time.Sleep(50 * time.Millisecond)
+
+	// Check if updateServiceTags was called as expected
+	if updateServiceTagsCalled.Load() < 2 || updateServiceTagsCalled.Load() > 3 {
+		t.Errorf("Expected updateServiceTags to be called 2 or 3 times, got %d", updateServiceTagsCalled)
+	}
+}