testserver.go

Documentation: k8s.io/kubernetes/cmd/kube-controller-manager/app/testing

     1  /*
     2  Copyright 2018 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package testing
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"net"
    23  	"os"
    24  	"time"
    25  
    26  	"github.com/spf13/pflag"
    27  
    28  	"k8s.io/apimachinery/pkg/util/wait"
    29  	"k8s.io/client-go/kubernetes"
    30  	restclient "k8s.io/client-go/rest"
    31  	logsapi "k8s.io/component-base/logs/api/v1"
    32  	"k8s.io/klog/v2"
    33  	"k8s.io/kubernetes/cmd/kube-controller-manager/app"
    34  	kubecontrollerconfig "k8s.io/kubernetes/cmd/kube-controller-manager/app/config"
    35  	"k8s.io/kubernetes/cmd/kube-controller-manager/app/options"
    36  )
    37  
    38  func init() {
    39  	// If instantiated more than once or together with other servers, the
    40  	// servers would try to modify the global logging state. This must get
    41  	// ignored during testing.
    42  	logsapi.ReapplyHandling = logsapi.ReapplyHandlingIgnoreUnchanged
    43  }
    44  
    45  // TearDownFunc is to be called to tear down a test server.
    46  type TearDownFunc func()
    47  
    48  // TestServer return values supplied by kube-test-ApiServer
    49  type TestServer struct {
    50  	LoopbackClientConfig *restclient.Config // Rest client config using the magic token
    51  	Options              *options.KubeControllerManagerOptions
    52  	Config               *kubecontrollerconfig.Config
    53  	TearDownFn           TearDownFunc // TearDown function
    54  	TmpDir               string       // Temp Dir used, by the apiserver
    55  }
    56  
    57  // StartTestServer starts a kube-controller-manager. A rest client config and a tear-down func,
    58  // and location of the tmpdir are returned.
    59  //
    60  // Note: we return a tear-down func instead of a stop channel because the later will leak temporary
    61  // files that because Golang testing's call to os.Exit will not give a stop channel go routine
    62  // enough time to remove temporary files.
    63  func StartTestServer(ctx context.Context, customFlags []string) (result TestServer, err error) {
    64  	logger := klog.FromContext(ctx)
    65  	ctx, cancel := context.WithCancel(ctx)
    66  	var errCh chan error
    67  	tearDown := func() {
    68  		cancel()
    69  
    70  		// If the kube-controller-manager was started, let's wait for
    71  		// it to shutdown cleanly.
    72  		if errCh != nil {
    73  			err, ok := <-errCh
    74  			if ok && err != nil {
    75  				logger.Error(err, "Failed to shutdown test server cleanly")
    76  			}
    77  		}
    78  		if len(result.TmpDir) != 0 {
    79  			os.RemoveAll(result.TmpDir)
    80  		}
    81  	}
    82  	defer func() {
    83  		if result.TearDownFn == nil {
    84  			tearDown()
    85  		}
    86  	}()
    87  
    88  	result.TmpDir, err = os.MkdirTemp("", "kube-controller-manager")
    89  	if err != nil {
    90  		return result, fmt.Errorf("failed to create temp dir: %v", err)
    91  	}
    92  
    93  	fs := pflag.NewFlagSet("test", pflag.PanicOnError)
    94  
    95  	s, err := options.NewKubeControllerManagerOptions()
    96  	if err != nil {
    97  		return TestServer{}, err
    98  	}
    99  	all, disabled, aliases := app.KnownControllers(), app.ControllersDisabledByDefault(), app.ControllerAliases()
   100  	namedFlagSets := s.Flags(all, disabled, aliases)
   101  	for _, f := range namedFlagSets.FlagSets {
   102  		fs.AddFlagSet(f)
   103  	}
   104  	fs.Parse(customFlags)
   105  
   106  	if s.SecureServing.BindPort != 0 {
   107  		s.SecureServing.Listener, s.SecureServing.BindPort, err = createListenerOnFreePort()
   108  		if err != nil {
   109  			return result, fmt.Errorf("failed to create listener: %v", err)
   110  		}
   111  		s.SecureServing.ServerCert.CertDirectory = result.TmpDir
   112  
   113  		logger.Info("kube-controller-manager will listen securely", "port", s.SecureServing.BindPort)
   114  	}
   115  
   116  	config, err := s.Config(all, disabled, aliases)
   117  	if err != nil {
   118  		return result, fmt.Errorf("failed to create config from options: %v", err)
   119  	}
   120  
   121  	errCh = make(chan error)
   122  	go func(ctx context.Context) {
   123  		defer close(errCh)
   124  
   125  		if err := app.Run(ctx, config.Complete()); err != nil {
   126  			errCh <- err
   127  		}
   128  	}(ctx)
   129  
   130  	logger.Info("Waiting for /healthz to be ok...")
   131  	client, err := kubernetes.NewForConfig(config.LoopbackClientConfig)
   132  	if err != nil {
   133  		return result, fmt.Errorf("failed to create a client: %v", err)
   134  	}
   135  	err = wait.PollWithContext(ctx, 100*time.Millisecond, 30*time.Second, func(ctx context.Context) (bool, error) {
   136  		select {
   137  		case <-ctx.Done():
   138  			return false, ctx.Err()
   139  		case err := <-errCh:
   140  			return false, err
   141  		default:
   142  		}
   143  
   144  		result := client.CoreV1().RESTClient().Get().AbsPath("/healthz").Do(context.TODO())
   145  		status := 0
   146  		result.StatusCode(&status)
   147  		if status == 200 {
   148  			return true, nil
   149  		}
   150  		return false, nil
   151  	})
   152  	if err != nil {
   153  		return result, fmt.Errorf("failed to wait for /healthz to return ok: %v", err)
   154  	}
   155  
   156  	// from here the caller must call tearDown
   157  	result.LoopbackClientConfig = config.LoopbackClientConfig
   158  	result.Options = s
   159  	result.Config = config
   160  	result.TearDownFn = tearDown
   161  
   162  	return result, nil
   163  }
   164  
   165  // StartTestServerOrDie calls StartTestServer t.Fatal if it does not succeed.
   166  func StartTestServerOrDie(ctx context.Context, flags []string) *TestServer {
   167  	result, err := StartTestServer(ctx, flags)
   168  	if err == nil {
   169  		return &result
   170  	}
   171  
   172  	panic(fmt.Errorf("failed to launch server: %v", err))
   173  }
   174  
   175  func createListenerOnFreePort() (net.Listener, int, error) {
   176  	ln, err := net.Listen("tcp", ":0")
   177  	if err != nil {
   178  		return nil, 0, err
   179  	}
   180  
   181  	// get port
   182  	tcpAddr, ok := ln.Addr().(*net.TCPAddr)
   183  	if !ok {
   184  		ln.Close()
   185  		return nil, 0, fmt.Errorf("invalid listen address: %q", ln.Addr().String())
   186  	}
   187  
   188  	return ln, tcpAddr.Port, nil
   189  }
   190
View as plain text