github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/pprof/pprof-test/main_test.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package main_test
     6  
     7  import (
     8  	"bufio"
     9  	"errors"
    10  	"net/http"
    11  	"os"
    12  	"path/filepath"
    13  	"reflect"
    14  	"sync"
    15  	"testing"
    16  	"time"
    17  
    18  	"github.com/grailbio/testutil"
    19  	"github.com/stretchr/testify/require"
    20  	"v.io/x/lib/gosh"
    21  )
    22  
    23  func retry(url string, duration time.Duration, iterations int) error {
    24  	i := 0
    25  	for {
    26  		if iterations == i {
    27  			return errors.New("max attempts exceeded for URL")
    28  		}
    29  
    30  		i++
    31  		response, err := http.Get(url)
    32  		if err == nil {
    33  			if response.StatusCode == 200 {
    34  				return nil
    35  			}
    36  
    37  			return errors.New("page not found")
    38  		}
    39  
    40  		time.Sleep(duration)
    41  	}
    42  }
    43  
    44  func find(t *testing.T, dir string, files []string) []string {
    45  	if len(files) == 0 {
    46  		return nil
    47  	}
    48  	d, err := os.Open(dir)
    49  	require.NoError(t, err)
    50  	entries, err := d.Readdirnames(0)
    51  	require.NoError(t, err)
    52  	found := []string{}
    53  	for _, want := range files {
    54  		for _, f := range entries {
    55  			if f == want {
    56  				found = append(found, want)
    57  				break
    58  			}
    59  		}
    60  	}
    61  	return found
    62  }
    63  
    64  func TestEverything(t *testing.T) {
    65  	sh := gosh.NewShell(nil)
    66  	bin := testutil.GoExecutable(t, "//go/src/github.com/Schaudge/grailbase/pprof/pprof-test/pprof-test")
    67  	tempDir := sh.MakeTempDir()
    68  	cpu := filepath.Join(tempDir, "c")
    69  	heap := filepath.Join(tempDir, "h")
    70  
    71  	for _, tst := range []struct {
    72  		name    string
    73  		args    []string
    74  		files   []string
    75  		success bool
    76  	}{
    77  		{"without", []string{}, nil, false},
    78  		{"with", []string{"-pprof=127.0.0.1:0"}, nil, true},
    79  		{"with-cl", []string{"-pprof=127.0.0.1:0", "-cpu-profile=" + cpu, "-heap-profile=" + heap}, []string{"c-00000.pprof", "h-00000.pprof"}, true},
    80  	} {
    81  		addrs := make(chan string)
    82  
    83  		var wg sync.WaitGroup
    84  		wg.Add(1)
    85  		go func() {
    86  			cmd := sh.Cmd(bin, tst.args...)
    87  			rc := cmd.StdoutPipe()
    88  			cmd.Start()
    89  			require.NoError(t, cmd.Err)
    90  			scanner := bufio.NewScanner(rc)
    91  			for cmd.Err == nil && scanner.Scan() {
    92  				a := scanner.Text()
    93  				addrs <- a
    94  			}
    95  			rc.Close()
    96  			cmd.Wait()
    97  			close(addrs)
    98  			wg.Done()
    99  		}()
   100  
   101  		mainAddr := <-addrs
   102  		debugAddr := <-addrs
   103  
   104  		// Make sure the main webserver is alive/test app is running.
   105  		if err := retry("http://"+mainAddr+"/alive/", 100*time.Millisecond, 50); err != nil {
   106  			t.Fatal(err)
   107  		}
   108  
   109  		// Make sure that /debug/ is not on the main webserver.
   110  		if err := retry("http://"+mainAddr+"/debug/pprof/", 50*time.Millisecond, 1); err == nil {
   111  			t.Fatal("found /debug/ on main server")
   112  		}
   113  
   114  		// Make sure that /debug/ is on the debug server, if configured.
   115  		if err := retry("http://"+debugAddr+"/debug/pprof/", 100*time.Millisecond, 50); tst.success != (err == nil) {
   116  			t.Fatalf("error detecting /debug/ %v %v %v %v %v", tst.success, err == nil, err, mainAddr, debugAddr)
   117  		}
   118  
   119  		// Quit the test app.
   120  		http.Get("http://" + mainAddr + "/quitquitquit/")
   121  
   122  		// Wait for the command to finish before testing for the files it
   123  		// should have written.
   124  		wg.Wait()
   125  
   126  		if got, want := find(t, tempDir, tst.files), tst.files; !reflect.DeepEqual(got, want) {
   127  			t.Errorf("%s: got %v, want %v", tst.name, got, want)
   128  		}
   129  	}
   130  }