move client functions

main.go

@@ -1,58 +1,19 @@
 package main
 
 import (
-	"context"
-	"time"
-
 	"github.com/docker/docker/api/types"
-	"github.com/docker/docker/api/types/events"
-	"github.com/docker/docker/api/types/filters"
-	"github.com/docker/docker/client"
 	"github.com/patrickmn/go-cache"
 	"github.com/rs/zerolog/log"
 	"github.com/shinebayar-g/ufw-docker-automated/logger"
 	"github.com/shinebayar-g/ufw-docker-automated/ufwhandler"
 )
 
-func createClient() (*context.Context, *client.Client, error) {
-	ctx := context.Background()
-	client, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
-	if err != nil {
-		return &ctx, client, err
-	}
-	_, err = client.Info(ctx)
-	return &ctx, client, err
-}
-
-func streamEvents(ctx *context.Context, c *client.Client) (<-chan events.Message, <-chan error) {
-	filter := filters.NewArgs()
-	filter.Add("type", "container")
-	filter.Add("label", "UFW_MANAGED=TRUE")
-	return c.Events(*ctx, types.EventsOptions{Filters: filter})
-}
-
-func reconnect() (*context.Context, *client.Client) {
-	var ctx *context.Context
-	var client *client.Client
-	var err error
-	for {
-		time.Sleep(5 * time.Second)
-		log.Info().Msg("ufw-docker-automated: Trying to reconnect..")
-		ctx, client, err = createClient()
-		if err == nil {
-			break
-		}
-	}
-	log.Info().Msg("ufw-docker-automated: Reconnected to the Docker Engine.")
-	return ctx, client
-}
-
 func main() {
 	logger.SetupLogger()
-	ctx, client, err := createClient()
+	ctx, client, err := ufwhandler.CreateClient()
 	if err != nil {
 		log.Error().Err(err).Msg("ufw-docker-automated: Client error.")
-		ctx, client = reconnect()
+		ctx, client = ufwhandler.Reconnect()
 	} else {
 		log.Info().Msg("ufw-docker-automated: Connected to the Docker Engine.")
 	}
@@ -65,7 +26,7 @@ func main() {
 	go ufwhandler.Cleanup(ctx, client)
 	go ufwhandler.Sync(ctx, createChannel, client)
 
-	messages, errors := streamEvents(ctx, client)
+	messages, errors := ufwhandler.StreamEvents(ctx, client)
 	for {
 		select {
 		case msg := <-messages:
@@ -83,9 +44,9 @@ func main() {
 		case err := <-errors:
 			if err != nil {
 				log.Error().Err(err).Msg("ufw-docker-automated: Event error.")
-				ctx, client = reconnect()
+				ctx, client = ufwhandler.Reconnect()
 				go ufwhandler.Sync(ctx, createChannel, client)
-				messages, errors = streamEvents(ctx, client)
+				messages, errors = ufwhandler.StreamEvents(ctx, client)
 			}
 		}
 	}

ufwhandler/client.go

@@ -0,0 +1,45 @@
+package ufwhandler
+
+import (
+	"context"
+	"time"
+
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/events"
+	"github.com/docker/docker/api/types/filters"
+	"github.com/docker/docker/client"
+	"github.com/rs/zerolog/log"
+)
+
+func CreateClient() (*context.Context, *client.Client, error) {
+	ctx := context.Background()
+	client, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
+	if err != nil {
+		return &ctx, client, err
+	}
+	_, err = client.Info(ctx)
+	return &ctx, client, err
+}
+
+func Reconnect() (*context.Context, *client.Client) {
+	var ctx *context.Context
+	var client *client.Client
+	var err error
+	for {
+		time.Sleep(5 * time.Second)
+		log.Info().Msg("ufw-docker-automated: Trying to reconnect..")
+		ctx, client, err = CreateClient()
+		if err == nil {
+			break
+		}
+	}
+	log.Info().Msg("ufw-docker-automated: Reconnected to the Docker Engine.")
+	return ctx, client
+}
+
+func StreamEvents(ctx *context.Context, c *client.Client) (<-chan events.Message, <-chan error) {
+	filter := filters.NewArgs()
+	filter.Add("type", "container")
+	filter.Add("label", "UFW_MANAGED=TRUE")
+	return c.Events(*ctx, types.EventsOptions{Filters: filter})
+}