testserver.go

Documentation: k8s.io/kubernetes/cmd/kube-scheduler/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  	"k8s.io/component-base/configz"
    32  	logsapi "k8s.io/component-base/logs/api/v1"
    33  	"k8s.io/kubernetes/cmd/kube-scheduler/app"
    34  	kubeschedulerconfig "k8s.io/kubernetes/cmd/kube-scheduler/app/config"
    35  	"k8s.io/kubernetes/cmd/kube-scheduler/app/options"
    36  
    37  	"k8s.io/klog/v2"
    38  )
    39  
    40  func init() {
    41  	// If instantiated more than once or together with other servers, the
    42  	// servers would try to modify the global logging state. This must get
    43  	// ignored during testing.
    44  	logsapi.ReapplyHandling = logsapi.ReapplyHandlingIgnoreUnchanged
    45  }
    46  
    47  // TearDownFunc is to be called to tear down a test server.
    48  type TearDownFunc func()
    49  
    50  // TestServer return values supplied by kube-test-ApiServer
    51  type TestServer struct {
    52  	LoopbackClientConfig *restclient.Config // Rest client config using the magic token
    53  	Options              *options.Options
    54  	Config               *kubeschedulerconfig.Config
    55  	TearDownFn           TearDownFunc // TearDown function
    56  	TmpDir               string       // Temp Dir used, by the apiserver
    57  }
    58  
    59  // StartTestServer starts a kube-scheduler. A rest client config and a tear-down func,
    60  // and location of the tmpdir are returned.
    61  //
    62  // Note: we return a tear-down func instead of a stop channel because the later will leak temporary
    63  //
    64  //	files that because Golang testing's call to os.Exit will not give a stop channel go routine
    65  //	enough time to remove temporary files.
    66  func StartTestServer(ctx context.Context, customFlags []string) (result TestServer, err error) {
    67  	logger := klog.FromContext(ctx)
    68  	ctx, cancel := context.WithCancel(ctx)
    69  
    70  	var errCh chan error
    71  	tearDown := func() {
    72  		cancel()
    73  
    74  		// If the scheduler was started, let's wait for it to
    75  		// shutdown clearly.
    76  		if errCh != nil {
    77  			err, ok := <-errCh
    78  			if ok && err != nil {
    79  				logger.Error(err, "Failed to shutdown test server clearly")
    80  			}
    81  		}
    82  		if len(result.TmpDir) != 0 {
    83  			os.RemoveAll(result.TmpDir)
    84  		}
    85  		configz.Delete("componentconfig")
    86  	}
    87  	defer func() {
    88  		if result.TearDownFn == nil {
    89  			tearDown()
    90  		}
    91  	}()
    92  
    93  	result.TmpDir, err = os.MkdirTemp("", "kube-scheduler")
    94  	if err != nil {
    95  		return result, fmt.Errorf("failed to create temp dir: %v", err)
    96  	}
    97  
    98  	fs := pflag.NewFlagSet("test", pflag.PanicOnError)
    99  
   100  	opts := options.NewOptions()
   101  	nfs := opts.Flags
   102  	for _, f := range nfs.FlagSets {
   103  		fs.AddFlagSet(f)
   104  	}
   105  	fs.Parse(customFlags)
   106  
   107  	if opts.SecureServing.BindPort != 0 {
   108  		opts.SecureServing.Listener, opts.SecureServing.BindPort, err = createListenerOnFreePort()
   109  		if err != nil {
   110  			return result, fmt.Errorf("failed to create listener: %v", err)
   111  		}
   112  		opts.SecureServing.ServerCert.CertDirectory = result.TmpDir
   113  
   114  		logger.Info("kube-scheduler will listen securely", "port", opts.SecureServing.BindPort)
   115  	}
   116  
   117  	cc, sched, err := app.Setup(ctx, opts)
   118  	if err != nil {
   119  		return result, fmt.Errorf("failed to create config from options: %v", err)
   120  	}
   121  
   122  	errCh = make(chan error)
   123  	go func(ctx context.Context) {
   124  		defer close(errCh)
   125  		if err := app.Run(ctx, cc, sched); err != nil {
   126  			errCh <- err
   127  		}
   128  	}(ctx)
   129  
   130  	logger.Info("Waiting for /healthz to be ok...")
   131  	client, err := kubernetes.NewForConfig(cc.LoopbackClientConfig)
   132  	if err != nil {
   133  		return result, fmt.Errorf("failed to create a client: %v", err)
   134  	}
   135  	err = wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, 30*time.Second, false, func(ctx context.Context) (bool, error) {
   136  		select {
   137  		case err := <-errCh:
   138  			return false, err
   139  		default:
   140  		}
   141  
   142  		result := client.CoreV1().RESTClient().Get().AbsPath("/healthz").Do(ctx)
   143  		status := 0
   144  		result.StatusCode(&status)
   145  		if status == 200 {
   146  			return true, nil
   147  		}
   148  		return false, nil
   149  	})
   150  	if err != nil {
   151  		return result, fmt.Errorf("failed to wait for /healthz to return ok: %v", err)
   152  	}
   153  
   154  	// from here the caller must call tearDown
   155  	result.LoopbackClientConfig = cc.LoopbackClientConfig
   156  	result.Options = opts
   157  	result.Config = cc.Config
   158  	result.TearDownFn = tearDown
   159  
   160  	return result, nil
   161  }
   162  
   163  // StartTestServerOrDie calls StartTestServer panic if it does not succeed.
   164  func StartTestServerOrDie(ctx context.Context, flags []string) *TestServer {
   165  	result, err := StartTestServer(ctx, flags)
   166  	if err == nil {
   167  		return &result
   168  	}
   169  
   170  	panic(fmt.Errorf("failed to launch server: %v", err))
   171  }
   172  
   173  func createListenerOnFreePort() (net.Listener, int, error) {
   174  	ln, err := net.Listen("tcp", ":0")
   175  	if err != nil {
   176  		return nil, 0, err
   177  	}
   178  
   179  	// get port
   180  	tcpAddr, ok := ln.Addr().(*net.TCPAddr)
   181  	if !ok {
   182  		ln.Close()
   183  		return nil, 0, fmt.Errorf("invalid listen address: %q", ln.Addr().String())
   184  	}
   185  
   186  	return ln, tcpAddr.Port, nil
   187  }
   188
View as plain text