github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/diagnostic/dump/dump_test.go (about)

     1  // Copyright 2019 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 dump
     6  
     7  import (
     8  	"archive/zip"
     9  	"bytes"
    10  	"context"
    11  	"errors"
    12  	"io"
    13  	"io/ioutil"
    14  	"net/http"
    15  	"net/http/httptest"
    16  	"reflect"
    17  	"regexp"
    18  	"sort"
    19  	"sync"
    20  	"testing"
    21  )
    22  
    23  func makeDumpConst(errC chan<- error, s string) Func {
    24  	return func(ctx context.Context, w io.Writer) error {
    25  		if _, err := w.Write([]byte(s)); err != nil {
    26  			// This should not happen, so we let the main test goroutine know.
    27  			errC <- err
    28  		}
    29  		return nil
    30  	}
    31  }
    32  
    33  func makeDumpError(errC chan<- error, s string) Func {
    34  	return func(ctx context.Context, w io.Writer) error {
    35  		// Fake a partial failed write.
    36  		s := s[:len(s)/2]
    37  		if _, err := w.Write([]byte(s)); err != nil {
    38  			// This should not happen, so we let the main test goroutine know.
    39  			errC <- err
    40  		}
    41  		return errors.New("dump func error")
    42  	}
    43  }
    44  
    45  func dumpSkipPart(_ context.Context, _ io.Writer) error {
    46  	return ErrSkipPart
    47  }
    48  func TestShellQuote(t *testing.T) {
    49  	for _, c := range []struct {
    50  		s    string
    51  		want string
    52  	}{
    53  		{``, `''`},
    54  		{`'`, `''\'''`},
    55  		{`hello`, `'hello'`},
    56  		{`hello world`, `'hello world'`},
    57  		{`hello'world`, `'hello'\''world'`},
    58  	} {
    59  		if got, want := shellQuote(c.s), c.want; got != want {
    60  			t.Errorf("got %q, want %q", got, want)
    61  		}
    62  	}
    63  }
    64  
    65  func verifyDump(t *testing.T, server *httptest.Server, dumpFuncErrC chan error, wantNames []string) {
    66  	var dumpFuncErr error
    67  	var wg sync.WaitGroup
    68  	wg.Add(1)
    69  	go func() {
    70  		defer wg.Done()
    71  		dumpFuncErr = <-dumpFuncErrC
    72  	}()
    73  
    74  	resp, err := http.Get(server.URL + "/dump.zip")
    75  	if err != nil {
    76  		t.Fatal(err)
    77  	}
    78  	defer resp.Body.Close()
    79  	if got, want := resp.StatusCode, http.StatusOK; got != want {
    80  		t.Fatalf("got %v, want %v", got, want)
    81  	}
    82  	// Read the whole body, so we can immediately make sure that our dump
    83  	// funcs worked.
    84  	body, err := ioutil.ReadAll(resp.Body)
    85  	if err != nil {
    86  		t.Fatalf("could not read dump body: %v", err)
    87  	}
    88  	close(dumpFuncErrC)
    89  	wg.Wait()
    90  	if dumpFuncErr != nil {
    91  		t.Fatalf("unexpected error writing dump part: %v", dumpFuncErr)
    92  	}
    93  	zr, err := zip.NewReader(bytes.NewReader(body), int64(len(body)))
    94  	if err != nil {
    95  		t.Fatal(err)
    96  	}
    97  	re := regexp.MustCompile(`.*/`)
    98  	var names []string
    99  	for _, entry := range zr.File {
   100  		// Strip the prefix to recover the original name.
   101  		name := re.ReplaceAllString(entry.Name, "")
   102  		names = append(names, name)
   103  		var contents bytes.Buffer
   104  		rc, err := entry.Open()
   105  		if err != nil {
   106  			t.Fatal(err)
   107  		}
   108  		if _, err := io.Copy(&contents, rc); err != nil {
   109  			t.Fatal(err)
   110  		}
   111  		if err := rc.Close(); err != nil {
   112  			t.Fatal(err)
   113  		}
   114  		// Assume contents are "<name>-contents", matching our known
   115  		// construction of the dump contents.
   116  		if got, want := contents.String(), name+"-contents"; got != want {
   117  			t.Errorf("got %v, want %v", got, want)
   118  		}
   119  	}
   120  	sort.Strings(names)
   121  	sort.Strings(wantNames)
   122  	if got, want := names, wantNames; !reflect.DeepEqual(got, want) {
   123  		t.Errorf("got %v, want %v", got, want)
   124  	}
   125  }
   126  
   127  func TestServeHTTP(t *testing.T) {
   128  	reg := NewRegistry("abc")
   129  	dumpFuncErrC := make(chan error)
   130  	reg.Register("foo", makeDumpConst(dumpFuncErrC, "foo-contents"))
   131  	reg.Register("bar", makeDumpConst(dumpFuncErrC, "bar-contents"))
   132  	reg.Register("baz", makeDumpConst(dumpFuncErrC, "baz-contents"))
   133  
   134  	mux := http.NewServeMux()
   135  	mux.Handle("/dump.zip", reg)
   136  	server := httptest.NewServer(mux)
   137  
   138  	verifyDump(t, server, dumpFuncErrC, []string{"foo", "bar", "baz"})
   139  }
   140  
   141  func TestServeHTTPFailedParts(t *testing.T) {
   142  	reg := NewRegistry("abc")
   143  	dumpFuncErrC := make(chan error)
   144  	reg.Register("foo", makeDumpConst(dumpFuncErrC, "foo-contents"))
   145  	// Note that the following dump part funcs will return an error.
   146  	reg.Register("bar", makeDumpError(dumpFuncErrC, "bar-contents"))
   147  	reg.Register("baz", makeDumpError(dumpFuncErrC, "baz-contents"))
   148  	reg.Register("skip", dumpSkipPart)
   149  
   150  	mux := http.NewServeMux()
   151  	mux.Handle("/dump.zip", reg)
   152  	server := httptest.NewServer(mux)
   153  
   154  	// Verify that only the successful dump part func is in the dump.
   155  	verifyDump(t, server, dumpFuncErrC, []string{"foo"})
   156  }