feat: add castai provider (#16)

internal/cast/cast.go → internal/castai/castai.go

@@ -1,20 +1,22 @@
 //go:generate mockgen -destination ./mock/client.go . Client
-package cast
+package castai
 
 import (
 	"bytes"
-	"castai-agent/internal/config"
 	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/go-resty/resty/v2"
-	"github.com/sirupsen/logrus"
 	"io"
 	"mime/multipart"
 	"net/http"
 	"net/textproto"
 	"net/url"
 	"time"
+
+	"github.com/go-resty/resty/v2"
+	"github.com/sirupsen/logrus"
+
+	"castai-agent/internal/config"
 )
 
 const (

internal/cast/cast_test.go → internal/castai/castai_test.go

@@ -1,4 +1,4 @@
-package cast
+package castai
 
 import (
 	"castai-agent/internal/services/collector"

internal/cast/types.go → internal/castai/types.go

@@ -1,4 +1,4 @@
-package cast
+package castai
 
 import "castai-agent/internal/services/collector"
 

internal/config/config.go

@@ -2,13 +2,15 @@ package config
 
 import (
 	"fmt"
+
 	"github.com/spf13/viper"
 )
 
 type Config struct {
 	API        API
 	Kubeconfig string
 	Provider   string
+	CASTAI     *CASTAI
 	EKS        *EKS
 }
 
@@ -17,6 +19,11 @@ type API struct {
 	URL string
 }
 
+type CASTAI struct {
+	ClusterID      string
+	OrganizationID string
+}
+
 type EKS struct {
 	AccountID   string
 	Region      string
@@ -25,6 +32,7 @@ type EKS struct {
 
 var cfg *Config
 
+// Get configuration bound to environment variables.
 func Get() Config {
 	if cfg != nil {
 		return *cfg
@@ -37,6 +45,9 @@ func Get() Config {
 
 	_ = viper.BindEnv("provider")
 
+	_ = viper.BindEnv("castai.clusterid", "CASTAI_CLUSTER_ID")
+	_ = viper.BindEnv("castai.organizationid", "CASTAI_ORGANIZATION_ID")
+
 	_ = viper.BindEnv("eks.accountid", "EKS_ACCOUNT_ID")
 	_ = viper.BindEnv("eks.region", "EKS_REGION")
 	_ = viper.BindEnv("eks.clustername", "EKS_CLUSTER_NAME")
@@ -53,6 +64,15 @@ func Get() Config {
 		required("API_URL")
 	}
 
+	if cfg.CASTAI != nil {
+		if cfg.CASTAI.ClusterID == "" {
+			requiredDiscoveryDisabled("CASTAI_CLUSTER_ID")
+		}
+		if cfg.CASTAI.OrganizationID == "" {
+			requiredDiscoveryDisabled("CASTAI_ORGANIZATION_ID")
+		}
+	}
+
 	if cfg.EKS != nil {
 		if cfg.EKS.AccountID == "" {
 			requiredDiscoveryDisabled("EKS_ACCOUNT_ID")
@@ -68,6 +88,11 @@ func Get() Config {
 	return *cfg
 }
 
+// Reset is used only for unit testing to reset configuration and rebind variables.
+func Reset() {
+	cfg = nil
+}
+
 func required(variable string) {
 	panic(fmt.Errorf("env variable %s is required", variable))
 }

internal/config/config_test.go

@@ -1,9 +1,10 @@
 package config
 
 import (
-	"github.com/stretchr/testify/require"
 	"os"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestConfig(t *testing.T) {

internal/services/collector/collector.go

@@ -4,12 +4,13 @@ package collector
 import (
 	"context"
 	"fmt"
+	"regexp"
+	"strconv"
+
 	"github.com/sirupsen/logrus"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/version"
 	"k8s.io/client-go/kubernetes"
-	"regexp"
-	"strconv"
 )
 
 // Collector is responsible for gathering K8s data from the cluster.

internal/services/providers/castai/castai.go

@@ -0,0 +1,61 @@
+package castai
+
+import (
+	"context"
+
+	"github.com/sirupsen/logrus"
+	"k8s.io/api/core/v1"
+
+	"castai-agent/internal/castai"
+	"castai-agent/internal/config"
+	"castai-agent/internal/services/providers/types"
+	"castai-agent/pkg/labels"
+)
+
+const (
+	Name = "castai"
+)
+
+func New(_ context.Context, log logrus.FieldLogger) (types.Provider, error) {
+	return &Provider{log: log}, nil
+}
+
+type Provider struct {
+	log logrus.FieldLogger
+}
+
+func (p *Provider) RegisterCluster(_ context.Context, _ castai.Client) (*types.ClusterRegistration, error) {
+	cfg := config.Get().CASTAI
+	return &types.ClusterRegistration{
+		ClusterID:      cfg.ClusterID,
+		OrganizationID: cfg.OrganizationID,
+	}, nil
+}
+
+func (p *Provider) Name() string {
+	return Name
+}
+
+func (p *Provider) FilterSpot(_ context.Context, nodes []*v1.Node) ([]*v1.Node, error) {
+	var spots []*v1.Node
+
+	for _, n := range nodes {
+		if val, ok := n.ObjectMeta.Labels[labels.Spot]; ok && val == "true" {
+			spots = append(spots, n)
+		}
+	}
+
+	return spots, nil
+}
+
+func (p *Provider) AccountID(_ context.Context) (string, error) {
+	return "", nil
+}
+
+func (p *Provider) ClusterName(_ context.Context) (string, error) {
+	return "", nil
+}
+
+func (p *Provider) ClusterRegion(_ context.Context) (string, error) {
+	return "", nil
+}

internal/services/providers/eks/client/aws.go

@@ -5,14 +5,15 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
+	"strings"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/sirupsen/logrus"
 	"k8s.io/utils/pointer"
-	"math"
-	"strings"
 )
 
 // Client is an abstraction on the AWS SDK to enable easier mocking and manipulation of request data.

internal/services/providers/eks/eks.go

@@ -1,23 +1,26 @@
 package eks
 
 import (
-	"castai-agent/internal/cast"
-	"castai-agent/internal/config"
-	"castai-agent/internal/services/providers/eks/client"
 	"context"
 	"fmt"
+	"time"
+
 	"github.com/patrickmn/go-cache"
 	"github.com/sirupsen/logrus"
 	"k8s.io/api/core/v1"
-	"time"
+
+	"castai-agent/internal/castai"
+	"castai-agent/internal/config"
+	"castai-agent/internal/services/providers/eks/client"
+	"castai-agent/internal/services/providers/types"
 )
 
 const (
 	Name = "eks"
 )
 
 // New configures and returns an EKS provider.
-func New(ctx context.Context, log logrus.FieldLogger) (*Provider, error) {
+func New(ctx context.Context, log logrus.FieldLogger) (types.Provider, error) {
 	log = log.WithField("provider", Name)
 
 	var opts []client.Opt
@@ -48,6 +51,40 @@ type Provider struct {
 	spotCache *cache.Cache
 }
 
+func (p *Provider) RegisterCluster(ctx context.Context, client castai.Client) (*types.ClusterRegistration, error) {
+	cn, err := p.awsClient.GetClusterName(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("getting cluster name: %w", err)
+	}
+	r, err := p.awsClient.GetRegion(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("getting region: %w", err)
+	}
+	accID, err := p.awsClient.GetAccountID(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("getting account id: %w", err)
+	}
+
+	req := &castai.RegisterClusterRequest{
+		Name: *cn,
+		EKS: castai.EKSParams{
+			ClusterName: *cn,
+			Region:      *r,
+			AccountID:   *accID,
+		},
+	}
+
+	resp, err := client.RegisterCluster(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("requesting castai api: %w", err)
+	}
+
+	return &types.ClusterRegistration{
+		ClusterID:      resp.ID,
+		OrganizationID: resp.OrganizationID,
+	}, nil
+}
+
 func (p *Provider) FilterSpot(ctx context.Context, nodes []*v1.Node) ([]*v1.Node, error) {
 	if p.spotCache == nil {
 		p.spotCache = cache.New(60*time.Minute, 10*time.Minute)
@@ -124,7 +161,7 @@ func (p *Provider) AccountID(ctx context.Context) (string, error) {
 	return *accID, nil
 }
 
-func (p *Provider) RegisterClusterRequest(ctx context.Context) (*cast.RegisterClusterRequest, error) {
+func (p *Provider) RegisterClusterRequest(ctx context.Context) (*castai.RegisterClusterRequest, error) {
 	cn, err := p.awsClient.GetClusterName(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("getting cluster name: %w", err)
@@ -137,9 +174,9 @@ func (p *Provider) RegisterClusterRequest(ctx context.Context) (*cast.RegisterCl
 	if err != nil {
 		return nil, fmt.Errorf("getting account id: %w", err)
 	}
-	return &cast.RegisterClusterRequest{
+	return &castai.RegisterClusterRequest{
 		Name: *cn,
-		EKS: cast.EKSParams{
+		EKS: castai.EKSParams{
 			ClusterName: *cn,
 			Region:      *r,
 			AccountID:   *accID,