diff --git a/cmd/cleanup.go b/cmd/cleanup.go new file mode 100644 index 0000000..a19b9e5 --- /dev/null +++ b/cmd/cleanup.go @@ -0,0 +1,58 @@ +/* +Copyright © 2024 Juliano Martinez + +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) +} diff --git a/cmd/systemd.go b/cmd/systemd.go index 9019e2b..2982860 100644 --- a/cmd/systemd.go +++ b/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(), diff --git a/pkg/tagit/tagit.go b/pkg/tagit/tagit.go index e2011c4..6d81dfe 100644 --- a/pkg/tagit/tagit.go +++ b/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{ diff --git a/pkg/tagit/tagit_test.go b/pkg/tagit/tagit_test.go index 11ce87c..c3938a7 100644 --- a/pkg/tagit/tagit_test.go +++ b/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) + } +}