golang.org/x/build@v0.0.0-20240506185731-218518f32b70/cmd/runqemubuildlet/heartbeat_test.go (about)

     1  // Copyright 2021 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"context"
     9  	"errors"
    10  	"fmt"
    11  	"net/http"
    12  	"net/http/httptest"
    13  	"net/url"
    14  	"testing"
    15  	"time"
    16  )
    17  
    18  func TestCheckBuildletHealth(t *testing.T) {
    19  	cases := []struct {
    20  		desc     string
    21  		respCode int
    22  		wantErr  bool
    23  	}{
    24  		{
    25  			desc:     "success",
    26  			respCode: http.StatusOK,
    27  		},
    28  		{
    29  			desc:     "failure",
    30  			respCode: http.StatusBadGateway,
    31  			wantErr:  true,
    32  		},
    33  	}
    34  	for _, c := range cases {
    35  		t.Run(c.desc, func(t *testing.T) {
    36  			m := http.NewServeMux()
    37  			m.HandleFunc("/healthz", func(w http.ResponseWriter, req *http.Request) {
    38  				w.WriteHeader(c.respCode)
    39  				fmt.Sprintln(w, "ok")
    40  			})
    41  			s := httptest.NewServer(m)
    42  			defer s.Close()
    43  			u, err := url.Parse(s.URL)
    44  			if err != nil {
    45  				t.Fatalf("url.Parse(%q) = %v, wanted no error", s.URL, err)
    46  			}
    47  			u.Path = "/healthz"
    48  
    49  			if err := checkBuildletHealth(context.Background(), u.String()); (err != nil) != c.wantErr {
    50  				t.Errorf("checkBuildletHealth(_, %q) = %v, wantErr: %t", s.URL, err, c.wantErr)
    51  			}
    52  		})
    53  	}
    54  }
    55  
    56  func TestHeartbeatContext(t *testing.T) {
    57  	ctx := context.Background()
    58  
    59  	didWork := make(chan interface{}, 2)
    60  	done := make(chan interface{})
    61  	ctx, cancel := heartbeatContext(ctx, time.Millisecond, 100*time.Millisecond, func(context.Context) error {
    62  		select {
    63  		case <-done:
    64  			return errors.New("heartbeat stopped")
    65  		case didWork <- nil:
    66  		default:
    67  		}
    68  		return nil
    69  	})
    70  	defer cancel()
    71  
    72  	select {
    73  	case <-time.After(5 * time.Second):
    74  		t.Errorf("heatbeatContext() never called f, wanted at least one call")
    75  	case <-didWork:
    76  	}
    77  
    78  	select {
    79  	case <-done:
    80  		t.Errorf("heartbeatContext() finished early, wanted it to still be testing")
    81  	case <-didWork:
    82  		close(done)
    83  	}
    84  
    85  	select {
    86  	case <-time.After(5 * time.Second):
    87  		t.Errorf("heartbeatContext() did not timeout, wanted timeout after failing over %v", time.Second)
    88  	case <-ctx.Done():
    89  		// heartbeatContext() successfully timed out after failing
    90  	}
    91  }