Support take over labels selector

cmd/fake-kubelet/main.go

@@ -26,6 +26,7 @@ var (
 	nodeIP                     = net.ParseIP(getEnv("NODE_IP", "196.168.0.1"))
 	nodeName                   = getEnv("NODE_NAME", "fake")
 	takeOverAll                = getEnvBool("TAKE_OVER_ALL", false)
+	takeOverLabelsSelector     = getEnv("TAKE_OVER_LABELS_SELECTOR", "type=fake-kubelet")
 	generateNodeName           = getEnv("GENERATE_NODE_NAME", "")
 	generateReplicas           = getEnv("GENERATE_REPLICAS", "0")
 	kubeconfig                 = getEnv("KUBECONFIG", "")
@@ -45,6 +46,7 @@ func init() {
 	pflag.IPVar(&nodeIP, "node_ip", nodeIP, "node ip")
 	pflag.StringVarP(&nodeName, "node_name", "n", nodeName, "node name")
 	pflag.BoolVar(&takeOverAll, "take_over_all", takeOverAll, "take over all node")
+	pflag.StringVar(&takeOverLabelsSelector, "take_over_labels_selector", takeOverLabelsSelector, "take over labels selector")
 	pflag.StringVar(&kubeconfig, "kubeconfig", kubeconfig, "kubeconfig")
 	pflag.StringVar(&master, "master", master, "master")
 	pflag.StringVar(&serverAddress, "server_address", serverAddress, "server address")
@@ -92,6 +94,7 @@ func main() {
 		ClientSet:                  clientset,
 		Nodes:                      nodes,
 		TakeOverAll:                takeOverAll,
+		TakeOverLabelsSelector:     takeOverLabelsSelector,
 		CIDR:                       cidr,
 		NodeIP:                     nodeIP.String(),
 		Logger:                     logger,

controller.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/client-go/kubernetes"
 	"sigs.k8s.io/yaml"
 )
@@ -49,6 +50,7 @@ type Config struct {
 	ClientSet                  kubernetes.Interface
 	Nodes                      []string
 	TakeOverAll                bool
+	TakeOverLabelsSelector     string
 	CIDR                       string
 	NodeIP                     string
 	Logger                     Logger
@@ -64,12 +66,19 @@ type Logger interface {
 
 // NewController creates a new fake kubelet controller
 func NewController(conf Config) (*Controller, error) {
-
 	var nodeSelectorFunc func(node *corev1.Node) bool
 	if conf.TakeOverAll {
 		nodeSelectorFunc = func(node *corev1.Node) bool {
 			return true
 		}
+	} else if conf.TakeOverLabelsSelector != "" {
+		selector, err := labels.Parse(conf.TakeOverLabelsSelector)
+		if err != nil {
+			return nil, err
+		}
+		nodeSelectorFunc = func(node *corev1.Node) bool {
+			return selector.Matches(labels.Set(node.Labels))
+		}
 	}
 
 	var lockPodsOnNodeFunc func(ctx context.Context, nodeName string) error

deploy.yaml

@@ -254,6 +254,10 @@ spec:
               value: fake-
             - name: GENERATE_REPLICAS
               value: "5"
+            - name: TAKE_OVER_LABELS_SELECTOR
+              value: type=fake-kubelet
+            - name: TAKE_OVER_ALL
+              value: "false"
             - name: CIDR
               value: 10.0.0.1/24
             - name: HEALTH_ADDRESS