refactor: use RunE for testable commands and improve coverage

cmd/root_test.go

@@ -0,0 +1,98 @@
+/*
+Copyright © 2022 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 (
+	"bytes"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRootCmd_Structure(t *testing.T) {
+	assert.Equal(t, "ballot", rootCmd.Use)
+	assert.NotEmpty(t, rootCmd.Short)
+	assert.NotEmpty(t, rootCmd.Long)
+}
+
+func TestRootCmd_HasRunSubcommand(t *testing.T) {
+	found := false
+	for _, cmd := range rootCmd.Commands() {
+		if cmd.Use == "run" {
+			found = true
+			break
+		}
+	}
+	assert.True(t, found, "rootCmd should have 'run' subcommand")
+}
+
+func TestRootCmd_ExecuteHelp(t *testing.T) {
+	// Capture output
+	buf := new(bytes.Buffer)
+	rootCmd.SetOut(buf)
+	rootCmd.SetErr(buf)
+	rootCmd.SetArgs([]string{"--help"})
+
+	err := rootCmd.Execute()
+	assert.NoError(t, err)
+	assert.Contains(t, buf.String(), "ballot")
+}
+
+func TestInitConfig_WithConfigFile(t *testing.T) {
+	// Create a temporary config file
+	tmpDir := t.TempDir()
+	configPath := filepath.Join(tmpDir, "test-config.yaml")
+	configContent := []byte(`
+election:
+  enabled:
+    - test_service
+`)
+	err := os.WriteFile(configPath, configContent, 0644)
+	assert.NoError(t, err)
+
+	// Reset viper and set config file
+	viper.Reset()
+	cfgFile = configPath
+
+	// Call initConfig
+	initConfig()
+
+	// Verify the config was loaded
+	enabled := viper.GetStringSlice("election.enabled")
+	assert.Contains(t, enabled, "test_service")
+
+	// Clean up
+	cfgFile = ""
+	viper.Reset()
+}
+
+func TestInitConfig_WithoutConfigFile(t *testing.T) {
+	// Reset viper
+	viper.Reset()
+	cfgFile = ""
+
+	// This should not panic even without a config file
+	initConfig()
+}
+
+func TestRootCmd_PersistentFlags(t *testing.T) {
+	flag := rootCmd.PersistentFlags().Lookup("config")
+	assert.NotNil(t, flag)
+	assert.Equal(t, "config", flag.Name)
+}

cmd/run.go

@@ -17,6 +17,7 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"os/signal"
 	"sync"
@@ -32,41 +33,56 @@ import (
 var runCmd = &cobra.Command{
 	Use:   "run",
 	Short: "Run the ballot and starts all the defined elections",
-	Run: func(cmd *cobra.Command, args []string) {
-		ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
-		defer cancel()
-
-		var wg sync.WaitGroup
-		enabledServices := viper.GetStringSlice("election.enabled")
-
-		for _, name := range enabledServices {
-			b, err := ballot.New(ctx, name)
-			if err != nil {
-				log.WithFields(log.Fields{
-					"caller":  "run",
-					"step":    "New",
-					"service": name,
-				}).Error(err)
-				os.Exit(1)
-			}
+	RunE:  runElection,
+}
+
+func runElection(cmd *cobra.Command, args []string) error {
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
+
+	return runElectionWithContext(ctx)
+}
+
+func runElectionWithContext(ctx context.Context) error {
+	var wg sync.WaitGroup
+	enabledServices := viper.GetStringSlice("election.enabled")
+
+	if len(enabledServices) == 0 {
+		return fmt.Errorf("no services enabled for election")
+	}
 
-			wg.Add(1)
-			go func(b *ballot.Ballot, name string) {
-				defer wg.Done()
-				err := b.Run()
-				if err != nil {
-					log.WithFields(log.Fields{
-						"caller":  "run",
-						"step":    "runCmd",
-						"service": name,
-					}).Error(err)
-				}
-			}(b, name)
+	errCh := make(chan error, len(enabledServices))
+
+	for _, name := range enabledServices {
+		b, err := ballot.New(ctx, name)
+		if err != nil {
+			return fmt.Errorf("failed to create ballot for service %s: %w", name, err)
 		}
 
-		wg.Wait()
-		log.Info("All elections stopped, shutting down")
-	},
+		wg.Add(1)
+		go func(b *ballot.Ballot, name string) {
+			defer wg.Done()
+			if err := b.Run(); err != nil {
+				errCh <- fmt.Errorf("service %s: %w", name, err)
+			}
+		}(b, name)
+	}
+
+	wg.Wait()
+	close(errCh)
+
+	// Collect any errors from running elections
+	var errs []error
+	for err := range errCh {
+		errs = append(errs, err)
+	}
+
+	if len(errs) > 0 {
+		return fmt.Errorf("election errors: %v", errs)
+	}
+
+	log.Info("All elections stopped, shutting down")
+	return nil
 }
 
 func init() {

cmd/run_test.go

@@ -0,0 +1,71 @@
+/*
+Copyright © 2022 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 (
+	"context"
+	"testing"
+
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRunElectionWithContext_NoServicesEnabled(t *testing.T) {
+	viper.Reset()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	err := runElectionWithContext(ctx)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "no services enabled for election")
+}
+
+func TestRunElectionWithContext_InvalidService(t *testing.T) {
+	viper.Reset()
+	viper.Set("election.enabled", []string{"invalid_service"})
+	// Don't set the service configuration, so ballot.New will fail
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	err := runElectionWithContext(ctx)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to create ballot for service invalid_service")
+}
+
+func TestRunCmd_Structure(t *testing.T) {
+	assert.Equal(t, "run", runCmd.Use)
+	assert.NotEmpty(t, runCmd.Short)
+	assert.NotNil(t, runCmd.RunE)
+}
+
+func TestRunElection_CancelledContext(t *testing.T) {
+	viper.Reset()
+	viper.Set("election.enabled", []string{"test_service"})
+	viper.Set("election.services.test_service.id", "test_id")
+	viper.Set("election.services.test_service.key", "election/test/leader")
+
+	// Create an already cancelled context
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	// This should handle the cancelled context gracefully
+	err := runElectionWithContext(ctx)
+	// Depending on timing, may get an error from ballot.New or from Run
+	// The important thing is it doesn't panic
+	_ = err
+}