Move almost everything out of autospotting.go

autospotting.go

@@ -4,91 +4,24 @@
 package main
 
 import (
-	"context"
-	"encoding/json"
 	"log"
-	"os"
 
 	autospotting "github.com/AutoSpotting/AutoSpotting/core"
-	"github.com/aws/aws-lambda-go/events"
-	"github.com/aws/aws-lambda-go/lambda"
 )
 
-var conf autospotting.Config
-
 // Version represents the build version being used
 var Version = "number missing"
 
 func main() {
-	if os.Getenv("AWS_LAMBDA_FUNCTION_NAME") != "" {
-		lambda.Start(Handler)
-	} else {
-		run()
+	conf := autospotting.Config{
+		Version: Version,
 	}
-}
-
-func run() {
+	autospotting.ParseConfig(&conf)
 
 	log.Println("Starting autospotting agent, build", Version)
 	log.Printf("Configuration flags: %#v", conf)
 
 	autospotting.Run(&conf)
-	log.Println("Execution completed, nothing left to do")
-}
-
-// this is the equivalent of a main for when running from Lambda, but on Lambda
-// the run() is executed within the handler function every time we have an event
-func init() {
-	conf = autospotting.Config{
-		Version: Version,
-	}
-	autospotting.ParseConfig(&conf)
-}
 
-// Handler implements the AWS Lambda handler
-func Handler(ctx context.Context, rawEvent json.RawMessage) {
-
-	var snsEvent events.SNSEvent
-	var cloudwatchEvent events.CloudWatchEvent
-	parseEvent := rawEvent
-
-	// Try to parse event as an Sns Message
-	if err := json.Unmarshal(parseEvent, &snsEvent); err != nil {
-		log.Println(err.Error())
-		return
-	}
-
-	// If event is from Sns - extract Cloudwatch's one
-	if snsEvent.Records != nil {
-		snsRecord := snsEvent.Records[0]
-		parseEvent = []byte(snsRecord.SNS.Message)
-	}
-
-	// Try to parse event as Cloudwatch Event Rule
-	if err := json.Unmarshal(parseEvent, &cloudwatchEvent); err != nil {
-		log.Println(err.Error())
-		return
-	}
-
-	// If event is Instance Spot Interruption
-	if cloudwatchEvent.DetailType == "EC2 Spot Instance Interruption Warning" {
-		instanceID, err := autospotting.GetInstanceIDDueForTermination(cloudwatchEvent)
-		if err != nil || instanceID == nil {
-			return
-		}
-
-		spotTermination := autospotting.NewSpotTermination(cloudwatchEvent.Region)
-		if spotTermination.IsInAutoSpottingASG(instanceID, conf.TagFilteringMode, conf.FilterByTags) {
-			err := spotTermination.ExecuteAction(instanceID, conf.TerminationNotificationAction)
-			if err != nil {
-				log.Printf("Error executing spot termination action: %s\n", err.Error())
-			}
-		} else {
-			log.Printf("Instance %s is not in AutoSpotting ASG\n", *instanceID)
-			return
-		}
-	} else {
-		// Event is Autospotting Cron Scheduling
-		run()
-	}
+	log.Println("Execution completed, nothing left to do")
 }

core/config.go

@@ -6,8 +6,10 @@ package autospotting
 import (
 	"fmt"
 	"io"
+	"io/ioutil"
 	"log"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws/endpoints"
@@ -92,6 +94,9 @@ type Config struct {
 	// the instance role when calling CloudFormation helpers instead of the standard CloudFormation
 	// authentication method
 	PatchBeanstalkUserdata string
+
+	logger *log.Logger
+	debug  *log.Logger
 }
 
 // ParseConfig loads configuration from command line flags, environments variables, and config files.
@@ -113,6 +118,13 @@ func ParseConfig(conf *Config) {
 	conf.MainRegion = region
 	conf.SleepMultiplier = 1
 
+	conf.logger = log.New(conf.LogFile, "", conf.LogFlag)
+	if os.Getenv("AUTOSPOTTING_DEBUG") == "true" {
+		conf.debug = log.New(conf.LogFile, "", conf.LogFlag)
+	} else {
+		conf.debug = log.New(ioutil.Discard, "", 0)
+	}
+
 	flagSet.StringVar(&conf.AllowedInstanceTypes, "allowed_instance_types", "",
 		"\n\tIf specified, the spot instances will be searched only among these types.\n\tIf missing, any instance type is allowed.\n"+
 			"\tAccepts a list of comma or whitespace separated instance types (supports globs).\n"+
@@ -194,4 +206,24 @@ func ParseConfig(conf *Config) {
 		log.Fatal(err.Error())
 	}
 	conf.InstanceData = data
+
+	addDefaultFilteringMode(conf)
+	addDefaultFilter(conf)
+}
+
+func addDefaultFilteringMode(conf *Config) {
+	if conf.TagFilteringMode != "opt-out" {
+		conf.TagFilteringMode = "opt-in"
+	}
+}
+
+func addDefaultFilter(conf *Config) {
+	if len(strings.TrimSpace(conf.FilterByTags)) == 0 {
+		switch conf.TagFilteringMode {
+		case "opt-out":
+			conf.FilterByTags = "spot-enabled=false"
+		default:
+			conf.FilterByTags = "spot-enabled=true"
+		}
+	}
 }

core/config_test.go

@@ -5,6 +5,7 @@ package autospotting
 
 import (
 	"os"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -77,3 +78,93 @@ func TestParseConfig(t *testing.T) {
 		}
 	}
 }
+
+func Test_addDefaultFilter(t *testing.T) {
+
+	tests := []struct {
+		name   string
+		config Config
+		want   string
+	}{
+		{
+			name:   "Default No ASG Tags",
+			config: Config{},
+			want:   "spot-enabled=true",
+		},
+		{
+			name: "Specified ASG Tags",
+			config: Config{
+				FilterByTags: "environment=dev",
+			},
+			want: "environment=dev",
+		},
+		{
+			name: "Specified ASG that is just whitespace",
+			config: Config{
+				FilterByTags: "         ",
+			},
+			want: "spot-enabled=true",
+		},
+		{
+			name:   "Default No ASG Tags",
+			config: Config{TagFilteringMode: "opt-out"},
+			want:   "spot-enabled=false",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			addDefaultFilter(&tt.config)
+
+			if !reflect.DeepEqual(tt.config.FilterByTags, tt.want) {
+				t.Errorf("addDefaultFilter() = %v, want %v", tt.config.FilterByTags, tt.want)
+			}
+		})
+	}
+}
+
+func Test_addDefaultFilteringMode(t *testing.T) {
+	tests := []struct {
+		name string
+		cfg  Config
+		want string
+	}{
+		{
+			name: "Missing FilterMode",
+			cfg:  Config{TagFilteringMode: ""},
+			want: "opt-in",
+		},
+		{
+			name: "Opt-in FilterMode",
+			cfg: Config{
+				TagFilteringMode: "opt-in",
+			},
+			want: "opt-in",
+		},
+		{
+			name: "Opt-out FilterMode",
+			cfg: Config{
+				TagFilteringMode: "opt-out",
+			},
+			want: "opt-out",
+		},
+		{
+			name: "Anything else gives the opt-in FilterMode",
+			cfg: Config{
+				TagFilteringMode: "whatever",
+			},
+			want: "opt-in",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			addDefaultFilteringMode(&tt.cfg)
+			if !reflect.DeepEqual(tt.cfg.TagFilteringMode, tt.want) {
+				t.Errorf("addDefaultFilteringMode() = %v, want %v",
+					tt.cfg.TagFilteringMode, tt.want)
+			}
+		})
+	}
+}