├── .github └── workflows │ ├── ci.yml │ └── daily.yml ├── .gitignore ├── .gitmodules ├── LICENSE ├── README.md ├── assert ├── assert.go └── assert_test.go ├── cmdlog ├── cmdlog.go ├── cmdlog_test.go ├── example │ ├── example.png │ └── main.go ├── export_test.go └── testdata │ └── TestLogger │ ├── COLOR=.exp.json │ ├── COLOR=1.exp.json │ ├── WithPrefix.exp.json │ ├── multiline.exp.json │ └── tty.exp.json ├── diff ├── diff.go ├── diff_test.go ├── example.png └── testdata │ ├── TestTestData.exp.json │ └── TestTestData │ └── TESTDATA_ACCEPT.exp.json ├── go.mod ├── go.sum ├── go2 └── go2.go ├── make.sh ├── mapfs ├── mapfs.go └── mapfs_test.go ├── xbrowser └── xbrowser.go ├── xcontext ├── mutex.go ├── xcontext.go └── xcontext_test.go ├── xdefer ├── xdefer.go └── xdefer_test.go ├── xexec └── xexec.go ├── xhttp ├── err.go ├── log.go └── serve.go ├── xjson └── xjson.go ├── xmain ├── flag_helpers.go ├── opts.go ├── stdlib_exec.go ├── xmain.go ├── xmaintest.go └── xmaintest_test.go ├── xos ├── env.go ├── env_test.go └── xos.go ├── xrand ├── xrand.go └── xrand_test.go └── xterm └── xterm.go /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: ci 2 | on: [push, pull_request] 3 | concurrency: 4 | group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} 5 | cancel-in-progress: true 6 | 7 | jobs: 8 | ci: 9 | runs-on: ubuntu-latest 10 | steps: 11 | - uses: actions/checkout@v3 12 | - uses: actions/setup-go@v4 13 | with: 14 | go-version-file: ./go.mod 15 | cache: true 16 | - run: COLOR=1 ./make.sh 17 | env: 18 | GITHUB_TOKEN: ${{ secrets._GITHUB_TOKEN }} 19 | SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }} 20 | nofixups: 21 | runs-on: ubuntu-latest 22 | steps: 23 | - uses: actions/checkout@v3 24 | - run: git submodule update --init 25 | - run: COLOR=1 ./ci/bin/nofixups.sh 26 | env: 27 | GITHUB_TOKEN: ${{ secrets._GITHUB_TOKEN }} 28 | SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }} 29 | -------------------------------------------------------------------------------- /.github/workflows/daily.yml: -------------------------------------------------------------------------------- 1 | name: daily 2 | on: 3 | workflow_dispatch: 4 | schedule: 5 | - cron: '42 0 * * *' # daily at 00:42 6 | concurrency: 7 | group: ${{ github.workflow }} 8 | cancel-in-progress: true 9 | 10 | jobs: 11 | ci: 12 | runs-on: ubuntu-latest 13 | steps: 14 | - uses: actions/checkout@v3 15 | - uses: actions/setup-go@v3 16 | with: 17 | go-version-file: ./go.mod 18 | cache: true 19 | - run: COLOR=1 CI_FORCE=1 ./make.sh 20 | env: 21 | GITHUB_TOKEN: ${{ secrets._GITHUB_TOKEN }} 22 | SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }} 23 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.got.* 2 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "ci"] 2 | path = ci 3 | url = https://github.com/terrastruct/ci 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2022 Terrastruct Inc. info@terrastruct.com 2 | 3 | Permission to use, copy, modify, and distribute this software for any 4 | purpose with or without fee is hereby granted, provided that the above 5 | copyright notice and this permission notice appear in all copies. 6 | 7 | THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES 8 | WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF 9 | MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR 10 | ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES 11 | WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN 12 | ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF 13 | OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 14 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # util-go 2 | 3 | [![godoc](https://pkg.go.dev/badge/oss.terrastruct.com/util-go.svg)](https://pkg.go.dev/oss.terrastruct.com/util-go) 4 | [![ci](https://github.com/terrastruct/util-go/actions/workflows/ci.yml/badge.svg)](https://github.com/terrastruct/util-go/actions/workflows/ci.yml) 5 | [![daily](https://github.com/terrastruct/util-go/actions/workflows/daily.yml/badge.svg)](https://github.com/terrastruct/util-go/actions/workflows/daily.yml) 6 | [![license](https://img.shields.io/github/license/terrastruct/util-go?color=9cf)](./LICENSE) 7 | 8 | Terrastruct's general purpose go libraries. 9 | 10 | See https://pkg.go.dev/oss.terrastruct.com/util-go for docs. 11 | 12 | If there's enough external demand for a single package to be split off into its 13 | own repo from this collection we will. Feel free to open an issue to request. 14 | 15 | 16 | - [./diff](#diff) 17 | - [./assert](#assert) 18 | - [./xdefer](#xdefer) 19 | - [./cmdlog](#cmdlog) 20 | - [./xterm](#xterm) 21 | - [./xos](#xos) 22 | - [./xrand](#xrand) 23 | - [./xcontext](#xcontext) 24 | - [./xjson](#xjson) 25 | - [./go2](#go2) 26 | - [./xbrowser](#xbrowser) 27 | - [./xexec](#xexec) 28 | - [./xhttp](#xhttp) 29 | - [./xmain](#xmain) 30 | - [./mapfs](#mapfs) 31 | 32 | godoc is the canonical reference but we've provided this index as the godoc UI is frankly 33 | garbage after the move to pkg.go.dev. It's nowhere near as clear and responsive as the old 34 | UI. If this feedback reaches the authors of pkg.go.dev, please revert the UI back to what 35 | it was with godoc.org. 36 | 37 | ### [./diff](./diff) 38 | 39 | diff providers functions to diff strings, files and general Go values with git diff. 40 | 41 | ### [./assert](./assert) 42 | 43 | assert provides test assertion helpers. It integrates with [./diff](#diff) to display 44 | beautiful diffs. 45 | 46 | note: `TestdataJSON` is extremely useful. 47 | 48 | ![example output](./diff/example.png) 49 | 50 | - Strings 51 | - Files 52 | - Runes 53 | - JSON 54 | - Testdata 55 | - TestdataJSON 56 | 57 | ### [./xdefer](./xdefer) 58 | 59 | xdefer annotates all errors returned from a function transparently. 60 | 61 | ### [./cmdlog](./cmdlog) 62 | 63 | cmdlog implements color leveled logging for command line tools. 64 | 65 | ![example output](./cmdlog/example/example.png) 66 | 67 | `cmdlog` supports arbitrary randomly colored prefixes just like 68 | [terrastruct/ci](https://github.com/terrastruct/ci). 69 | 70 | Example is in [./cmdlog/example/main.go](./cmdlog/example/main.go). 71 | 72 | See [./cmdlog/cmdlog_test.go](./cmdlog/cmdlog_test.go) for further usage. 73 | 74 | You can log in tests with `NewTB`. 75 | 76 | - `$COLOR` is obeyed to force enable/disable colored output. 77 | - `$DEBUG` is obeyed to enable/disable debug logs. 78 | 79 | ### [./xterm](./xterm) 80 | 81 | xterm implements outputting formatted text to a terminal. 82 | 83 | ### [./xos](./xos) 84 | 85 | xos provides OS helpers. 86 | 87 | ### [./xrand](./xrand) 88 | 89 | xrand provides helpers for generating useful random values. 90 | We use it mainly for generating inputs to tests. 91 | 92 | ### [./xcontext](./xcontext) 93 | 94 | xcontext implements indispensable context helpers. 95 | 96 | ### [./xjson](./xjson) 97 | 98 | xjson implements basic JSON helpers. 99 | 100 | ### [./go2](./go2) 101 | 102 | go2 contains general utility helpers that should've been in Go. Maybe they'll be in Go 2.0. 103 | 104 | ### [./xbrowser](./xbrowser) 105 | 106 | xbrowser enables opening a user's GUI browser to a URL. 107 | 108 | ### [./xexec](./xexec) 109 | 110 | xexec provides exec helpers. 111 | 112 | ### [./xhttp](./xhttp) 113 | 114 | xhttp provides HTTP helpers. 115 | 116 | ### [./xmain](./xmain) 117 | 118 | xmain implements helpers for building CLI tools. 119 | 120 | ### [./mapfs](./mapfs) 121 | 122 | Package mapfs takes in a description of a filesystem as a `map[string]string` and writes it to a temp directory so that it may be used as an io/fs.FS. 123 | -------------------------------------------------------------------------------- /assert/assert.go: -------------------------------------------------------------------------------- 1 | // Package assert provides test assertion helpers. 2 | package assert 3 | 4 | import ( 5 | "io" 6 | "os" 7 | "path/filepath" 8 | "testing" 9 | 10 | "go.uber.org/multierr" 11 | 12 | "oss.terrastruct.com/util-go/diff" 13 | "oss.terrastruct.com/util-go/xjson" 14 | ) 15 | 16 | func Success(tb testing.TB, err error) { 17 | tb.Helper() 18 | if err != nil { 19 | tb.Fatalf("unexpected error: %v", err) 20 | } 21 | } 22 | 23 | func Error(tb testing.TB, err error) { 24 | tb.Helper() 25 | if err == nil { 26 | tb.Fatal("expected error") 27 | } 28 | } 29 | 30 | func ErrorString(tb testing.TB, err error, msg string) { 31 | tb.Helper() 32 | if err == nil { 33 | tb.Fatalf("expected error containing %q", msg) 34 | } 35 | String(tb, msg, err.Error()) 36 | } 37 | 38 | func StringJSON(tb testing.TB, exp string, got interface{}) { 39 | tb.Helper() 40 | String(tb, exp, string(xjson.Marshal(got))) 41 | } 42 | 43 | func String(tb testing.TB, exp, got string) { 44 | tb.Helper() 45 | diff, err := diff.Strings(exp, got) 46 | Success(tb, err) 47 | if diff != "" { 48 | tb.Fatalf("\n%s", diff) 49 | } 50 | } 51 | 52 | func JSON(tb testing.TB, exp, got interface{}) { 53 | tb.Helper() 54 | diff, err := diff.JSON(exp, got) 55 | Success(tb, err) 56 | if diff != "" { 57 | tb.Fatalf("\n%s", diff) 58 | } 59 | } 60 | 61 | func Runes(tb testing.TB, exp, got string) { 62 | tb.Helper() 63 | err := diff.Runes(exp, got) 64 | Success(tb, err) 65 | } 66 | 67 | func TestdataJSON(tb testing.TB, got interface{}) { 68 | tb.Helper() 69 | err := diff.TestdataJSON(filepath.Join("testdata", tb.Name()), got) 70 | Success(tb, err) 71 | } 72 | 73 | func Testdata(tb testing.TB, ext string, got []byte) { 74 | tb.Helper() 75 | err := diff.Testdata(filepath.Join("testdata", tb.Name()), ext, got) 76 | Success(tb, err) 77 | } 78 | 79 | func TestdataDir(tb testing.TB, dir string) { 80 | tb.Helper() 81 | err := diff.TestdataDir(filepath.Join("testdata", tb.Name()), dir) 82 | if err != nil { 83 | for _, err = range multierr.Errors(err) { 84 | tb.Error(err) 85 | } 86 | } 87 | if tb.Failed() { 88 | tb.FailNow() 89 | } 90 | } 91 | 92 | func Close(tb testing.TB, c io.Closer) { 93 | tb.Helper() 94 | err := c.Close() 95 | if err != nil { 96 | tb.Fatalf("failed to close %T: %v", c, err) 97 | } 98 | } 99 | 100 | func Equal(tb testing.TB, exp, got interface{}) { 101 | tb.Helper() 102 | if exp == got { 103 | return 104 | } 105 | exps, ok := exp.(string) 106 | if ok { 107 | gots, ok := got.(string) 108 | if ok { 109 | String(tb, exps, gots) 110 | return 111 | } 112 | } 113 | tb.Fatalf("expected %#v but got %#v", exp, got) 114 | } 115 | 116 | func NotEqual(tb testing.TB, v1, v2 interface{}) { 117 | tb.Helper() 118 | if v1 != v2 { 119 | return 120 | } 121 | tb.Fatalf("did not expect %#v", v2) 122 | } 123 | 124 | func True(tb testing.TB, v bool) { 125 | tb.Helper() 126 | Equal(tb, true, v) 127 | } 128 | 129 | func False(tb testing.TB, v bool) { 130 | tb.Helper() 131 | Equal(tb, false, v) 132 | } 133 | 134 | func TempDir(tb testing.TB) (dir string, cleanup func()) { 135 | tb.Helper() 136 | 137 | dir, err := os.MkdirTemp("", "util-go.assert.TempDir") 138 | Success(tb, err) 139 | return dir, func() { 140 | tb.Helper() 141 | err = os.RemoveAll(dir) 142 | Success(tb, err) 143 | } 144 | } 145 | 146 | func WriteFile(tb testing.TB, fp string, data []byte, perms os.FileMode) { 147 | tb.Helper() 148 | 149 | err := os.WriteFile(fp, data, perms) 150 | Success(tb, err) 151 | } 152 | 153 | func ReadFile(tb testing.TB, fp string) (data []byte) { 154 | tb.Helper() 155 | 156 | data, err := os.ReadFile(fp) 157 | Success(tb, err) 158 | return data 159 | } 160 | 161 | func Remove(tb testing.TB, fp string) { 162 | tb.Helper() 163 | 164 | err := os.Remove(fp) 165 | Success(tb, err) 166 | } 167 | 168 | func RemoveAll(tb testing.TB, fp string) { 169 | tb.Helper() 170 | 171 | err := os.RemoveAll(fp) 172 | Success(tb, err) 173 | } 174 | -------------------------------------------------------------------------------- /assert/assert_test.go: -------------------------------------------------------------------------------- 1 | package assert_test 2 | 3 | import ( 4 | "testing" 5 | 6 | "oss.terrastruct.com/util-go/assert" 7 | ) 8 | 9 | func TestStringJSON(t *testing.T) { 10 | t.Parallel() 11 | 12 | gen := func() (m1, m2 map[string]interface{}) { 13 | m1 = map[string]interface{}{ 14 | "one": 1, 15 | "two": 2, 16 | "three": 3, 17 | "four": 4, 18 | "five": map[string]interface{}{ 19 | "yes": "yes", 20 | "no": "yes", 21 | "five": map[string]interface{}{ 22 | "yes": "no", 23 | "no": "yes", 24 | }, 25 | }, 26 | } 27 | 28 | m2 = map[string]interface{}{ 29 | "one": 1, 30 | "two": 2, 31 | "three": 3, 32 | "four": 4, 33 | "five": map[string]interface{}{ 34 | "yes": "yes", 35 | "no": "yes", 36 | "five": map[string]interface{}{ 37 | "yes": "no", 38 | "no": "yes", 39 | }, 40 | }, 41 | } 42 | 43 | return m1, m2 44 | } 45 | 46 | t.Run("equal", func(t *testing.T) { 47 | t.Parallel() 48 | 49 | m1, m2 := gen() 50 | assert.JSON(t, m1, m2) 51 | }) 52 | 53 | t.Run("diff", func(t *testing.T) { 54 | t.Parallel() 55 | 56 | m1, m2 := gen() 57 | m2["five"].(map[string]interface{})["five"].(map[string]interface{})["no"] = "ys" 58 | 59 | fataledWithDiff := false 60 | ftb := &fakeTB{ 61 | TB: t, 62 | fatalf: func(f string, v ...interface{}) { 63 | t.Helper() 64 | if len(v) == 1 { 65 | t.Logf(f, v...) 66 | fataledWithDiff = true 67 | return 68 | } 69 | 70 | t.Fatalf(f, v...) 71 | }} 72 | 73 | defer func() { 74 | if t.Failed() || !fataledWithDiff { 75 | t.Error("expected assert.StringJSON to fatal with correct diff") 76 | } 77 | }() 78 | assert.JSON(ftb, m1, m2) 79 | }) 80 | } 81 | 82 | type fakeTB struct { 83 | fatalf func(string, ...interface{}) 84 | testing.TB 85 | } 86 | 87 | func (ftb *fakeTB) Fatalf(f string, v ...interface{}) { 88 | ftb.TB.Helper() 89 | ftb.fatalf(f, v...) 90 | } 91 | -------------------------------------------------------------------------------- /cmdlog/cmdlog.go: -------------------------------------------------------------------------------- 1 | // Package cmdlog implements color leveled logging for command line tools. 2 | package cmdlog 3 | 4 | import ( 5 | "bytes" 6 | "io" 7 | "log" 8 | "os" 9 | "sync" 10 | "sync/atomic" 11 | "testing" 12 | "time" 13 | 14 | "oss.terrastruct.com/util-go/xos" 15 | "oss.terrastruct.com/util-go/xterm" 16 | ) 17 | 18 | var timeNow = time.Now 19 | 20 | const defaultTSFormat = "15:04:05" 21 | 22 | func init() { 23 | l := New(xos.NewEnv(os.Environ()), os.Stderr) 24 | l.SetTS(true) 25 | l = l.WithPrefix(xterm.Blue, "stdlog") 26 | 27 | log.SetOutput(l.NoLevel.Writer()) 28 | log.SetPrefix(l.NoLevel.Prefix()) 29 | log.SetFlags(l.NoLevel.Flags()) 30 | } 31 | 32 | type Logger struct { 33 | env *xos.Env 34 | w io.Writer 35 | tsw *tsWriter 36 | dw *debugWriter 37 | 38 | NoLevel *log.Logger 39 | Debug *log.Logger 40 | Success *log.Logger 41 | Info *log.Logger 42 | Warn *log.Logger 43 | Error *log.Logger 44 | } 45 | 46 | func (l *Logger) GetTS() bool { 47 | l.tsw.mu.Lock() 48 | defer l.tsw.mu.Unlock() 49 | return l.tsw.enabled 50 | } 51 | 52 | func (l *Logger) GetTSFormat() string { 53 | l.tsw.mu.Lock() 54 | defer l.tsw.mu.Unlock() 55 | return l.tsw.tsfmt 56 | } 57 | 58 | func (l *Logger) GetDebug() bool { 59 | return l.dw.debug() 60 | } 61 | 62 | func (l *Logger) SetTS(enabled bool) { 63 | l.tsw.mu.Lock() 64 | l.tsw.enabled = enabled 65 | l.tsw.mu.Unlock() 66 | } 67 | 68 | func (l *Logger) SetTSFormat(tsfmt string) { 69 | l.tsw.mu.Lock() 70 | l.tsw.tsfmt = tsfmt 71 | l.tsw.mu.Unlock() 72 | } 73 | 74 | func (l *Logger) SetDebug(enabled bool) { 75 | vi := int64(0) 76 | if enabled { 77 | vi = 1 78 | } 79 | atomic.StoreInt64(&l.dw.flag, vi) 80 | } 81 | 82 | func New(env *xos.Env, w io.Writer) *Logger { 83 | tsw := &tsWriter{w: w, tsfmt: defaultTSFormat} 84 | dw := &debugWriter{w: tsw, env: env} 85 | l := &Logger{ 86 | env: env, 87 | w: w, 88 | dw: dw, 89 | tsw: tsw, 90 | } 91 | l.init("") 92 | return l 93 | } 94 | 95 | func (l *Logger) init(prefix string) { 96 | l.NoLevel = log.New(prefixWriter{l.tsw, prefix}, "", 0) 97 | 98 | if prefix != "" { 99 | prefix += " " 100 | } 101 | l.Debug = log.New(prefixWriter{l.dw, prefix + xterm.Prefix(l.env, l.w, "", "debug")}, "", 0) 102 | l.Success = log.New(prefixWriter{l.tsw, prefix + xterm.Prefix(l.env, l.w, xterm.Green, "success")}, "", 0) 103 | l.Info = log.New(prefixWriter{l.tsw, prefix + xterm.Prefix(l.env, l.w, xterm.Blue, "info")}, "", 0) 104 | l.Warn = log.New(prefixWriter{l.tsw, prefix + xterm.Prefix(l.env, l.w, xterm.Yellow, "warn")}, "", 0) 105 | l.Error = log.New(prefixWriter{l.tsw, prefix + xterm.Prefix(l.env, l.w, xterm.Red, "err")}, "", 0) 106 | } 107 | 108 | type prefixWriter struct { 109 | w io.Writer 110 | prefix string 111 | } 112 | 113 | func (pw prefixWriter) Write(p []byte) (int, error) { 114 | lines := bytes.Split(p, []byte("\n")) 115 | p2 := make([]byte, 0, (len(pw.prefix)+1)*len(lines)+len(p)) 116 | 117 | for _, l := range lines[:len(lines)-1] { 118 | prefix := pw.prefix 119 | if len(l) > 0 { 120 | prefix += " " 121 | } 122 | p2 = append(p2, prefix...) 123 | p2 = append(p2, l...) 124 | p2 = append(p2, '\n') 125 | } 126 | 127 | n, err := pw.w.Write(p2) 128 | if n > len(p) { 129 | n = len(p) 130 | } 131 | return n, err 132 | } 133 | 134 | type debugWriter struct { 135 | w io.Writer 136 | flag int64 137 | env *xos.Env 138 | } 139 | 140 | func (dw *debugWriter) debug() bool { 141 | if atomic.LoadInt64(&dw.flag) == 0 { 142 | return dw.env.Debug() 143 | } 144 | return true 145 | } 146 | 147 | func (dw *debugWriter) Write(p []byte) (int, error) { 148 | if !dw.debug() { 149 | return len(p), nil 150 | } 151 | return dw.w.Write(p) 152 | } 153 | 154 | type tsWriter struct { 155 | w io.Writer 156 | 157 | mu sync.Mutex 158 | tsfmt string 159 | enabled bool 160 | } 161 | 162 | func (tsw *tsWriter) Write(p []byte) (int, error) { 163 | tsw.mu.Lock() 164 | enabled := tsw.enabled 165 | tsfmt := tsw.tsfmt 166 | tsw.mu.Unlock() 167 | 168 | if !enabled { 169 | return tsw.w.Write(p) 170 | } 171 | 172 | ts := timeNow().Format(tsfmt) 173 | prefix := []byte("[" + ts + "]") 174 | 175 | lines := bytes.Split(p, []byte("\n")) 176 | p2 := make([]byte, 0, (len(prefix)+1)*len(lines)+len(p)) 177 | 178 | for _, l := range lines[:len(lines)-1] { 179 | prefix := prefix 180 | if len(l) > 0 { 181 | prefix = append(prefix, ' ') 182 | } 183 | p2 = append(p2, prefix...) 184 | p2 = append(p2, l...) 185 | p2 = append(p2, '\n') 186 | } 187 | 188 | n, err := tsw.w.Write(p2) 189 | if n > len(p) { 190 | n = len(p) 191 | } 192 | return n, err 193 | } 194 | 195 | func NewTB(env *xos.Env, tb testing.TB) *Logger { 196 | return New(env, tbWriter{tb}) 197 | } 198 | 199 | type tbWriter struct { 200 | tb testing.TB 201 | } 202 | 203 | func (tbw tbWriter) Write(p []byte) (int, error) { 204 | tbw.tb.Logf("%s", p) 205 | return len(p), nil 206 | } 207 | 208 | // Allows detection as a terminal. 209 | func (tbWriter) Fd() uintptr { 210 | return os.Stderr.Fd() 211 | } 212 | 213 | func (l *Logger) WithCCPrefix(s string) *Logger { 214 | return l.withPrefix(xterm.CCPrefix(l.env, l.w, s)) 215 | } 216 | 217 | func (l *Logger) WithPrefix(caps, s string) *Logger { 218 | return l.withPrefix(xterm.Prefix(l.env, l.w, caps, s)) 219 | } 220 | 221 | func (l *Logger) withPrefix(s string) *Logger { 222 | l2 := new(Logger) 223 | *l2 = *l 224 | 225 | prefix := l.NoLevel.Writer().(prefixWriter).prefix 226 | if len(s) > 0 { 227 | if len(prefix) > 0 { 228 | prefix += " " 229 | } 230 | prefix += s 231 | } 232 | l2.init(prefix) 233 | return l2 234 | } 235 | -------------------------------------------------------------------------------- /cmdlog/cmdlog_test.go: -------------------------------------------------------------------------------- 1 | package cmdlog_test 2 | 3 | import ( 4 | "bytes" 5 | "context" 6 | "fmt" 7 | "io" 8 | "testing" 9 | "time" 10 | 11 | "github.com/creack/pty" 12 | 13 | "oss.terrastruct.com/util-go/assert" 14 | "oss.terrastruct.com/util-go/cmdlog" 15 | "oss.terrastruct.com/util-go/xos" 16 | ) 17 | 18 | func TestLogger(t *testing.T) { 19 | t.Parallel() 20 | 21 | var tca = []struct { 22 | name string 23 | run func(t *testing.T, ctx context.Context, env *xos.Env) 24 | }{ 25 | { 26 | name: "COLOR=1", 27 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 28 | b := &bytes.Buffer{} 29 | env.Setenv("COLOR", "1") 30 | l := cmdlog.New(env, b) 31 | 32 | testLogger(l) 33 | 34 | t.Log(b.String()) 35 | assert.TestdataJSON(t, b.String()) 36 | }, 37 | }, 38 | { 39 | name: "COLOR=", 40 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 41 | b := &bytes.Buffer{} 42 | l := cmdlog.New(env, b) 43 | 44 | testLogger(l) 45 | 46 | t.Log(b.String()) 47 | assert.TestdataJSON(t, b.String()) 48 | }, 49 | }, 50 | { 51 | name: "tty", 52 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 53 | ptmx, tty, err := pty.Open() 54 | if err != nil { 55 | t.Fatalf("failed to open pty: %v", err) 56 | } 57 | defer assert.Close(t, ptmx) 58 | defer assert.Close(t, tty) 59 | 60 | l := cmdlog.New(env, tty) 61 | testLogger(l) 62 | 63 | timer := time.AfterFunc(time.Second*5, func() { 64 | // For some reason ptmx.SetDeadline() does not work. 65 | // tty has to be closed for a read on ptmx to unblock. 66 | tty.Close() 67 | t.Error("read took too long, update expLen") 68 | }) 69 | defer timer.Stop() 70 | // If the expected output changes, increase this to 9999, rerun and then update 71 | // for new length. 72 | const expLen = 415 73 | out, err := io.ReadAll(io.LimitReader(ptmx, expLen)) 74 | if err != nil { 75 | t.Fatalf("failed to read log output: %v", err) 76 | } 77 | t.Log(len(out), string(out)) 78 | assert.TestdataJSON(t, string(out)) 79 | }, 80 | }, 81 | { 82 | name: "testing.TB", 83 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 84 | ft := &fakeTB{ 85 | TB: t, 86 | logf: func(f string, v ...interface{}) { 87 | t.Helper() 88 | assert.String(t, "info: what's up\n", fmt.Sprintf(f, v...)) 89 | }, 90 | } 91 | 92 | env.Setenv("COLOR", "0") 93 | l := cmdlog.NewTB(env, ft) 94 | l.Info.Printf("what's up") 95 | }, 96 | }, 97 | { 98 | name: "WithPrefix", 99 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 100 | b := &bytes.Buffer{} 101 | env.Setenv("COLOR", "1") 102 | l := cmdlog.New(env, b) 103 | 104 | l2 := l.WithCCPrefix("lochness") 105 | if l2 == l { 106 | t.Fatalf("expected l and l2 to be different loggers") 107 | } 108 | l2 = l2.WithCCPrefix("imgbundler") 109 | l2 = l2.WithCCPrefix("cache") 110 | 111 | testLogger(l) 112 | testLogger(l2) 113 | 114 | t.Log(b.String()) 115 | assert.TestdataJSON(t, b.String()) 116 | }, 117 | }, 118 | { 119 | name: "multiline", 120 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 121 | b := &bytes.Buffer{} 122 | env.Setenv("COLOR", "1") 123 | l := cmdlog.New(env, b) 124 | 125 | l.NoLevel.Print("") 126 | l.SetTS(true) 127 | l.NoLevel.Print("") 128 | l.SetTS(false) 129 | 130 | l2 := l.WithCCPrefix("lochness") 131 | l2 = l2.WithCCPrefix("imgbundler") 132 | l2 = l2.WithCCPrefix("cache") 133 | 134 | l2.Warn.Print(``) 135 | l2.Warn.Print("\n\n\n") 136 | l2.SetTS(true) 137 | l2.Warn.Printf(`yes %d 138 | yes %d`, 3, 4) 139 | 140 | t.Log(b.String()) 141 | assert.TestdataJSON(t, b.String()) 142 | }, 143 | }, 144 | } 145 | 146 | ctx := context.Background() 147 | for _, tc := range tca { 148 | tc := tc 149 | t.Run(tc.name, func(t *testing.T) { 150 | t.Parallel() 151 | 152 | ctx, cancel := context.WithCancel(ctx) 153 | defer cancel() 154 | 155 | env := xos.NewEnv(nil) 156 | tc.run(t, ctx, env) 157 | }) 158 | } 159 | } 160 | 161 | func testLogger(l *cmdlog.Logger) { 162 | l.NoLevel.Println("Somehow, the world always affects you more than you affect it.") 163 | 164 | l.SetDebug(true) 165 | l.Debug.Println("Man is a rational animal who always loses his temper when he is called upon.") 166 | 167 | l.SetDebug(false) 168 | l.Debug.Println("You can never trust a woman; she may be true to you.") 169 | 170 | l.SetTS(true) 171 | l.Success.Println("An alcoholic is someone you don't like who drinks as much as you do.") 172 | l.Info.Println("There once was this swami who lived above a delicatessan.") 173 | 174 | l.SetTSFormat(time.UnixDate) 175 | l.Warn.Println("Telephone books are like dictionaries -- if you know the answer before.") 176 | 177 | l.SetTS(false) 178 | l.Error.Println("Nothing can be done in one trip.") 179 | l.Error.Println(`Good day to let down old friends who need help. 180 | I believe in getting into hot water; it keeps you clean.`) 181 | } 182 | 183 | type fakeTB struct { 184 | testing.TB 185 | logf func(string, ...interface{}) 186 | } 187 | 188 | func (ftb *fakeTB) Logf(f string, v ...interface{}) { 189 | ftb.TB.Helper() 190 | ftb.logf(f, v...) 191 | } 192 | -------------------------------------------------------------------------------- /cmdlog/example/example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/terrastruct/util-go/243d8661088abfd1cc6d1722615fed0b4800b133/cmdlog/example/example.png -------------------------------------------------------------------------------- /cmdlog/example/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "os" 5 | "time" 6 | 7 | "oss.terrastruct.com/util-go/cmdlog" 8 | "oss.terrastruct.com/util-go/xos" 9 | ) 10 | 11 | func main() { 12 | l := cmdlog.New(xos.NewEnv(os.Environ()), os.Stderr) 13 | l = l.WithCCPrefix("lochness") 14 | l = l.WithCCPrefix("imgbundler") 15 | l = l.WithCCPrefix("cache") 16 | 17 | l.NoLevel.Println("Somehow, the world always affects you more than you affect it.") 18 | 19 | l.SetDebug(true) 20 | l.Debug.Println("Man is a rational animal who always loses his temper when he is called upon.") 21 | 22 | l.SetDebug(false) 23 | l.Debug.Println("You can never trust a woman; she may be true to you.") 24 | 25 | l.SetTS(true) 26 | l.Success.Println("An alcoholic is someone you don't like who drinks as much as you do.") 27 | l.Info.Println("There once was this swami who lived above a delicatessan.") 28 | 29 | l.SetTSFormat(time.UnixDate) 30 | l.Warn.Println("Telephone books are like dictionaries -- if you know the answer before.") 31 | 32 | l.SetTS(false) 33 | l.Error.Println("Nothing can be done in one trip.") 34 | l.Error.Println(`Good day to let down old friends who need help. 35 | I believe in getting into hot water; it keeps you clean.`) 36 | } 37 | -------------------------------------------------------------------------------- /cmdlog/export_test.go: -------------------------------------------------------------------------------- 1 | package cmdlog 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "log" 7 | "os" 8 | "testing" 9 | "time" 10 | 11 | "oss.terrastruct.com/util-go/assert" 12 | "oss.terrastruct.com/util-go/xos" 13 | "oss.terrastruct.com/util-go/xterm" 14 | ) 15 | 16 | func init() { 17 | timeNow = func() time.Time { 18 | return time.Date(2000, time.January, 1, 1, 1, 1, 1, time.UTC) 19 | } 20 | } 21 | 22 | func TestStdlog(t *testing.T) { 23 | pw, ok := log.Default().Writer().(prefixWriter) 24 | if !ok { 25 | t.Fatalf("unexpected log.Default().Writer(): %T", log.Default().Writer()) 26 | } 27 | tsw, ok := pw.w.(*tsWriter) 28 | if !ok { 29 | t.Fatalf("unexpected pw.w: %T", pw.w) 30 | } 31 | b := &bytes.Buffer{} 32 | ow := tsw.w 33 | tsw.w = b 34 | defer func() { 35 | tsw.w = ow 36 | }() 37 | 38 | log.Print("testing stdlog") 39 | exp := fmt.Sprintf("[01:01:01] %s testing stdlog\n", 40 | xterm.Prefix(xos.NewEnv(os.Environ()), os.Stderr, xterm.Blue, "stdlog"), 41 | ) 42 | assert.String(t, exp, b.String()) 43 | } 44 | -------------------------------------------------------------------------------- /cmdlog/testdata/TestLogger/COLOR=.exp.json: -------------------------------------------------------------------------------- 1 | " Somehow, the world always affects you more than you affect it.\ndebug: Man is a rational animal who always loses his temper when he is called upon.\n[01:01:01] success: An alcoholic is someone you don't like who drinks as much as you do.\n[01:01:01] info: There once was this swami who lived above a delicatessan.\n[Sat Jan 1 01:01:01 UTC 2000] warn: Telephone books are like dictionaries -- if you know the answer before.\nerr: Nothing can be done in one trip.\nerr: Good day to let down old friends who need help.\nerr: I believe in getting into hot water; it keeps you clean.\n" 2 | -------------------------------------------------------------------------------- /cmdlog/testdata/TestLogger/COLOR=1.exp.json: -------------------------------------------------------------------------------- 1 | " Somehow, the world always affects you more than you affect it.\ndebug: Man is a rational animal who always loses his temper when he is called upon.\n[01:01:01] \u001b[32msuccess\u001b[0m: An alcoholic is someone you don't like who drinks as much as you do.\n[01:01:01] \u001b[34minfo\u001b[0m: There once was this swami who lived above a delicatessan.\n[Sat Jan 1 01:01:01 UTC 2000] \u001b[33mwarn\u001b[0m: Telephone books are like dictionaries -- if you know the answer before.\n\u001b[31merr\u001b[0m: Nothing can be done in one trip.\n\u001b[31merr\u001b[0m: Good day to let down old friends who need help.\n\u001b[31merr\u001b[0m: I believe in getting into hot water; it keeps you clean.\n" 2 | -------------------------------------------------------------------------------- /cmdlog/testdata/TestLogger/WithPrefix.exp.json: -------------------------------------------------------------------------------- 1 | " Somehow, the world always affects you more than you affect it.\ndebug: Man is a rational animal who always loses his temper when he is called upon.\n[01:01:01] \u001b[32msuccess\u001b[0m: An alcoholic is someone you don't like who drinks as much as you do.\n[01:01:01] \u001b[34minfo\u001b[0m: There once was this swami who lived above a delicatessan.\n[Sat Jan 1 01:01:01 UTC 2000] \u001b[33mwarn\u001b[0m: Telephone books are like dictionaries -- if you know the answer before.\n\u001b[31merr\u001b[0m: Nothing can be done in one trip.\n\u001b[31merr\u001b[0m: Good day to let down old friends who need help.\n\u001b[31merr\u001b[0m: I believe in getting into hot water; it keeps you clean.\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: Somehow, the world always affects you more than you affect it.\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: debug: Man is a rational animal who always loses his temper when he is called upon.\n[Sat Jan 1 01:01:01 UTC 2000] \u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[32msuccess\u001b[0m: An alcoholic is someone you don't like who drinks as much as you do.\n[Sat Jan 1 01:01:01 UTC 2000] \u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[34minfo\u001b[0m: There once was this swami who lived above a delicatessan.\n[Sat Jan 1 01:01:01 UTC 2000] \u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[33mwarn\u001b[0m: Telephone books are like dictionaries -- if you know the answer before.\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[31merr\u001b[0m: Nothing can be done in one trip.\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[31merr\u001b[0m: Good day to let down old friends who need help.\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[31merr\u001b[0m: I believe in getting into hot water; it keeps you clean.\n" 2 | -------------------------------------------------------------------------------- /cmdlog/testdata/TestLogger/multiline.exp.json: -------------------------------------------------------------------------------- 1 | "\n[01:01:01]\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[33mwarn\u001b[0m:\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[33mwarn\u001b[0m:\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[33mwarn\u001b[0m:\n\u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[33mwarn\u001b[0m:\n[01:01:01] \u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[33mwarn\u001b[0m: yes 3\n[01:01:01] \u001b[93mlochness\u001b[0m: \u001b[32mimgbundler\u001b[0m: \u001b[95mcache\u001b[0m: \u001b[33mwarn\u001b[0m: yes 4\n" 2 | -------------------------------------------------------------------------------- /cmdlog/testdata/TestLogger/tty.exp.json: -------------------------------------------------------------------------------- 1 | " Somehow, the world always affects you more than you affect it.\r\ndebug: Man is a rational animal who always loses his temper when he is called upon.\r\n[01:01:01] \u001b[32msuccess\u001b[0m: An alcoholic is someone you don't like who drinks as much as you do.\r\n[01:01:01] \u001b[34minfo\u001b[0m: There once was this swami who lived above a delicatessan.\r\n[Sat Jan 1 01:01:01 UTC 2000] \u001b[33mwarn\u001b[0m: Telephone books are like dictionari" 2 | -------------------------------------------------------------------------------- /diff/diff.go: -------------------------------------------------------------------------------- 1 | // package diff contains diff generation helpers, particularly useful for tests. 2 | // 3 | // - Strings 4 | // - Files 5 | // - Runes 6 | // - JSON 7 | // - Testdata 8 | // - TestdataJSON 9 | package diff 10 | 11 | import ( 12 | "context" 13 | "errors" 14 | "fmt" 15 | "io/ioutil" 16 | "os" 17 | "os/exec" 18 | "path/filepath" 19 | "strings" 20 | "time" 21 | 22 | "go.uber.org/multierr" 23 | 24 | "oss.terrastruct.com/util-go/xdefer" 25 | "oss.terrastruct.com/util-go/xjson" 26 | ) 27 | 28 | // Strings diffs exp with got in a git style diff. 29 | // 30 | // The git style diff header will contain real paths to exp and got 31 | // on the file system so that you can easily inspect them. 32 | // 33 | // This behavior is particularly useful for when you need to update 34 | // a test with the new got. You can just copy and paste from the got 35 | // file in the diff header. 36 | // 37 | // It uses Files under the hood. 38 | func Strings(exp, got string) (ds string, err error) { 39 | defer xdefer.Errorf(&err, "failed to diff text") 40 | 41 | if exp == got { 42 | return "", nil 43 | } 44 | 45 | d, err := ioutil.TempDir("", "ts_d2_diff") 46 | if err != nil { 47 | return "", err 48 | } 49 | 50 | expPath := filepath.Join(d, "exp") 51 | gotPath := filepath.Join(d, "got") 52 | 53 | err = ioutil.WriteFile(expPath, []byte(exp), 0644) 54 | if err != nil { 55 | return "", err 56 | } 57 | err = ioutil.WriteFile(gotPath, []byte(got), 0644) 58 | if err != nil { 59 | return "", err 60 | } 61 | 62 | return Files(expPath, gotPath) 63 | } 64 | 65 | // Files diffs expPath with gotPath and prints a git style diff header. 66 | // 67 | // It uses git under the hood. 68 | func Files(expPath, gotPath string) (ds string, err error) { 69 | defer xdefer.Errorf(&err, "failed to diff files") 70 | 71 | _, err = os.Stat(expPath) 72 | if os.IsNotExist(err) { 73 | expPath = "/dev/null" 74 | } 75 | _, err = os.Stat(gotPath) 76 | if os.IsNotExist(err) { 77 | gotPath = "/dev/null" 78 | } 79 | 80 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute) 81 | defer cancel() 82 | cmd := exec.CommandContext(ctx, "git", "-c", "diff.color=always", "diff", 83 | // Use the best diff-algorithm and highlight trailing whitespace. 84 | "--diff-algorithm=histogram", 85 | "--ws-error-highlight=all", 86 | "--no-index", 87 | expPath, gotPath) 88 | cmd.Env = append(cmd.Env, "GIT_CONFIG_NOSYSTEM=1", "HOME=") 89 | 90 | diffBytes, err := cmd.CombinedOutput() 91 | var ee *exec.ExitError 92 | if err != nil && !errors.As(err, &ee) { 93 | return "", fmt.Errorf("git diff failed: out=%q: %w", diffBytes, err) 94 | } 95 | ds = string(diffBytes) 96 | 97 | // Strips the diff header before --- 98 | // 99 | // diff --git a/var/folders/tf/yp9nqwbx4g5djqjxms03cvx80000gn/T/d2parser_test916758829/exp b/var/folders/tf/yp9nqwbx4g5djqjxms03cvx80000gn/T/d2parser_test916758829/got 100 | // index d48c704b..dbe709e6 100644 101 | // --- a/var/folders/tf/yp9nqwbx4g5djqjxms03cvx80000gn/T/d2parser_test916758829/exp 102 | // +++ b/var/folders/tf/yp9nqwbx4g5djqjxms03cvx80000gn/T/d2parser_test916758829/got 103 | // @@ -1,5 +1,5 @@ 104 | // 105 | // becomes: 106 | // 107 | // --- a/var/folders/tf/yp9nqwbx4g5djqjxms03cvx80000gn/T/d2parser_test916758829/exp 108 | // +++ b/var/folders/tf/yp9nqwbx4g5djqjxms03cvx80000gn/T/d2parser_test916758829/got 109 | // @@ -1,5 +1,5 @@ 110 | i := strings.Index(ds, "index") 111 | if i > -1 { 112 | j := strings.IndexByte(ds[i:], '\n') 113 | if j > -1 { 114 | ds = ds[i+j+1:] 115 | } 116 | } 117 | return strings.TrimSpace(ds), nil 118 | } 119 | 120 | // Runes is like Strings but formats exp and got with each unicode codepoint on a separate 121 | // line and generates a diff of that. It's useful for autogenerated UTF-8 with 122 | // xrand.String as Strings won't generate a coherent diff with undisplayable characters. 123 | func Runes(exp, got string) error { 124 | if exp == got { 125 | return nil 126 | } 127 | expRunes := formatRunes(exp) 128 | gotRunes := formatRunes(got) 129 | ds, err := Strings(expRunes, gotRunes) 130 | if err != nil { 131 | return err 132 | } 133 | if ds != "" { 134 | return errors.New(ds) 135 | } 136 | return nil 137 | } 138 | 139 | func formatRunes(s string) string { 140 | return strings.Join(strings.Split(fmt.Sprintf("%#v", []rune(s)), ", "), "\n") 141 | } 142 | 143 | // TestdataJSON is for when you have JSON that is too large to easily keep embedded by the 144 | // tests in _test.go files. As well, it makes the acceptance of large changes trivial 145 | // unlike say fs/embed. 146 | // 147 | // TestdataJSON encodes got as JSON and diffs it against the stored json in path.exp.json. 148 | // The got JSON is stored in path.got.json. If the diff is empty, it returns nil. 149 | // 150 | // Otherwise it returns an error containing the diff. 151 | // 152 | // In order to accept changes path.got.json has to become path.exp.json. You can use 153 | // ./ci/testdata/accept.sh to rename all non stale path.got.json files to path.exp.json. 154 | // 155 | // You can scope it to a single test or folder, see ./ci/testdata/accept.sh --help 156 | // 157 | // Also see ./ci/testdata/clean.sh --help for cleaning the repository of all 158 | // path.got.json and path.exp.json files. 159 | // 160 | // You can also use $TESTDATA_ACCEPT=1 to update all path.exp.json files on the fly. 161 | // This is useful when you're regenerating the repository's testdata. You can't easily 162 | // use the accept script without rerunning go test multiple times as go test will return 163 | // after too many test failures and will not continue until they are fixed. 164 | // 165 | // You'll want to use -count=1 to disable go test's result caching if you do use 166 | // $TESTDATA_ACCEPT. 167 | // 168 | // TestdataJSON will automatically create nonexistent directories in path. 169 | // 170 | // Here's an example that you can play with to better understand the behaviour: 171 | // 172 | // err = diff.TestdataJSON(filepath.Join("testdata", t.Name()), "change me") 173 | // if err != nil { 174 | // t.Fatal(err) 175 | // } 176 | // 177 | // Normally you want to use t.Name() as path for clarity but you can pass in any string. 178 | // e.g. a single test could persist two json objects into testdata with: 179 | // 180 | // err = diff.TestdataJSON(filepath.Join("testdata", t.Name(), "1"), "change me 1") 181 | // if err != nil { 182 | // t.Fatal(err) 183 | // } 184 | // err = diff.TestdataJSON(filepath.Join("testdata", t.Name(), "2"), "change me 2") 185 | // if err != nil { 186 | // t.Fatal(err) 187 | // } 188 | // 189 | // These would persist in testdata/${t.Name()}/1.exp.json and testdata/${t.Name()}/2.exp.json 190 | // 191 | // It uses Files under the hood. 192 | // 193 | // note: testdata is the canonical Go directory for such persistent test only files. 194 | // It is unfortunately poorly documented. See https://pkg.go.dev/cmd/go/internal/test 195 | // So normally you'd want path to be filepath.Join("testdata", t.Name()). 196 | // This is also the reason this function is named "TestdataJSON". 197 | func TestdataJSON(path string, got interface{}) error { 198 | gotb := xjson.Marshal(got) 199 | gotb = append(gotb, '\n') 200 | return Testdata(path, ".json", gotb) 201 | } 202 | 203 | // ext includes period like path.Ext() 204 | func Testdata(path, ext string, got []byte) error { 205 | expPath := fmt.Sprintf("%s.exp%s", path, ext) 206 | gotPath := fmt.Sprintf("%s.got%s", path, ext) 207 | 208 | err := os.MkdirAll(filepath.Dir(gotPath), 0755) 209 | if err != nil { 210 | return err 211 | } 212 | err = ioutil.WriteFile(gotPath, []byte(got), 0600) 213 | if err != nil { 214 | return err 215 | } 216 | 217 | ds, err := Files(expPath, gotPath) 218 | if err != nil { 219 | return err 220 | } 221 | 222 | if ds != "" { 223 | if os.Getenv("TESTDATA_ACCEPT") != "" || os.Getenv("TA") != "" { 224 | return os.Rename(gotPath, expPath) 225 | } 226 | if os.Getenv("NO_DIFF") != "" || os.Getenv("ND") != "" { 227 | ds = "diff hidden with $NO_DIFF=1 or $ND=1" 228 | } 229 | return fmt.Errorf("diff (rerun with $TESTDATA_ACCEPT=1 or $TA=1 to accept):\n%s", ds) 230 | } 231 | return os.Remove(gotPath) 232 | } 233 | 234 | func JSON(exp, got interface{}) (string, error) { 235 | return Strings(string(xjson.Marshal(exp)), string(xjson.Marshal(got))) 236 | } 237 | 238 | func TestdataDir(testName, dir string) (err error) { 239 | defer xdefer.Errorf(&err, "failed to commit testdata dir %v", dir) 240 | testdataDir(&err, testName, dir) 241 | return err 242 | } 243 | 244 | func testdataDir(errs *error, testName, dir string) { 245 | ea, err := os.ReadDir(dir) 246 | if err != nil { 247 | *errs = multierr.Combine(*errs, err) 248 | return 249 | } 250 | 251 | for _, e := range ea { 252 | if e.IsDir() { 253 | testdataDir(errs, filepath.Join(testName, e.Name()), filepath.Join(dir, e.Name())) 254 | } else { 255 | ext := filepath.Ext(e.Name()) 256 | name := strings.TrimSuffix(e.Name(), ext) 257 | got, err := os.ReadFile(filepath.Join(dir, e.Name())) 258 | if err != nil { 259 | *errs = multierr.Combine(*errs, err) 260 | continue 261 | } 262 | err = Testdata(filepath.Join(testName, name), ext, got) 263 | if err != nil { 264 | *errs = multierr.Combine(*errs, err) 265 | } 266 | } 267 | } 268 | } 269 | -------------------------------------------------------------------------------- /diff/diff_test.go: -------------------------------------------------------------------------------- 1 | package diff_test 2 | 3 | import ( 4 | "os" 5 | "path/filepath" 6 | "testing" 7 | 8 | "oss.terrastruct.com/util-go/assert" 9 | "oss.terrastruct.com/util-go/diff" 10 | ) 11 | 12 | func TestTestData(t *testing.T) { 13 | t.Run("TESTDATA_ACCEPT", testTestDataAccept) 14 | 15 | os.Unsetenv("TESTDATA_ACCEPT") 16 | os.Unsetenv("TA") 17 | 18 | m1 := map[string]interface{}{ 19 | "one": 1, 20 | "two": 2, 21 | "three": 3, 22 | "four": 4, 23 | "five": map[string]interface{}{ 24 | "yes": "yes", 25 | "no": "yes", 26 | "five": map[string]interface{}{ 27 | "yes": "no", 28 | "no": "yes", 29 | }, 30 | }, 31 | } 32 | 33 | err := os.Remove("testdata/TestTestData.exp.json") 34 | if err != nil && !os.IsNotExist(err) { 35 | t.Fatalf("unexpected error: %v", err) 36 | } 37 | 38 | err = diff.TestdataJSON(filepath.Join("testdata", t.Name()), m1) 39 | assert.Error(t, err) 40 | exp := `diff (rerun with $TESTDATA_ACCEPT=1 or $TA=1 to accept): 41 | --- /dev/null 42 | +++ b/testdata/TestTestData.got.json 43 | @@ -0,0 +1,14 @@ 44 | +{ 45 | + "five": { 46 | + "five": { 47 | + "no": "yes", 48 | + "yes": "no" 49 | + }, 50 | + "no": "yes", 51 | + "yes": "yes" 52 | + }, 53 | + "four": 4, 54 | + "one": 1, 55 | + "three": 3, 56 | + "two": 2 57 | +}` 58 | got := err.Error() 59 | ds, err := diff.Strings(exp, got) 60 | if err != nil { 61 | t.Fatalf("unable to generate exp diff: %v", err) 62 | } 63 | if ds != "" { 64 | t.Fatalf("expected no diff:\n%s", ds) 65 | } 66 | err = diff.Runes(exp, got) 67 | if err != nil { 68 | t.Fatalf("expected no rune diff: %v", err) 69 | } 70 | 71 | err = os.Rename("testdata/TestTestData.got.json", "testdata/TestTestData.exp.json") 72 | assert.Success(t, err) 73 | 74 | m1["five"].(map[string]interface{})["five"].(map[string]interface{})["no"] = "ys" 75 | 76 | err = diff.TestdataJSON(filepath.Join("testdata", t.Name()), m1) 77 | if err == nil { 78 | t.Fatalf("expected err: %#v", err) 79 | } 80 | exp = `diff (rerun with $TESTDATA_ACCEPT=1 or $TA=1 to accept): 81 | --- a/testdata/TestTestData.exp.json 82 | +++ b/testdata/TestTestData.got.json 83 | @@ -1,7 +1,7 @@ 84 | { 85 |  "five": { 86 |  "five": { 87 | - "no": "yes", 88 | + "no": "ys", 89 |  "yes": "no" 90 |  }, 91 |  "no": "yes",` 92 | got = err.Error() 93 | ds, err = diff.Strings(exp, got) 94 | assert.Success(t, err) 95 | if ds != "" { 96 | t.Fatalf("expected no diff:\n%s", ds) 97 | } 98 | 99 | exp += "a" 100 | ds, err = diff.Strings(exp, got) 101 | assert.Success(t, err) 102 | if ds == "" { 103 | t.Fatalf("expected incorrect diff:\n%s", ds) 104 | } 105 | err = diff.Runes(exp, got) 106 | assert.Error(t, err) 107 | } 108 | 109 | func testTestDataAccept(t *testing.T) { 110 | m1 := map[string]interface{}{ 111 | "one": 1, 112 | } 113 | 114 | os.Setenv("TESTDATA_ACCEPT", "1") 115 | os.Setenv("TA", "1") 116 | err := diff.TestdataJSON(filepath.Join("testdata", t.Name()), m1) 117 | assert.Success(t, err) 118 | 119 | m1["one"] = 2 120 | 121 | os.Setenv("TESTDATA_ACCEPT", "") 122 | os.Setenv("TA", "") 123 | err = diff.TestdataJSON(filepath.Join("testdata", t.Name()), m1) 124 | assert.Error(t, err) 125 | exp := `diff (rerun with $TESTDATA_ACCEPT=1 or $TA=1 to accept): 126 | --- a/testdata/TestTestData/TESTDATA_ACCEPT.exp.json 127 | +++ b/testdata/TestTestData/TESTDATA_ACCEPT.got.json 128 | @@ -1,3 +1,3 @@ 129 | { 130 | - "one": 1 131 | + "one": 2 132 | }` 133 | ds, err := diff.Strings(exp, err.Error()) 134 | assert.Success(t, err) 135 | if ds != "" { 136 | t.Fatalf("expected no diff:\n%s", ds) 137 | } 138 | 139 | err = os.Remove(filepath.Join("testdata", t.Name()) + ".got.json") 140 | assert.Success(t, err) 141 | } 142 | -------------------------------------------------------------------------------- /diff/example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/terrastruct/util-go/243d8661088abfd1cc6d1722615fed0b4800b133/diff/example.png -------------------------------------------------------------------------------- /diff/testdata/TestTestData.exp.json: -------------------------------------------------------------------------------- 1 | { 2 | "five": { 3 | "five": { 4 | "no": "yes", 5 | "yes": "no" 6 | }, 7 | "no": "yes", 8 | "yes": "yes" 9 | }, 10 | "four": 4, 11 | "one": 1, 12 | "three": 3, 13 | "two": 2 14 | } 15 | -------------------------------------------------------------------------------- /diff/testdata/TestTestData/TESTDATA_ACCEPT.exp.json: -------------------------------------------------------------------------------- 1 | { 2 | "one": 1 3 | } 4 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module oss.terrastruct.com/util-go 2 | 3 | go 1.18 4 | 5 | require ( 6 | github.com/creack/pty v1.1.18 7 | github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 8 | github.com/spf13/pflag v1.0.5 9 | go.uber.org/multierr v1.9.0 10 | golang.org/x/exp v0.0.0-20221126150942-6ab00d035af9 11 | golang.org/x/term v0.2.0 12 | golang.org/x/text v0.4.0 13 | golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 14 | ) 15 | 16 | require ( 17 | go.uber.org/atomic v1.7.0 // indirect 18 | golang.org/x/sys v0.2.0 // indirect 19 | ) 20 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/creack/pty v1.1.18 h1:n56/Zwd5o6whRC5PMGretI4IdRLlmBXYNjScPaBgsbY= 2 | github.com/creack/pty v1.1.18/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4= 3 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 4 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 5 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 6 | github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 h1:KoWmjvw+nsYOo29YJK9vDA65RGE3NrOnUtO7a+RF9HU= 7 | github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8/go.mod h1:HKlIX3XHQyzLZPlr7++PzdhaXEj94dEiJgZDTsxEqUI= 8 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 9 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 10 | github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= 11 | github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= 12 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 13 | github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= 14 | github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= 15 | go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw= 16 | go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= 17 | go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI= 18 | go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ= 19 | golang.org/x/exp v0.0.0-20221126150942-6ab00d035af9 h1:yZNXmy+j/JpX19vZkVktWqAo7Gny4PBWYYK3zskGpx4= 20 | golang.org/x/exp v0.0.0-20221126150942-6ab00d035af9/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= 21 | golang.org/x/sys v0.0.0-20210616045830-e2b7044e8c71/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 22 | golang.org/x/sys v0.2.0 h1:ljd4t30dBnAvMZaQCevtY0xLLD0A+bRZXbgLMLU1F/A= 23 | golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 24 | golang.org/x/term v0.2.0 h1:z85xZCsEl7bi/KwbNADeBYoOP0++7W1ipu+aGnpwzRM= 25 | golang.org/x/term v0.2.0/go.mod h1:TVmDHMZPmdnySmBfhjOoOdhjzdE1h4u1VwSiw2l1Nuc= 26 | golang.org/x/text v0.4.0 h1:BrVqGRd7+k1DiOgtnFvAkoQEWQvBc25ouMJM6429SFg= 27 | golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= 28 | golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 h1:H2TDz8ibqkAF6YGhCdN3jS9O0/s90v0rJh3X/OLHEUk= 29 | golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2/go.mod h1:K8+ghG5WaK9qNqU5K3HdILfMLy1f3aNYFI/wnl100a8= 30 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 31 | -------------------------------------------------------------------------------- /go2/go2.go: -------------------------------------------------------------------------------- 1 | // Package go2 contains general utility helpers that should've been in Go. Maybe they'll be in Go 2.0. 2 | package go2 3 | 4 | import ( 5 | "hash/fnv" 6 | "math" 7 | 8 | "golang.org/x/exp/constraints" 9 | ) 10 | 11 | func Pointer[T any](v T) *T { 12 | return &v 13 | } 14 | 15 | func Min[T constraints.Ordered](a, b T) T { 16 | if a < b { 17 | return a 18 | } 19 | return b 20 | } 21 | 22 | func Max[T constraints.Ordered](a, b T) T { 23 | if a > b { 24 | return a 25 | } 26 | return b 27 | } 28 | 29 | func StringToIntHash(s string) int { 30 | h := fnv.New32a() 31 | h.Write([]byte(s)) 32 | return int(h.Sum32()) 33 | } 34 | 35 | func Contains[T comparable](els []T, el T) bool { 36 | for _, el2 := range els { 37 | if el2 == el { 38 | return true 39 | } 40 | } 41 | return false 42 | } 43 | 44 | func Filter[T any](els []T, fn func(T) bool) []T { 45 | out := []T{} 46 | for _, el := range els { 47 | if fn(el) { 48 | out = append(out, el) 49 | } 50 | } 51 | return out 52 | } 53 | 54 | func IntMax(x, y int) int { 55 | return int(math.Max(float64(x), float64(y))) 56 | } 57 | 58 | func IntMin(x, y int) int { 59 | return int(math.Min(float64(x), float64(y))) 60 | } 61 | -------------------------------------------------------------------------------- /make.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | set -eu 3 | if [ ! -e "$(dirname "$0")/ci/.git" ]; then 4 | set -x 5 | git submodule update --init 6 | set +x 7 | fi 8 | . "$(dirname "$0")/ci/lib.sh" 9 | cd "$(dirname "$0")" 10 | 11 | job_parseflags "$@" 12 | runjob fmt ./ci/bin/fmt.sh & 13 | runjob lint ci_go_lint & 14 | runjob build 'go build ./...' & 15 | runjob test 'go test ./...' & 16 | ci_waitjobs 17 | -------------------------------------------------------------------------------- /mapfs/mapfs.go: -------------------------------------------------------------------------------- 1 | // Package mapfs takes in a description of a filesystem as a map[string]string 2 | // and writes it to a temp directory so that it may be used as an io/fs.FS. 3 | package mapfs 4 | 5 | import ( 6 | "errors" 7 | "fmt" 8 | "io/fs" 9 | "os" 10 | "path" 11 | ) 12 | 13 | type FS struct { 14 | dir string 15 | fs.FS 16 | } 17 | 18 | func New(m map[string]string) (*FS, error) { 19 | tempDir, err := os.MkdirTemp("", "mapfs-*") 20 | if err != nil { 21 | return nil, fmt.Errorf("failed to create root mapfs dir: %w", err) 22 | } 23 | for p, s := range m { 24 | p = path.Join(tempDir, p) 25 | err = os.MkdirAll(path.Dir(p), 0755) 26 | if err != nil { 27 | return nil, fmt.Errorf("failed to create mapfs dir %q: %w", path.Dir(p), err) 28 | } 29 | err = os.WriteFile(p, []byte(s), 0644) 30 | if err != nil { 31 | return nil, fmt.Errorf("failed to write mapfs file %q: %w", p, err) 32 | } 33 | } 34 | return &FS{ 35 | dir: tempDir, 36 | FS: os.DirFS(tempDir), 37 | }, nil 38 | } 39 | 40 | func (fs *FS) Close() error { 41 | err := os.RemoveAll(fs.dir) 42 | if errors.Is(err, os.ErrNotExist) { 43 | return nil 44 | } 45 | if err != nil { 46 | return fmt.Errorf("failed to close mapfs.FS: %w", err) 47 | } 48 | return nil 49 | } 50 | -------------------------------------------------------------------------------- /mapfs/mapfs_test.go: -------------------------------------------------------------------------------- 1 | package mapfs_test 2 | 3 | import ( 4 | "io/fs" 5 | "testing" 6 | 7 | "oss.terrastruct.com/util-go/assert" 8 | "oss.terrastruct.com/util-go/mapfs" 9 | ) 10 | 11 | func TestMapFS(t *testing.T) { 12 | t.Parallel() 13 | 14 | m := map[string]string{ 15 | "index": " I installed 'Linux 6.1', doesn't that make me a unix", 16 | "d2/imports": "Do your part to help preserve life on Earth -- by trying to preserve your own.", 17 | "d2/globs": "I'm going to raise an issue and stick it in your ear.", 18 | "nested/nested/nested/nested": "Yuppie Wannabes", 19 | } 20 | 21 | mapfs, err := mapfs.New(m) 22 | assert.Success(t, err) 23 | t.Cleanup(func() { 24 | err := mapfs.Close() 25 | assert.Success(t, err) 26 | }) 27 | 28 | for p, s := range m { 29 | b, err := fs.ReadFile(mapfs, p) 30 | assert.Success(t, err) 31 | assert.Equal(t, s, string(b)) 32 | } 33 | 34 | _, err = fs.ReadFile(mapfs, "../escape") 35 | assert.ErrorString(t, err, "stat ../escape: invalid argument") 36 | _, err = fs.ReadFile(mapfs, "/root") 37 | assert.ErrorString(t, err, "stat /root: invalid argument") 38 | } 39 | -------------------------------------------------------------------------------- /xbrowser/xbrowser.go: -------------------------------------------------------------------------------- 1 | package xbrowser 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "os/exec" 7 | 8 | "github.com/pkg/browser" 9 | 10 | "oss.terrastruct.com/util-go/xos" 11 | ) 12 | 13 | func Open(ctx context.Context, env *xos.Env, url string) error { 14 | browserEnv := env.Getenv("BROWSER") 15 | if browserEnv == "0" || browserEnv == "false" { 16 | return nil 17 | } 18 | if browserEnv != "" && browserEnv != "1" && browserEnv != "true" { 19 | browserSh := fmt.Sprintf(`%s "$1"`, browserEnv) 20 | cmd := exec.CommandContext(ctx, "sh", "-sc", browserSh, "--", url) 21 | out, err := cmd.CombinedOutput() 22 | if err != nil { 23 | return fmt.Errorf("failed to run %v (out: %q): %w", cmd.Args, out, err) 24 | } 25 | return nil 26 | } 27 | return browser.OpenURL(url) 28 | } 29 | -------------------------------------------------------------------------------- /xcontext/mutex.go: -------------------------------------------------------------------------------- 1 | package xcontext 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | ) 7 | 8 | type Mutex struct { 9 | ch chan struct{} 10 | } 11 | 12 | func NewMutex() *Mutex { 13 | return &Mutex{ 14 | ch: make(chan struct{}, 1), 15 | } 16 | } 17 | 18 | func (m *Mutex) TryLock() bool { 19 | select { 20 | case m.ch <- struct{}{}: 21 | return true 22 | default: 23 | return false 24 | } 25 | } 26 | 27 | func (m *Mutex) Lock(ctx context.Context) error { 28 | select { 29 | case <-ctx.Done(): 30 | return fmt.Errorf("failed to acquire lock: %w", ctx.Err()) 31 | case m.ch <- struct{}{}: 32 | return nil 33 | } 34 | } 35 | 36 | func (m *Mutex) Unlock() { 37 | select { 38 | case <-m.ch: 39 | default: 40 | panic("xcontext.Mutex: Unlock before Lock") 41 | } 42 | } 43 | -------------------------------------------------------------------------------- /xcontext/xcontext.go: -------------------------------------------------------------------------------- 1 | // Package xcontext implements indispensable context helpers. 2 | package xcontext 3 | 4 | import ( 5 | "context" 6 | "time" 7 | ) 8 | 9 | // WithoutCancel returns a context derived from ctx that may 10 | // never be cancelled. 11 | func WithoutCancel(ctx context.Context) context.Context { 12 | return withoutCancel{ctx: ctx} 13 | } 14 | 15 | type withoutCancel struct { 16 | ctx context.Context 17 | } 18 | 19 | func (c withoutCancel) Deadline() (time.Time, bool) { 20 | return time.Time{}, false 21 | } 22 | 23 | func (c withoutCancel) Done() <-chan struct{} { 24 | return nil 25 | } 26 | 27 | func (c withoutCancel) Err() error { 28 | return nil 29 | } 30 | 31 | func (c withoutCancel) Value(key interface{}) interface{} { 32 | return c.ctx.Value(key) 33 | } 34 | 35 | // WithoutValues creates a new context derived from ctx that does not inherit its values 36 | // but does pass on cancellation. 37 | func WithoutValues(ctx context.Context) context.Context { 38 | return withoutValues{ctx: ctx} 39 | } 40 | 41 | type withoutValues struct { 42 | ctx context.Context 43 | } 44 | 45 | func (c withoutValues) Deadline() (time.Time, bool) { 46 | return c.ctx.Deadline() 47 | } 48 | 49 | func (c withoutValues) Done() <-chan struct{} { 50 | return c.ctx.Done() 51 | } 52 | 53 | func (c withoutValues) Err() error { 54 | return c.ctx.Err() 55 | } 56 | 57 | func (c withoutValues) Value(key interface{}) interface{} { 58 | return nil 59 | } 60 | -------------------------------------------------------------------------------- /xcontext/xcontext_test.go: -------------------------------------------------------------------------------- 1 | package xcontext_test 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "oss.terrastruct.com/util-go/assert" 8 | "oss.terrastruct.com/util-go/xcontext" 9 | ) 10 | 11 | func TestWithoutCancel(t *testing.T) { 12 | t.Parallel() 13 | 14 | s := "meow" 15 | 16 | ctx := context.Background() 17 | ctx = stringWith(ctx, s) 18 | assert.Success(t, ctx.Err()) 19 | 20 | t.Run("no_cancel", func(t *testing.T) { 21 | t.Parallel() 22 | 23 | ctx := xcontext.WithoutCancel(ctx) 24 | assert.Success(t, ctx.Err()) 25 | 26 | s2 := stringFrom(ctx) 27 | assert.String(t, s, s2) 28 | }) 29 | 30 | t.Run("cancel_before", func(t *testing.T) { 31 | t.Parallel() 32 | 33 | ctx, cancel := context.WithCancel(ctx) 34 | cancel() 35 | 36 | assert.Error(t, ctx.Err()) 37 | 38 | ctx = xcontext.WithoutCancel(ctx) 39 | assert.Success(t, ctx.Err()) 40 | 41 | s2 := stringFrom(ctx) 42 | assert.String(t, s, s2) 43 | }) 44 | 45 | t.Run("cancel_after", func(t *testing.T) { 46 | t.Parallel() 47 | 48 | ctx, cancel := context.WithCancel(ctx) 49 | ctx = xcontext.WithoutCancel(ctx) 50 | cancel() 51 | assert.Success(t, ctx.Err()) 52 | 53 | s2 := stringFrom(ctx) 54 | assert.String(t, s, s2) 55 | }) 56 | } 57 | 58 | func TestWithoutValues(t *testing.T) { 59 | t.Parallel() 60 | 61 | ctx := context.Background() 62 | 63 | const k = "Death is nature's way of saying `Howdy'." 64 | const exp = "Proposed Additions to the PDP-11 Instruction Set" 65 | const exp2 = "character density, n.:" 66 | 67 | t.Run("no_value", func(t *testing.T) { 68 | t.Parallel() 69 | 70 | v := ctx.Value(k) 71 | assert.JSON(t, nil, v) 72 | 73 | ctx := xcontext.WithoutValues(ctx) 74 | 75 | v = ctx.Value(k) 76 | assert.JSON(t, nil, v) 77 | }) 78 | 79 | t.Run("with_value", func(t *testing.T) { 80 | t.Parallel() 81 | 82 | ctxv := context.WithValue(ctx, k, exp) 83 | ctx := xcontext.WithoutValues(ctxv) 84 | 85 | // ctxv contains k but ctx doesn't. 86 | v := ctxv.Value(k) 87 | assert.JSON(t, exp, v) 88 | 89 | v = ctx.Value(k) 90 | assert.JSON(t, nil, v) 91 | 92 | ctx = context.WithValue(ctx, k, exp2) 93 | v = ctx.Value(k) 94 | assert.JSON(t, exp2, v) 95 | }) 96 | 97 | t.Run("cancel", func(t *testing.T) { 98 | t.Parallel() 99 | 100 | t.Run("before", func(t *testing.T) { 101 | t.Parallel() 102 | 103 | ctx, cancel := context.WithCancel(ctx) 104 | cancel() 105 | 106 | ctx = xcontext.WithoutValues(ctx) 107 | assert.Error(t, ctx.Err()) 108 | }) 109 | 110 | t.Run("after", func(t *testing.T) { 111 | t.Parallel() 112 | 113 | ctx, cancel := context.WithCancel(ctx) 114 | defer cancel() 115 | 116 | ctx = xcontext.WithoutValues(ctx) 117 | assert.Success(t, ctx.Err()) 118 | 119 | cancel() 120 | assert.Error(t, ctx.Err()) 121 | }) 122 | }) 123 | } 124 | 125 | type stringKey struct{} 126 | 127 | func stringFrom(ctx context.Context) string { 128 | return ctx.Value(stringKey{}).(string) 129 | } 130 | 131 | func stringWith(ctx context.Context, s string) context.Context { 132 | return context.WithValue(ctx, stringKey{}, s) 133 | } 134 | -------------------------------------------------------------------------------- /xdefer/xdefer.go: -------------------------------------------------------------------------------- 1 | // Package xdefer implements an extremely useful function, Errorf, to annotate all errors returned from a function transparently. 2 | package xdefer 3 | 4 | import ( 5 | "fmt" 6 | "strings" 7 | 8 | "golang.org/x/xerrors" 9 | ) 10 | 11 | type deferError struct { 12 | s string 13 | err error 14 | frame xerrors.Frame 15 | } 16 | 17 | var _ interface { 18 | xerrors.Wrapper 19 | xerrors.Formatter 20 | Is(error) bool 21 | } = deferError{} 22 | 23 | func (e deferError) Unwrap() error { 24 | return e.err 25 | } 26 | 27 | func (e deferError) Format(f fmt.State, c rune) { 28 | xerrors.FormatError(e, f, c) 29 | } 30 | 31 | // Used to detect if there is a duplicate frame as a result 32 | // of using xdefer and if so to ignore it. 33 | type fakeXerrorsPrinter struct { 34 | s []string 35 | } 36 | 37 | func (fp *fakeXerrorsPrinter) Print(v ...interface{}) { 38 | fp.s = append(fp.s, fmt.Sprint(v...)) 39 | } 40 | 41 | func (fp *fakeXerrorsPrinter) Printf(f string, v ...interface{}) { 42 | fp.s = append(fp.s, fmt.Sprintf(f, v...)) 43 | } 44 | 45 | func (fp *fakeXerrorsPrinter) Detail() bool { 46 | return true 47 | } 48 | 49 | func (e deferError) shouldPrintFrame(p xerrors.Printer) bool { 50 | fm, ok := e.err.(xerrors.Formatter) 51 | if !ok { 52 | return true 53 | } 54 | 55 | fp := &fakeXerrorsPrinter{} 56 | e.frame.Format(fp) 57 | fp2 := &fakeXerrorsPrinter{} 58 | _ = fm.FormatError(fp2) 59 | if len(fp.s) >= 2 && len(fp2.s) >= 3 { 60 | if fp.s[1] == fp2.s[2] { 61 | // We don't need to print our frame into the real 62 | // xerrors printer as the next error will have it. 63 | return false 64 | } 65 | } 66 | return true 67 | } 68 | 69 | func (e deferError) FormatError(p xerrors.Printer) error { 70 | if e.s == "" { 71 | if e.shouldPrintFrame(p) { 72 | e.frame.Format(p) 73 | } 74 | return e.err 75 | } 76 | 77 | p.Print(e.s) 78 | if p.Detail() && e.shouldPrintFrame(p) { 79 | e.frame.Format(p) 80 | } 81 | return e.err 82 | } 83 | 84 | func (e deferError) Is(err error) bool { 85 | return xerrors.Is(e.err, err) 86 | } 87 | 88 | func (e deferError) Error() string { 89 | if e.s == "" { 90 | fp := &fakeXerrorsPrinter{} 91 | e.frame.Format(fp) 92 | if len(fp.s) < 1 { 93 | return e.err.Error() 94 | } 95 | return fmt.Sprintf("%v: %v", strings.TrimSpace(fp.s[0]), e.err) 96 | } 97 | return fmt.Sprintf("%v: %v", e.s, e.err) 98 | } 99 | 100 | // Errorf makes it easy to defer annotate an error for all return paths in a function. 101 | // See the tests for how it's used. 102 | // 103 | // Pass s == "" to only annotate the location of the return. 104 | func Errorf(err *error, s string, v ...interface{}) { 105 | if *err != nil { 106 | *err = deferError{ 107 | s: fmt.Sprintf(s, v...), 108 | err: *err, 109 | frame: xerrors.Caller(1), 110 | } 111 | } 112 | } 113 | -------------------------------------------------------------------------------- /xdefer/xdefer_test.go: -------------------------------------------------------------------------------- 1 | package xdefer_test 2 | 3 | import ( 4 | "fmt" 5 | "runtime" 6 | "testing" 7 | 8 | "golang.org/x/xerrors" 9 | 10 | "oss.terrastruct.com/util-go/assert" 11 | "oss.terrastruct.com/util-go/xdefer" 12 | ) 13 | 14 | func TestErrorf(t *testing.T) { 15 | t.Parallel() 16 | 17 | err := func() (err error) { 18 | defer xdefer.Errorf(&err, "second wrap %#v", []int{99, 3}) 19 | 20 | err = xerrors.New("ola amigo") 21 | if err != nil { 22 | // This is the first line that should be reported on xdefer. 23 | return xerrors.Errorf("first wrap: %w", err) 24 | } 25 | 26 | return nil 27 | }() 28 | 29 | _, fp, _, ok := runtime.Caller(0) 30 | if !ok { 31 | t.Fatal("runtime.Caller failed") 32 | } 33 | exp := fmt.Sprintf(`second wrap []int{99, 3}: 34 | - first wrap: 35 | oss.terrastruct.com/util-go/xdefer_test.TestErrorf.func1 36 | %v:23 37 | - ola amigo: 38 | oss.terrastruct.com/util-go/xdefer_test.TestErrorf.func1 39 | %[1]v:20`, 40 | fp, 41 | ) 42 | 43 | got := fmt.Sprintf("%+v", err) 44 | assert.String(t, exp, got) 45 | } 46 | 47 | func TestEmptyErrorf(t *testing.T) { 48 | t.Parallel() 49 | 50 | err := func() (err error) { 51 | defer xdefer.Errorf(&err, "") 52 | 53 | err = xerrors.New("ola amigo") 54 | if err != nil { 55 | return err 56 | } 57 | 58 | return nil 59 | }() 60 | 61 | _, fp, _, ok := runtime.Caller(0) 62 | if !ok { 63 | t.Fatal("runtime.Caller failed") 64 | } 65 | exp := fmt.Sprintf(`oss.terrastruct.com/util-go/xdefer_test.TestEmptyErrorf.func1 66 | %v:55 67 | - ola amigo: 68 | oss.terrastruct.com/util-go/xdefer_test.TestEmptyErrorf.func1 69 | %[1]v:53`, 70 | fp, 71 | ) 72 | 73 | got := fmt.Sprintf("%+v", err) 74 | assert.String(t, exp, got) 75 | 76 | exp = err.Error() 77 | got = "oss.terrastruct.com/util-go/xdefer_test.TestEmptyErrorf.func1: ola amigo" 78 | assert.String(t, exp, got) 79 | } 80 | -------------------------------------------------------------------------------- /xexec/xexec.go: -------------------------------------------------------------------------------- 1 | package xexec 2 | 3 | import ( 4 | "io/fs" 5 | "os" 6 | "path/filepath" 7 | "runtime" 8 | "strings" 9 | ) 10 | 11 | // findExecutable is from package os/exec 12 | func findExecutable(file string) error { 13 | d, err := os.Stat(file) 14 | if err != nil { 15 | return err 16 | } 17 | m := d.Mode() 18 | if !m.IsDir() && m&0111 != 0 { 19 | return nil 20 | } 21 | return fs.ErrPermission 22 | } 23 | 24 | // SearchPath searches for all executables that have prefix in their names in 25 | // the directories named by the PATH environment variable. 26 | func SearchPath(prefix string) ([]string, error) { 27 | var matches []string 28 | envPath := os.Getenv("PATH") 29 | dirSet := make(map[string]struct{}) 30 | for _, dir := range filepath.SplitList(envPath) { 31 | if dir == "" { 32 | // From exec package: 33 | // Unix shell semantics: path element "" means "." 34 | dir = "." 35 | } 36 | if _, ok := dirSet[dir]; ok { 37 | continue 38 | } 39 | dirSet[dir] = struct{}{} 40 | files, err := os.ReadDir(dir) 41 | if err != nil { 42 | continue 43 | } 44 | for _, f := range files { 45 | if strings.HasPrefix(f.Name(), prefix) { 46 | match := filepath.Join(dir, f.Name()) 47 | // Unideal but I don't want to maintain two separate implementations of this 48 | // function like os/exec. 49 | if runtime.GOOS == "windows" { 50 | matches = append(matches, match) 51 | continue 52 | } 53 | err = findExecutable(match) 54 | if err == nil { 55 | matches = append(matches, match) 56 | } 57 | } 58 | } 59 | 60 | } 61 | return matches, nil 62 | } 63 | -------------------------------------------------------------------------------- /xhttp/err.go: -------------------------------------------------------------------------------- 1 | package xhttp 2 | 3 | import ( 4 | "encoding/json" 5 | "errors" 6 | "fmt" 7 | "log" 8 | "net/http" 9 | 10 | "oss.terrastruct.com/util-go/cmdlog" 11 | ) 12 | 13 | // Error represents an HTTP error. 14 | // It's exported only for comparison in tests. 15 | type Error struct { 16 | Code int 17 | Resp interface{} 18 | Err error 19 | } 20 | 21 | var _ interface { 22 | Is(error) bool 23 | Unwrap() error 24 | } = Error{} 25 | 26 | // Errorf creates a new error with code, resp, msg and v. 27 | // 28 | // When returned from an xhttp.HandlerFunc, it will be correctly logged 29 | // and written to the connection. See xhttp.WrapHandlerFunc 30 | func Errorf(code int, resp interface{}, msg string, v ...interface{}) error { 31 | return errorWrap(code, resp, fmt.Errorf(msg, v...)) 32 | } 33 | 34 | // ErrorWrap wraps err with the code and resp for xhttp.HandlerFunc. 35 | // 36 | // When returned from an xhttp.HandlerFunc, it will be correctly logged 37 | // and written to the connection. See xhttp.WrapHandlerFunc 38 | func ErrorWrap(code int, resp interface{}, err error) error { 39 | return errorWrap(code, resp, err) 40 | } 41 | 42 | func errorWrap(code int, resp interface{}, err error) error { 43 | if resp == nil { 44 | resp = http.StatusText(code) 45 | } 46 | return Error{code, resp, err} 47 | } 48 | 49 | func (e Error) Unwrap() error { 50 | return e.Err 51 | } 52 | 53 | func (e Error) Is(err error) bool { 54 | e2, ok := err.(Error) 55 | if !ok { 56 | return false 57 | } 58 | return e.Code == e2.Code && e.Resp == e2.Resp && errors.Is(e.Err, e2.Err) 59 | } 60 | 61 | func (e Error) Error() string { 62 | return fmt.Sprintf("http error with code %v and resp %#v: %v", e.Code, e.Resp, e.Err) 63 | } 64 | 65 | // HandlerFunc is like http.HandlerFunc but returns an error. 66 | // See Errorf and ErrorWrap. 67 | type HandlerFunc func(w http.ResponseWriter, r *http.Request) error 68 | 69 | type HandlerFuncAdapter struct { 70 | Log *cmdlog.Logger 71 | Func HandlerFunc 72 | } 73 | 74 | // ServeHTTP adapts xhttp.HandlerFunc into http.Handler for usage with standard 75 | // HTTP routers like chi. 76 | // 77 | // It logs and writes any error from xhttp.HandlerFunc to the connection. 78 | // 79 | // If err was created with xhttp.Errorf or wrapped with xhttp.WrapError, then the error 80 | // will be logged at the correct level for the status code and xhttp.JSON will be called 81 | // with the code and resp. 82 | // 83 | // 400s are logged as warns and 500s as errors. 84 | // 85 | // If the error was not created with the xhttp helpers then a 500 will be written. 86 | // 87 | // If resp is nil, then resp is set to http.StatusText(code) 88 | // 89 | // If the code is not a 400 or a 500, then an error about about the unexpected error code 90 | // will be logged and a 500 will be written. The original error will also be logged. 91 | func (a HandlerFuncAdapter) ServeHTTP(w http.ResponseWriter, r *http.Request) { 92 | var h http.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 93 | err := a.Func(w, r) 94 | if err != nil { 95 | handleError(a.Log, w, err) 96 | } 97 | }) 98 | 99 | h.ServeHTTP(w, r) 100 | } 101 | 102 | func handleError(clog *cmdlog.Logger, w http.ResponseWriter, err error) { 103 | var herr Error 104 | ok := errors.As(err, &herr) 105 | if !ok { 106 | herr = ErrorWrap(http.StatusInternalServerError, nil, err).(Error) 107 | } 108 | 109 | var logger *log.Logger 110 | switch { 111 | case 400 <= herr.Code && herr.Code < 500: 112 | logger = clog.Warn 113 | case 500 <= herr.Code && herr.Code < 600: 114 | logger = clog.Error 115 | default: 116 | logger = clog.Error 117 | 118 | clog.Error.Printf("unexpected non error http status code %d with resp: %#v", herr.Code, herr.Resp) 119 | 120 | herr.Code = http.StatusInternalServerError 121 | herr.Resp = nil 122 | } 123 | 124 | if herr.Resp == nil { 125 | herr.Resp = http.StatusText(herr.Code) 126 | } 127 | 128 | logger.Printf("error handling http request: %v", err) 129 | 130 | ww, ok := w.(writtenResponseWriter) 131 | if !ok { 132 | clog.Warn.Printf("response writer does not implement Written, double write logs possible: %#v", w) 133 | } else if ww.Written() { 134 | // Avoid double writes if an error occurred while the response was 135 | // being written. 136 | return 137 | } 138 | 139 | JSON(clog, w, herr.Code, map[string]interface{}{ 140 | "error": herr.Resp, 141 | }) 142 | } 143 | 144 | type writtenResponseWriter interface { 145 | Written() bool 146 | } 147 | 148 | func JSON(clog *cmdlog.Logger, w http.ResponseWriter, code int, v interface{}) { 149 | if v == nil { 150 | v = map[string]interface{}{ 151 | "status": http.StatusText(code), 152 | } 153 | } 154 | 155 | b, err := json.Marshal(v) 156 | if err != nil { 157 | clog.Error.Printf("json marshal error: %v", err) 158 | http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) 159 | return 160 | } 161 | 162 | w.Header().Set("Content-Type", "application/json; charset=utf-8") 163 | w.WriteHeader(code) 164 | _, _ = w.Write(b) 165 | } 166 | -------------------------------------------------------------------------------- /xhttp/log.go: -------------------------------------------------------------------------------- 1 | package xhttp 2 | 3 | import ( 4 | "bufio" 5 | "errors" 6 | "fmt" 7 | "log" 8 | "net" 9 | "net/http" 10 | "runtime/debug" 11 | "time" 12 | 13 | "golang.org/x/text/message" 14 | 15 | "oss.terrastruct.com/util-go/cmdlog" 16 | ) 17 | 18 | type ResponseWriter interface { 19 | http.ResponseWriter 20 | http.Hijacker 21 | http.Flusher 22 | writtenResponseWriter 23 | } 24 | 25 | var _ ResponseWriter = &responseWriter{} 26 | 27 | type responseWriter struct { 28 | rw http.ResponseWriter 29 | 30 | written bool 31 | status int 32 | length int 33 | } 34 | 35 | func (rw *responseWriter) Header() http.Header { 36 | return rw.rw.Header() 37 | } 38 | 39 | func (rw *responseWriter) WriteHeader(statusCode int) { 40 | if !rw.written { 41 | rw.written = true 42 | rw.status = statusCode 43 | } 44 | rw.rw.WriteHeader(statusCode) 45 | } 46 | 47 | func (rw *responseWriter) Write(p []byte) (int, error) { 48 | if !rw.written && len(p) > 0 { 49 | rw.written = true 50 | if rw.status == 0 { 51 | rw.status = http.StatusOK 52 | } 53 | } 54 | rw.length += len(p) 55 | return rw.rw.Write(p) 56 | } 57 | 58 | func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { 59 | hj, ok := rw.rw.(http.Hijacker) 60 | if !ok { 61 | return nil, nil, fmt.Errorf("underlying response writer does not implement http.Hijacker: %T", rw.rw) 62 | } 63 | return hj.Hijack() 64 | } 65 | 66 | func (rw *responseWriter) Flush() { 67 | f, ok := rw.rw.(http.Flusher) 68 | if !ok { 69 | return 70 | } 71 | f.Flush() 72 | } 73 | 74 | func (rw *responseWriter) Written() bool { 75 | return rw.written 76 | } 77 | 78 | func Log(clog *cmdlog.Logger, next http.Handler) http.Handler { 79 | englishPrinter := message.NewPrinter(message.MatchLanguage("en")) 80 | return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 81 | defer func() { 82 | rec := recover() 83 | if rec != nil { 84 | clog.Error.Printf("caught panic: %#v\n%s", rec, debug.Stack()) 85 | JSON(clog, w, http.StatusInternalServerError, map[string]interface{}{ 86 | "error": http.StatusText(http.StatusInternalServerError), 87 | }) 88 | } 89 | }() 90 | 91 | rw := &responseWriter{ 92 | rw: w, 93 | } 94 | 95 | start := time.Now() 96 | next.ServeHTTP(rw, r) 97 | dur := time.Since(start) 98 | 99 | if !rw.Written() { 100 | _, err := rw.Write(nil) 101 | if errors.Is(err, http.ErrHijacked) { 102 | clog.Success.Printf("%s %s %v: hijacked", r.Method, r.URL, dur) 103 | return 104 | } 105 | 106 | clog.Warn.Printf("%s %s %v: no response written", r.Method, r.URL, dur) 107 | return 108 | } 109 | 110 | var statusLogger *log.Logger 111 | switch { 112 | case 100 <= rw.status && rw.status <= 299: 113 | statusLogger = clog.Success 114 | case 300 <= rw.status && rw.status <= 399: 115 | statusLogger = clog.Info 116 | case 400 <= rw.status && rw.status <= 499: 117 | statusLogger = clog.Warn 118 | case 500 <= rw.status && rw.status <= 599: 119 | statusLogger = clog.Error 120 | } 121 | lengthStr := englishPrinter.Sprint(rw.length) 122 | statusLogger.Printf("%s %s %d %sB %v", r.Method, r.URL, rw.status, lengthStr, dur) 123 | }) 124 | } 125 | -------------------------------------------------------------------------------- /xhttp/serve.go: -------------------------------------------------------------------------------- 1 | // Package xhttp implements http helpers. 2 | package xhttp 3 | 4 | import ( 5 | "context" 6 | "log" 7 | "net" 8 | "net/http" 9 | "time" 10 | 11 | "oss.terrastruct.com/util-go/xcontext" 12 | ) 13 | 14 | func NewServer(log *log.Logger, h http.Handler) *http.Server { 15 | return &http.Server{ 16 | MaxHeaderBytes: 1 << 18, // 262,144B 17 | ReadTimeout: time.Minute, 18 | WriteTimeout: time.Minute, 19 | IdleTimeout: time.Hour, 20 | ErrorLog: log, 21 | Handler: http.MaxBytesHandler(h, 1<<20), // 1,048,576B 22 | } 23 | } 24 | 25 | func Serve(ctx context.Context, shutdownTimeout time.Duration, s *http.Server, l net.Listener) error { 26 | s.BaseContext = func(net.Listener) context.Context { 27 | return ctx 28 | } 29 | 30 | done := make(chan error, 1) 31 | go func() { 32 | done <- s.Serve(l) 33 | }() 34 | 35 | select { 36 | case err := <-done: 37 | return err 38 | case <-ctx.Done(): 39 | shutdownCtx := xcontext.WithoutCancel(ctx) 40 | shutdownCtx, cancel := context.WithTimeout(shutdownCtx, shutdownTimeout) 41 | defer cancel() 42 | shutdownErr := s.Shutdown(shutdownCtx) 43 | serveErr := <-done 44 | if serveErr != nil && serveErr != http.ErrServerClosed { 45 | return serveErr 46 | } 47 | return shutdownErr 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /xjson/xjson.go: -------------------------------------------------------------------------------- 1 | // Package xjson implements JSON helpers. 2 | package xjson 3 | 4 | import ( 5 | "bytes" 6 | "encoding/json" 7 | ) 8 | 9 | // Marshal is like json.MarshalIndent but does not escape HTML. 10 | // And does not return an error 11 | func Marshal(v interface{}) []byte { 12 | var b bytes.Buffer 13 | e := json.NewEncoder(&b) 14 | // Allows < and > in JSON strings without escaping which we do with SrcArrow and 15 | // DstArrow. See https://stackoverflow.com/a/28596225 16 | e.SetEscapeHTML(false) 17 | e.SetIndent("", " ") 18 | err := e.Encode(v) 19 | if err != nil { 20 | buf, _ := json.Marshal(err.Error()) 21 | return buf 22 | } 23 | buf := b.Bytes() 24 | // Remove trailing newline. 25 | return buf[:len(buf)-1] 26 | } 27 | -------------------------------------------------------------------------------- /xmain/flag_helpers.go: -------------------------------------------------------------------------------- 1 | // flag_helpers.go are private functions from pflag/flag.go 2 | package xmain 3 | 4 | import "strings" 5 | 6 | func wrap(i, w int, s string) string { 7 | if w == 0 { 8 | return strings.Replace(s, "\n", "\n"+strings.Repeat(" ", i), -1) 9 | } 10 | wrap := w - i 11 | var r, l string 12 | if wrap < 24 { 13 | i = 16 14 | wrap = w - i 15 | r += "\n" + strings.Repeat(" ", i) 16 | } 17 | if wrap < 24 { 18 | return strings.Replace(s, "\n", r, -1) 19 | } 20 | slop := 5 21 | wrap = wrap - slop 22 | l, s = wrapN(wrap, slop, s) 23 | r = r + strings.Replace(l, "\n", "\n"+strings.Repeat(" ", i), -1) 24 | for s != "" { 25 | var t string 26 | t, s = wrapN(wrap, slop, s) 27 | r = r + "\n" + strings.Repeat(" ", i) + strings.Replace(t, "\n", "\n"+strings.Repeat(" ", i), -1) 28 | } 29 | return r 30 | } 31 | 32 | func wrapN(i, slop int, s string) (string, string) { 33 | if i+slop > len(s) { 34 | return s, "" 35 | } 36 | w := strings.LastIndexAny(s[:i], " \t\n") 37 | if w <= 0 { 38 | return s, "" 39 | } 40 | nlPos := strings.LastIndex(s[:i], "\n") 41 | if nlPos > 0 && nlPos < w { 42 | return s[:nlPos], s[nlPos+1:] 43 | } 44 | return s[:w], s[w+1:] 45 | } 46 | -------------------------------------------------------------------------------- /xmain/opts.go: -------------------------------------------------------------------------------- 1 | package xmain 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "io" 7 | "strconv" 8 | "strings" 9 | 10 | "github.com/spf13/pflag" 11 | 12 | "oss.terrastruct.com/util-go/xos" 13 | ) 14 | 15 | type Opts struct { 16 | Args []string 17 | Flags *pflag.FlagSet 18 | env *xos.Env 19 | 20 | flagEnv map[string]string 21 | } 22 | 23 | func NewOpts(env *xos.Env, args []string) *Opts { 24 | flags := pflag.NewFlagSet("", pflag.ContinueOnError) 25 | flags.SortFlags = false 26 | flags.Usage = func() {} 27 | flags.SetOutput(io.Discard) 28 | return &Opts{ 29 | Args: args, 30 | Flags: flags, 31 | env: env, 32 | flagEnv: make(map[string]string), 33 | } 34 | } 35 | 36 | // Mostly copy pasted pasted from pflag.FlagUsagesWrapped 37 | // with modifications for env var 38 | func (o *Opts) Defaults() string { 39 | buf := new(bytes.Buffer) 40 | 41 | var lines []string 42 | 43 | maxlen := 0 44 | maxEnvLen := 0 45 | o.Flags.VisitAll(func(flag *pflag.Flag) { 46 | if flag.Hidden { 47 | return 48 | } 49 | 50 | line := "" 51 | if flag.Shorthand != "" && flag.ShorthandDeprecated == "" { 52 | line = fmt.Sprintf(" -%s, --%s", flag.Shorthand, flag.Name) 53 | } else { 54 | line = fmt.Sprintf(" --%s", flag.Name) 55 | } 56 | 57 | varname, usage := pflag.UnquoteUsage(flag) 58 | if varname != "" { 59 | line += " " + varname 60 | } 61 | if flag.NoOptDefVal != "" { 62 | switch flag.Value.Type() { 63 | case "string": 64 | line += fmt.Sprintf("[=\"%s\"]", flag.NoOptDefVal) 65 | case "bool": 66 | if flag.NoOptDefVal != "true" { 67 | line += fmt.Sprintf("[=%s]", flag.NoOptDefVal) 68 | } 69 | case "count": 70 | if flag.NoOptDefVal != "+1" { 71 | line += fmt.Sprintf("[=%s]", flag.NoOptDefVal) 72 | } 73 | default: 74 | line += fmt.Sprintf("[=%s]", flag.NoOptDefVal) 75 | } 76 | } 77 | 78 | line += "\x00" 79 | 80 | if len(line) > maxlen { 81 | maxlen = len(line) 82 | } 83 | 84 | if e, ok := o.flagEnv[flag.Name]; ok { 85 | line += fmt.Sprintf("$%s", e) 86 | } 87 | 88 | line += "\x01" 89 | 90 | if len(line) > maxEnvLen { 91 | maxEnvLen = len(line) 92 | } 93 | 94 | line += usage 95 | if flag.Value.Type() == "string" { 96 | line += fmt.Sprintf(" (default %q)", flag.DefValue) 97 | } else { 98 | line += fmt.Sprintf(" (default %s)", flag.DefValue) 99 | } 100 | if len(flag.Deprecated) != 0 { 101 | line += fmt.Sprintf(" (DEPRECATED: %s)", flag.Deprecated) 102 | } 103 | 104 | lines = append(lines, line) 105 | }) 106 | 107 | for _, line := range lines { 108 | sidx1 := strings.Index(line, "\x00") 109 | sidx2 := strings.Index(line, "\x01") 110 | spacing1 := strings.Repeat(" ", maxlen-sidx1) 111 | spacing2 := strings.Repeat(" ", (maxEnvLen-maxlen)-sidx2+sidx1) 112 | fmt.Fprintln(buf, line[:sidx1], spacing1, line[sidx1+1:sidx2], spacing2, wrap(maxEnvLen+3, 0, line[sidx2+1:])) 113 | } 114 | 115 | return buf.String() 116 | } 117 | 118 | func (o *Opts) getEnv(flag, k string) string { 119 | if k != "" { 120 | o.flagEnv[flag] = k 121 | return o.env.Getenv(k) 122 | } 123 | return "" 124 | } 125 | 126 | func (o *Opts) Int64(envKey, flag, shortFlag string, defaultVal int64, usage string) (*int64, error) { 127 | if env := o.getEnv(flag, envKey); env != "" { 128 | envVal, err := strconv.ParseInt(env, 10, 64) 129 | if err != nil { 130 | return nil, UsageErrorf(`invalid environment variable %s. Expected int64. Found "%v".`, envKey, envVal) 131 | } 132 | defaultVal = envVal 133 | } 134 | 135 | return o.Flags.Int64P(flag, shortFlag, defaultVal, usage), nil 136 | } 137 | 138 | func (o *Opts) Int64Slice(envKey, flag, shortFlag string, defaultVal []int64, usage string) (*[]int64, error) { 139 | if env := o.getEnv(flag, envKey); env != "" { 140 | split := strings.Split(env, ",") 141 | defaultVal = make([]int64, len(split)) 142 | for i, part := range split { 143 | val, err := strconv.ParseInt(strings.TrimSpace(part), 10, 64) 144 | if err != nil { 145 | return nil, UsageErrorf(`invalid environment variable %s. Expected []int64. Found "%v".`, envKey, env) 146 | } 147 | defaultVal[i] = val 148 | } 149 | } 150 | 151 | return o.Flags.Int64SliceP(flag, shortFlag, defaultVal, usage), nil 152 | } 153 | 154 | func (o *Opts) Float64(envKey, flag, shortFlag string, defaultVal float64, usage string) (*float64, error) { 155 | if env := o.getEnv(flag, envKey); env != "" { 156 | envVal, err := strconv.ParseFloat(env, 64) 157 | if err != nil { 158 | return nil, UsageErrorf(`invalid environment variable %s. Expected float64. Found "%v".`, envKey, envVal) 159 | } 160 | defaultVal = envVal 161 | } 162 | 163 | return o.Flags.Float64P(flag, shortFlag, defaultVal, usage), nil 164 | } 165 | 166 | func (o *Opts) String(envKey, flag, shortFlag string, defaultVal, usage string) *string { 167 | if env := o.getEnv(flag, envKey); env != "" { 168 | defaultVal = env 169 | } 170 | 171 | return o.Flags.StringP(flag, shortFlag, defaultVal, usage) 172 | } 173 | 174 | func (o *Opts) Bool(envKey, flag, shortFlag string, defaultVal bool, usage string) (*bool, error) { 175 | if env := o.getEnv(flag, envKey); env != "" { 176 | if !boolyEnv(env) { 177 | return nil, UsageErrorf(`invalid environment variable %s. Expected bool. Found "%s".`, envKey, env) 178 | } 179 | if truthyEnv(env) { 180 | defaultVal = true 181 | } else { 182 | defaultVal = false 183 | } 184 | } 185 | 186 | return o.Flags.BoolP(flag, shortFlag, defaultVal, usage), nil 187 | } 188 | 189 | func boolyEnv(s string) bool { 190 | return falseyEnv(s) || truthyEnv(s) 191 | } 192 | 193 | func falseyEnv(s string) bool { 194 | return s == "0" || s == "false" 195 | } 196 | 197 | func truthyEnv(s string) bool { 198 | return s == "1" || s == "true" 199 | } 200 | -------------------------------------------------------------------------------- /xmain/stdlib_exec.go: -------------------------------------------------------------------------------- 1 | package xmain 2 | 3 | import ( 4 | "bytes" 5 | "strconv" 6 | ) 7 | 8 | // Code here was copied from src/os/exec/exec.go. 9 | 10 | // prefixSuffixSaver is an io.Writer which retains the first N bytes 11 | // and the last N bytes written to it. The Bytes() methods reconstructs 12 | // it with a pretty error message. 13 | type prefixSuffixSaver struct { 14 | N int // max size of prefix or suffix 15 | prefix []byte 16 | suffix []byte // ring buffer once len(suffix) == N 17 | suffixOff int // offset to write into suffix 18 | skipped int64 19 | } 20 | 21 | func (w *prefixSuffixSaver) Write(p []byte) (n int, err error) { 22 | lenp := len(p) 23 | p = w.fill(&w.prefix, p) 24 | 25 | // Only keep the last w.N bytes of suffix data. 26 | if overage := len(p) - w.N; overage > 0 { 27 | p = p[overage:] 28 | w.skipped += int64(overage) 29 | } 30 | p = w.fill(&w.suffix, p) 31 | 32 | // w.suffix is full now if p is non-empty. Overwrite it in a circle. 33 | for len(p) > 0 { // 0, 1, or 2 iterations. 34 | n := copy(w.suffix[w.suffixOff:], p) 35 | p = p[n:] 36 | w.skipped += int64(n) 37 | w.suffixOff += n 38 | if w.suffixOff == w.N { 39 | w.suffixOff = 0 40 | } 41 | } 42 | return lenp, nil 43 | } 44 | 45 | // fill appends up to len(p) bytes of p to *dst, such that *dst does not 46 | // grow larger than w.N. It returns the un-appended suffix of p. 47 | func (w *prefixSuffixSaver) fill(dst *[]byte, p []byte) (pRemain []byte) { 48 | if remain := w.N - len(*dst); remain > 0 { 49 | add := minInt(len(p), remain) 50 | *dst = append(*dst, p[:add]...) 51 | p = p[add:] 52 | } 53 | return p 54 | } 55 | 56 | func (w *prefixSuffixSaver) Bytes() []byte { 57 | if w.suffix == nil { 58 | return w.prefix 59 | } 60 | if w.skipped == 0 { 61 | return append(w.prefix, w.suffix...) 62 | } 63 | var buf bytes.Buffer 64 | buf.Grow(len(w.prefix) + len(w.suffix) + 50) 65 | buf.Write(w.prefix) 66 | buf.WriteString("\n... omitting ") 67 | buf.WriteString(strconv.FormatInt(w.skipped, 10)) 68 | buf.WriteString(" bytes ...\n") 69 | buf.Write(w.suffix[w.suffixOff:]) 70 | buf.Write(w.suffix[:w.suffixOff]) 71 | return buf.Bytes() 72 | } 73 | 74 | func minInt(a, b int) int { 75 | if a < b { 76 | return a 77 | } 78 | return b 79 | } 80 | -------------------------------------------------------------------------------- /xmain/xmain.go: -------------------------------------------------------------------------------- 1 | // Package xmain provides a standard stub for the main of a command handling logging, 2 | // flags, signals and shutdown. 3 | package xmain 4 | 5 | import ( 6 | "context" 7 | "errors" 8 | "fmt" 9 | "io" 10 | "os" 11 | "os/signal" 12 | "path/filepath" 13 | "strings" 14 | "syscall" 15 | "time" 16 | 17 | "oss.terrastruct.com/util-go/cmdlog" 18 | "oss.terrastruct.com/util-go/xos" 19 | ) 20 | 21 | type RunFunc func(context.Context, *State) error 22 | 23 | func Main(run RunFunc) { 24 | name := "" 25 | args := []string(nil) 26 | if len(os.Args) > 0 { 27 | name = os.Args[0] 28 | args = os.Args[1:] 29 | } 30 | 31 | ms := &State{ 32 | Name: name, 33 | 34 | Stdin: os.Stdin, 35 | Stdout: os.Stdout, 36 | Stderr: os.Stderr, 37 | 38 | Env: xos.NewEnv(os.Environ()), 39 | } 40 | ms.Log = cmdlog.New(ms.Env, ms.Stderr) 41 | ms.Opts = NewOpts(ms.Env, args) 42 | 43 | wd, err := os.Getwd() 44 | if err != nil { 45 | ms.mainFatal(err) 46 | } 47 | ms.PWD = wd 48 | 49 | sigs := make(chan os.Signal, 1) 50 | signal.Notify(sigs, os.Interrupt, syscall.SIGTERM) 51 | 52 | err = ms.Main(context.Background(), sigs, run) 53 | if err != nil { 54 | ms.mainFatal(err) 55 | } 56 | } 57 | 58 | func (ms *State) mainFatal(err error) { 59 | code := 1 60 | msg := "" 61 | usage := false 62 | 63 | var eerr ExitError 64 | var uerr UsageError 65 | if errors.As(err, &eerr) { 66 | code = eerr.Code 67 | msg = eerr.Message 68 | } else if errors.As(err, &uerr) { 69 | msg = err.Error() 70 | usage = true 71 | } else { 72 | msg = err.Error() 73 | } 74 | 75 | if msg != "" { 76 | ms.Log.Error.Print(msg) 77 | if usage { 78 | ms.Log.Error.Print("Run with --help to see usage.") 79 | } 80 | } 81 | os.Exit(code) 82 | } 83 | 84 | type State struct { 85 | Name string 86 | 87 | Stdin io.Reader 88 | Stdout io.WriteCloser 89 | Stderr io.WriteCloser 90 | 91 | Log *cmdlog.Logger 92 | Env *xos.Env 93 | Opts *Opts 94 | 95 | PWD string 96 | } 97 | 98 | func (ms *State) Main(ctx context.Context, sigs <-chan os.Signal, run func(context.Context, *State) error) error { 99 | ctx, cancel := context.WithCancel(ctx) 100 | defer cancel() 101 | 102 | done := make(chan error, 1) 103 | go func() { 104 | defer close(done) 105 | done <- run(ctx, ms) 106 | }() 107 | 108 | select { 109 | case err := <-done: 110 | return err 111 | case sig := <-sigs: 112 | ms.Log.Warn.Printf("received signal %v: shutting down...", sig) 113 | cancel() 114 | select { 115 | case err := <-done: 116 | if err != nil && !errors.Is(err, context.Canceled) { 117 | return fmt.Errorf("failed to shutdown: %w", err) 118 | } 119 | if sig == syscall.SIGTERM { 120 | // We successfully shutdown. 121 | return nil 122 | } 123 | return ExitError{Code: 1} 124 | case <-time.After(time.Minute): 125 | return ExitError{ 126 | Code: 1, 127 | Message: "took longer than 1 minute to shutdown: exiting forcefully", 128 | } 129 | } 130 | } 131 | } 132 | 133 | type ExitError struct { 134 | Code int `json:"code"` 135 | Message string `json:"message"` 136 | } 137 | 138 | func ExitErrorf(code int, msg string, v ...interface{}) ExitError { 139 | return ExitError{ 140 | Code: code, 141 | Message: fmt.Sprintf(msg, v...), 142 | } 143 | } 144 | 145 | func (ee ExitError) Error() string { 146 | s := fmt.Sprintf("exiting with code %d", ee.Code) 147 | if ee.Message != "" { 148 | s += ": " + ee.Message 149 | } 150 | return s 151 | } 152 | 153 | type UsageError struct { 154 | Message string `json:"message"` 155 | } 156 | 157 | func UsageErrorf(msg string, v ...interface{}) UsageError { 158 | return UsageError{ 159 | Message: fmt.Sprintf(msg, v...), 160 | } 161 | } 162 | 163 | func (ue UsageError) Error() string { 164 | return fmt.Sprintf("bad usage: %s", ue.Message) 165 | } 166 | 167 | func (ms *State) ReadPath(fp string) ([]byte, error) { 168 | if fp == "-" { 169 | return io.ReadAll(ms.Stdin) 170 | } 171 | return os.ReadFile(fp) 172 | } 173 | 174 | func (ms *State) WritePath(fp string, p []byte) error { 175 | if fp == "-" { 176 | _, err := ms.Stdout.Write(p) 177 | if err != nil { 178 | return err 179 | } 180 | return ms.Stdout.Close() 181 | } 182 | return os.WriteFile(fp, p, 0644) 183 | } 184 | 185 | func (ms *State) AtomicWritePath(fp string, p []byte) error { 186 | if fp == "-" { 187 | return ms.WritePath(fp, p) 188 | } 189 | 190 | dir := filepath.Dir(fp) 191 | base := filepath.Base(fp) 192 | tempFile, err := os.CreateTemp(dir, "tmp-"+base+"-") 193 | if err != nil { 194 | return err 195 | } 196 | defer func() { 197 | // Clean up temporary file if it still exists and there's an error 198 | if err != nil { 199 | os.Remove(tempFile.Name()) 200 | } 201 | }() 202 | 203 | if _, err = tempFile.Write(p); err != nil { 204 | return err 205 | } 206 | 207 | if err = tempFile.Close(); err != nil { 208 | return err 209 | } 210 | 211 | if err = os.Rename(tempFile.Name(), fp); err != nil { 212 | return err 213 | } 214 | 215 | return nil 216 | } 217 | 218 | // AbsPath joins the PWD with fp to give the absolute path to fp. 219 | func (ms *State) AbsPath(fp string) string { 220 | if fp == "-" || filepath.IsAbs(fp) { 221 | return fp 222 | } 223 | return filepath.Join(ms.PWD, fp) 224 | } 225 | 226 | // HumanPath makes absolute path fp more suitable for human consumption 227 | // by replacing $HOME in fp with ~ and making it relative to the current PWD. 228 | func (ms *State) HumanPath(fp string) string { 229 | if fp == "-" { 230 | return fp 231 | } 232 | fp = ms.AbsPath(fp) 233 | 234 | if strings.HasPrefix(fp, ms.Env.Getenv("HOME")) { 235 | fp = filepath.Join("~", strings.TrimPrefix(fp, ms.Env.Getenv("HOME"))) 236 | } 237 | pwd := ms.PWD 238 | if strings.HasPrefix(pwd, ms.Env.Getenv("HOME")) { 239 | pwd = filepath.Join("~", strings.TrimPrefix(pwd, ms.Env.Getenv("HOME"))) 240 | } 241 | 242 | rel, err := filepath.Rel(pwd, fp) 243 | if err != nil { 244 | return fp 245 | } 246 | return rel 247 | } 248 | -------------------------------------------------------------------------------- /xmain/xmaintest.go: -------------------------------------------------------------------------------- 1 | package xmain 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "io" 8 | "os" 9 | "sync" 10 | "testing" 11 | "time" 12 | 13 | "oss.terrastruct.com/util-go/assert" 14 | "oss.terrastruct.com/util-go/cmdlog" 15 | "oss.terrastruct.com/util-go/xdefer" 16 | "oss.terrastruct.com/util-go/xos" 17 | ) 18 | 19 | type TestState struct { 20 | Run func(context.Context, *State) error 21 | Env *xos.Env 22 | Args []string 23 | PWD string 24 | 25 | Stdin io.Reader 26 | Stdout io.Writer 27 | Stderr io.Writer 28 | 29 | ms *State 30 | sigs chan os.Signal 31 | done chan struct{} 32 | doneErr *error 33 | } 34 | 35 | func (ts *TestState) StdinPipe() (pw io.WriteCloser) { 36 | ts.Stdin, pw = io.Pipe() 37 | return pw 38 | } 39 | 40 | func (ts *TestState) StdoutPipe() (pr io.Reader) { 41 | pr, ts.Stdout = io.Pipe() 42 | return pr 43 | } 44 | 45 | func (ts *TestState) StderrPipe() (pr io.Reader) { 46 | pr, ts.Stderr = io.Pipe() 47 | return pr 48 | } 49 | 50 | func (ts *TestState) Start(tb testing.TB, ctx context.Context) { 51 | tb.Helper() 52 | 53 | if ts.done != nil { 54 | tb.Fatal("xmain.TestState.Start cannot be called twice") 55 | } 56 | 57 | if ts.Env == nil { 58 | ts.Env = xos.NewEnv(nil) 59 | } 60 | var tempDirCleanup func() 61 | if ts.PWD == "" { 62 | ts.PWD, tempDirCleanup = assert.TempDir(tb) 63 | } 64 | 65 | ts.sigs = make(chan os.Signal, 1) 66 | ts.done = make(chan struct{}) 67 | 68 | name := "" 69 | args := []string(nil) 70 | if len(ts.Args) > 0 { 71 | name = ts.Args[0] 72 | args = ts.Args[1:] 73 | } 74 | log := cmdlog.NewTB(ts.Env, tb) 75 | ts.ms = &State{ 76 | Name: name, 77 | 78 | Log: log, 79 | Env: ts.Env, 80 | Opts: NewOpts(ts.Env, args), 81 | PWD: ts.PWD, 82 | } 83 | 84 | if ts.Stdin == nil { 85 | ts.ms.Stdin = io.LimitReader(nil, 0) 86 | } else if rc, ok := ts.Stdin.(io.ReadCloser); ok { 87 | ts.ms.Stdin = rc 88 | } else { 89 | var pw io.WriteCloser 90 | ts.ms.Stdin, pw = io.Pipe() 91 | go func() { 92 | defer pw.Close() 93 | io.Copy(pw, ts.Stdin) 94 | }() 95 | } 96 | 97 | var pipeWG sync.WaitGroup 98 | if ts.Stdout == nil { 99 | ts.ms.Stdout = nopWriterCloser{io.Discard} 100 | } else if wc, ok := ts.Stdout.(io.WriteCloser); ok { 101 | ts.ms.Stdout = wc 102 | } else { 103 | var pr io.Reader 104 | pr, ts.ms.Stdout = io.Pipe() 105 | pipeWG.Add(1) 106 | go func() { 107 | defer pipeWG.Done() 108 | io.Copy(ts.Stdout, pr) 109 | }() 110 | } 111 | if ts.Stderr == nil { 112 | ts.ms.Stderr = nopWriterCloser{&prefixSuffixSaver{N: 1 << 25}} 113 | } else if wc, ok := ts.Stderr.(io.WriteCloser); ok { 114 | ts.ms.Stderr = wc 115 | } else { 116 | var pr io.Reader 117 | pr, ts.ms.Stderr = io.Pipe() 118 | pipeWG.Add(1) 119 | go func() { 120 | defer pipeWG.Done() 121 | io.Copy(ts.Stderr, pr) 122 | }() 123 | } 124 | ts.ms.Log = cmdlog.New(ts.ms.Env, ts.ms.Stderr) 125 | 126 | go func() { 127 | var err error 128 | defer func() { 129 | ts.closeStdin() 130 | ts.ms.Stdout.Close() 131 | ts.ms.Stderr.Close() 132 | pipeWG.Wait() 133 | if tempDirCleanup != nil { 134 | tempDirCleanup() 135 | } 136 | ts.doneErr = &err 137 | close(ts.done) 138 | }() 139 | err = ts.ms.Main(ctx, ts.sigs, ts.Run) 140 | if err != nil { 141 | if ts.Stderr == nil { 142 | stderr := ts.ms.Stderr.(nopWriterCloser).Writer.(*prefixSuffixSaver).Bytes() 143 | if len(stderr) > 0 { 144 | err = fmt.Errorf("%w; stderr: %s", err, stderr) 145 | } 146 | } 147 | } 148 | }() 149 | } 150 | 151 | func (ts *TestState) closeStdin() { 152 | if rc, ok := ts.ms.Stdin.(io.ReadCloser); ok { 153 | rc.Close() 154 | } 155 | } 156 | 157 | func (ts *TestState) Cleanup(tb testing.TB) { 158 | tb.Helper() 159 | 160 | select { 161 | case <-ts.done: 162 | // Already exited. 163 | return 164 | default: 165 | } 166 | 167 | ts.closeStdin() 168 | 169 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute) 170 | defer cancel() 171 | err := ts.Signal(ctx, os.Interrupt) 172 | if err != nil { 173 | tb.Errorf("failed to os.Interrupt xmain test: %v", err) 174 | } 175 | err = ts.Wait(ctx) 176 | if errors.Is(err, context.DeadlineExceeded) { 177 | err = ts.Signal(ctx, os.Kill) 178 | if err != nil { 179 | tb.Errorf("failed to kill xmain test: %v", err) 180 | } 181 | ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) 182 | defer cancel() 183 | err = ts.Wait(ctx) 184 | } 185 | assert.Success(tb, err) 186 | } 187 | 188 | func (ts *TestState) Signal(ctx context.Context, sig os.Signal) (err error) { 189 | defer xdefer.Errorf(&err, "failed to signal xmain test: %v", ts.ms.Name) 190 | 191 | select { 192 | case <-ctx.Done(): 193 | return ctx.Err() 194 | case <-ts.done: 195 | return fmt.Errorf("xmain test exited: %w", *ts.doneErr) 196 | case ts.sigs <- sig: 197 | return nil 198 | } 199 | } 200 | 201 | func (ts *TestState) Wait(ctx context.Context) (err error) { 202 | defer xdefer.Errorf(&err, "failed to wait xmain test: %v", ts.ms.Name) 203 | 204 | select { 205 | case <-ctx.Done(): 206 | return ctx.Err() 207 | case <-ts.done: 208 | return *ts.doneErr 209 | } 210 | } 211 | 212 | type nopWriterCloser struct { 213 | io.Writer 214 | } 215 | 216 | func (c nopWriterCloser) Close() error { 217 | return nil 218 | } 219 | -------------------------------------------------------------------------------- /xmain/xmaintest_test.go: -------------------------------------------------------------------------------- 1 | package xmain_test 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "io" 8 | "path/filepath" 9 | "strings" 10 | "testing" 11 | 12 | "github.com/spf13/pflag" 13 | 14 | "oss.terrastruct.com/util-go/assert" 15 | "oss.terrastruct.com/util-go/xmain" 16 | "oss.terrastruct.com/util-go/xos" 17 | ) 18 | 19 | func TestTesting(t *testing.T) { 20 | t.Parallel() 21 | 22 | tca := []struct { 23 | name string 24 | run func(t *testing.T, ctx context.Context, env *xos.Env) 25 | }{ 26 | { 27 | name: "base", 28 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 29 | ts := &xmain.TestState{ 30 | Run: helloWorldRun, 31 | Env: env, 32 | Args: []string{"helloWorldRun"}, 33 | } 34 | 35 | ts.Start(t, ctx) 36 | defer ts.Cleanup(t) 37 | 38 | err := ts.Wait(ctx) 39 | assert.ErrorString(t, err, `failed to wait xmain test: helloWorldRun: bad usage: $HELLO_FLAG or -flag missing`) 40 | }, 41 | }, 42 | { 43 | name: "help", 44 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 45 | stdout := &strings.Builder{} 46 | ts := &xmain.TestState{ 47 | Run: helloWorldRun, 48 | Env: env, 49 | Args: []string{"helloWorldRun", "-help"}, 50 | Stdout: stdout, 51 | } 52 | 53 | ts.Start(t, ctx) 54 | defer ts.Cleanup(t) 55 | 56 | err := ts.Wait(ctx) 57 | assert.Success(t, err) 58 | 59 | assert.Equal(t, `Usage: 60 | helloWorldRun [-flag=val] 61 | 62 | helloWorldRun prints the value of -flag to stdout. $HELLO_FLAG is equivalent to -flag. 63 | `, stdout.String()) 64 | }, 65 | }, 66 | { 67 | name: "envPriority", 68 | run: func(t *testing.T, ctx context.Context, env *xos.Env) { 69 | env.Setenv("HELLO_FLAG", "world") 70 | stdout := &strings.Builder{} 71 | ts := &xmain.TestState{ 72 | Run: helloWorldRun, 73 | Env: env, 74 | Args: []string{"helloWorldRun", "hello"}, 75 | Stdout: stdout, 76 | } 77 | 78 | ts.Start(t, ctx) 79 | defer ts.Cleanup(t) 80 | 81 | err := ts.Wait(ctx) 82 | assert.Success(t, err) 83 | 84 | assert.Equal(t, "world", stdout.String()) 85 | }, 86 | }, 87 | } 88 | 89 | ctx := context.Background() 90 | for _, tc := range tca { 91 | tc := tc 92 | t.Run(tc.name, func(t *testing.T) { 93 | t.Parallel() 94 | 95 | ctx, cancel := context.WithCancel(ctx) 96 | defer cancel() 97 | 98 | tc.run(t, ctx, xos.NewEnv(nil)) 99 | }) 100 | } 101 | } 102 | 103 | func helloWorldRun(ctx context.Context, ms *xmain.State) error { 104 | flag := ms.Opts.String("HELLO_FLAG", "flag", "f", "", "") 105 | err := ms.Opts.Flags.Parse(ms.Opts.Args) 106 | if errors.Is(err, pflag.ErrHelp) { 107 | fmt.Fprintf(ms.Stdout, `Usage: 108 | %[1]s [-flag=val] 109 | 110 | %[1]s prints the value of -flag to stdout. $HELLO_FLAG is equivalent to -flag. 111 | `, filepath.Base(ms.Name)) 112 | return nil 113 | } 114 | if *flag == "" { 115 | return xmain.UsageErrorf("$HELLO_FLAG or -flag missing") 116 | } 117 | _, err = io.WriteString(ms.Stdout, *flag) 118 | return err 119 | } 120 | -------------------------------------------------------------------------------- /xos/env.go: -------------------------------------------------------------------------------- 1 | package xos 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "path/filepath" 7 | "strings" 8 | "sync" 9 | ) 10 | 11 | type Env struct { 12 | environMu sync.RWMutex 13 | environ []string 14 | } 15 | 16 | func NewEnv(environ []string) *Env { 17 | return &Env{ 18 | environ: environ, 19 | } 20 | } 21 | 22 | func (e *Env) Environ() []string { 23 | e.environMu.RLock() 24 | defer e.environMu.RUnlock() 25 | 26 | environ2 := make([]string, 0, len(e.environ)) 27 | return append(environ2, e.environ...) 28 | } 29 | 30 | func (e *Env) Setenv(name, value string) { 31 | e.environMu.Lock() 32 | defer e.environMu.Unlock() 33 | 34 | l := fmt.Sprintf("%s=%s", name, value) 35 | 36 | for i, l2 := range e.environ { 37 | j := strings.Index(l2, "=") 38 | if j == -1 { 39 | continue 40 | } 41 | name2 := l2[:j] 42 | if name != name2 { 43 | continue 44 | } 45 | e.environ[i] = l 46 | return 47 | } 48 | e.environ = append(e.environ, l) 49 | } 50 | 51 | func (e *Env) Getenv(name string) string { 52 | e.environMu.RLock() 53 | defer e.environMu.RUnlock() 54 | 55 | for _, l := range e.environ { 56 | i := strings.Index(l, "=") 57 | if i == -1 { 58 | continue 59 | } 60 | name2 := l[:i] 61 | if name == name2 { 62 | return l[i+1:] 63 | } 64 | } 65 | return "" 66 | } 67 | 68 | func (e *Env) Bool(name string) (*bool, error) { 69 | ev := e.Getenv(name) 70 | if ev == "" { 71 | return nil, nil 72 | } 73 | eb := new(bool) 74 | if ev == "0" || ev == "false" { 75 | return eb, nil 76 | } 77 | if ev == "1" || ev == "true" { 78 | *eb = true 79 | return eb, nil 80 | } 81 | return nil, fmt.Errorf("$%s must be 0, 1, false or true but got %q", name, ev) 82 | } 83 | 84 | func (e *Env) HumanPath(fp string) string { 85 | if strings.HasPrefix(fp, e.Getenv("HOME")) { 86 | return filepath.Join("~", strings.TrimPrefix(fp, e.Getenv("HOME"))) 87 | } 88 | return fp 89 | } 90 | 91 | func (e *Env) Debug() bool { 92 | debug, err := e.Bool("DEBUG") 93 | if err != nil { 94 | os.Stderr.WriteString(fmt.Sprintf("env: %v", err)) 95 | return false 96 | } 97 | if debug == nil || !*debug { 98 | return false 99 | } 100 | return true 101 | } 102 | 103 | func (e *Env) CI() bool { 104 | ci, err := e.Bool("CI") 105 | if err != nil { 106 | os.Stderr.WriteString(fmt.Sprintf("env: %v", err)) 107 | return false 108 | } 109 | if ci == nil || !*ci { 110 | return false 111 | } 112 | return true 113 | } 114 | -------------------------------------------------------------------------------- /xos/env_test.go: -------------------------------------------------------------------------------- 1 | package xos_test 2 | 3 | import ( 4 | "testing" 5 | 6 | "oss.terrastruct.com/util-go/assert" 7 | "oss.terrastruct.com/util-go/xos" 8 | ) 9 | 10 | func TestEnv(t *testing.T) { 11 | t.Parallel() 12 | 13 | e := xos.NewEnv(nil) 14 | assert.String(t, "", e.Getenv("NONE")) 15 | 16 | e.Setenv("DEBUG", "MEOW") 17 | e.Setenv("DEBUG2", "TWO") 18 | e.Setenv("DEBUG3", "THREE") 19 | e.Setenv("DEBUG4", "FOUR") 20 | 21 | assert.StringJSON(t, `[ 22 | "DEBUG=MEOW", 23 | "DEBUG2=TWO", 24 | "DEBUG3=THREE", 25 | "DEBUG4=FOUR" 26 | ]`, e.Environ()) 27 | 28 | e.Setenv("DEBUG", "D2") 29 | assert.StringJSON(t, `[ 30 | "DEBUG=D2", 31 | "DEBUG2=TWO", 32 | "DEBUG3=THREE", 33 | "DEBUG4=FOUR" 34 | ]`, e.Environ()) 35 | assert.String(t, "D2", e.Getenv("DEBUG")) 36 | } 37 | 38 | func TestBool(t *testing.T) { 39 | t.Parallel() 40 | 41 | env := xos.NewEnv(nil) 42 | 43 | eb, err := env.Bool("MY_BOOL") 44 | assert.Success(t, err) 45 | assert.StringJSON(t, `null`, eb) 46 | 47 | env.Setenv("MY_BOOL", "0") 48 | eb, err = env.Bool("MY_BOOL") 49 | assert.Success(t, err) 50 | assert.StringJSON(t, `false`, eb) 51 | 52 | env.Setenv("MY_BOOL", "1") 53 | eb, err = env.Bool("MY_BOOL") 54 | assert.Success(t, err) 55 | assert.StringJSON(t, `true`, eb) 56 | 57 | env.Setenv("MY_BOOL", "false") 58 | eb, err = env.Bool("MY_BOOL") 59 | assert.Success(t, err) 60 | assert.StringJSON(t, `false`, eb) 61 | 62 | env.Setenv("MY_BOOL", "true") 63 | eb, err = env.Bool("MY_BOOL") 64 | assert.Success(t, err) 65 | assert.StringJSON(t, `true`, eb) 66 | 67 | env.Setenv("MY_BOOL", "TRUE") 68 | eb, err = env.Bool("MY_BOOL") 69 | assert.Error(t, err) 70 | assert.ErrorString(t, err, `$MY_BOOL must be 0, 1, false or true but got "TRUE"`) 71 | } 72 | -------------------------------------------------------------------------------- /xos/xos.go: -------------------------------------------------------------------------------- 1 | // Package xos provides useful os helpers. 2 | package xos 3 | -------------------------------------------------------------------------------- /xrand/xrand.go: -------------------------------------------------------------------------------- 1 | // Package xrand provides helpers for generating useful random values. 2 | package xrand 3 | 4 | import ( 5 | "crypto/rand" 6 | "encoding/base64" 7 | "fmt" 8 | "io" 9 | math_rand "math/rand" 10 | "strings" 11 | "time" 12 | "unicode/utf8" 13 | ) 14 | 15 | func init() { 16 | math_rand.Seed(time.Now().UnixNano()) 17 | } 18 | 19 | // 0 is a valid unicode codepoint but certain programs 20 | // won't allow 0 to avoid clashing with c strings. e.g. pq. 21 | func randRune() rune { 22 | for { 23 | if Bool() { 24 | // Generate plain ASCII half the time. 25 | if math_rand.Int31n(100) == 0 { 26 | // Generate newline 1% of the time. 27 | return '\n' 28 | } 29 | return math_rand.Int31n(128) + 1 30 | } 31 | r := math_rand.Int31n(utf8.MaxRune+1) + 1 32 | if utf8.ValidRune(r) { 33 | return r 34 | } 35 | } 36 | } 37 | 38 | func String(n int, exclude []rune) string { 39 | var b strings.Builder 40 | for i := 0; i < n; i++ { 41 | r := randRune() 42 | excluded := false 43 | for _, xr := range exclude { 44 | if r == xr { 45 | excluded = true 46 | break 47 | } 48 | } 49 | if excluded { 50 | i-- 51 | continue 52 | } 53 | b.WriteRune(r) 54 | } 55 | return b.String() 56 | } 57 | 58 | func Bool() bool { 59 | return math_rand.Intn(2) == 0 60 | } 61 | 62 | func Base64(n int) string { 63 | var b strings.Builder 64 | b.Grow(n) 65 | wc := base64.NewEncoder(base64.URLEncoding, &b) 66 | 67 | _, err := io.CopyN(wc, rand.Reader, int64(n)) 68 | if err != nil { 69 | panic(fmt.Sprintf("error encoding base64: %v", err)) 70 | } 71 | 72 | err = wc.Close() 73 | if err != nil { 74 | panic(fmt.Sprintf("error encoding base64: %v", err)) 75 | } 76 | 77 | return b.String() 78 | } 79 | -------------------------------------------------------------------------------- /xrand/xrand_test.go: -------------------------------------------------------------------------------- 1 | package xrand_test 2 | 3 | import ( 4 | "testing" 5 | 6 | "oss.terrastruct.com/util-go/xrand" 7 | ) 8 | 9 | func TestString(t *testing.T) { 10 | t.Parallel() 11 | 12 | t.Run("invalids", func(t *testing.T) { 13 | t.Parallel() 14 | 15 | s := xrand.String(0, nil) 16 | if s != "" { 17 | t.Fatalf("expected empty string: %q", s) 18 | } 19 | s = xrand.String(-1, nil) 20 | if s != "" { 21 | t.Fatalf("expected empty string: %q", s) 22 | } 23 | }) 24 | 25 | s := xrand.String(20, nil) 26 | t.Logf("%d %s", len(s), s) 27 | } 28 | -------------------------------------------------------------------------------- /xterm/xterm.go: -------------------------------------------------------------------------------- 1 | package xterm 2 | 3 | import ( 4 | "fmt" 5 | "hash/crc64" 6 | "io" 7 | "math/rand" 8 | "os" 9 | 10 | "golang.org/x/term" 11 | 12 | "oss.terrastruct.com/util-go/xos" 13 | ) 14 | 15 | // See https://invisible-island.net/xterm/ctlseqs/ctlseqs.html#h3-Functions-using-CSI-_-ordered-by-the-final-character_s_ 16 | const ( 17 | csi = "\x1b[" 18 | reset = csi + "0m" 19 | 20 | Bold = csi + "1m" 21 | 22 | Red = csi + "31m" 23 | Green = csi + "32m" 24 | Yellow = csi + "33m" 25 | Blue = csi + "34m" 26 | Magenta = csi + "35m" 27 | Cyan = csi + "36m" 28 | 29 | BrightRed = csi + "91m" 30 | BrightGreen = csi + "92m" 31 | BrightYellow = csi + "93m" 32 | BrightBlue = csi + "94m" 33 | BrightMagenta = csi + "95m" 34 | BrightCyan = csi + "96m" 35 | ) 36 | 37 | var colors = [...]string{ 38 | Red, 39 | Green, 40 | Yellow, 41 | Blue, 42 | Magenta, 43 | Cyan, 44 | 45 | BrightRed, 46 | BrightGreen, 47 | BrightYellow, 48 | BrightBlue, 49 | BrightMagenta, 50 | BrightCyan, 51 | } 52 | 53 | // isTTY checks whether the given writer is a *os.File TTY. 54 | func isTTY(w io.Writer) bool { 55 | f, ok := w.(interface { 56 | Fd() uintptr 57 | }) 58 | return ok && term.IsTerminal(int(f.Fd())) 59 | } 60 | 61 | func shouldColor(env *xos.Env, w io.Writer) bool { 62 | eb, err := env.Bool("COLOR") 63 | if eb != nil { 64 | return *eb 65 | } 66 | if err != nil { 67 | os.Stderr.WriteString(fmt.Sprintf("xterm: %v", err)) 68 | } 69 | if env.Getenv("TERM") == "dumb" { 70 | return false 71 | } 72 | return isTTY(w) 73 | } 74 | 75 | func Tput(env *xos.Env, w io.Writer, caps, s string) string { 76 | if caps == "" { 77 | return s 78 | } 79 | if !shouldColor(env, w) { 80 | return s 81 | } 82 | return caps + s + reset 83 | } 84 | 85 | func Prefix(env *xos.Env, w io.Writer, caps, s string) string { 86 | s = fmt.Sprintf("%s", s) 87 | return Tput(env, w, caps, s) + ":" 88 | } 89 | 90 | var crc64Table = crc64.MakeTable(crc64.ISO) 91 | 92 | // CC meaning constant color. So constant color prefix. 93 | func CCPrefix(env *xos.Env, w io.Writer, s string) string { 94 | sum := crc64.Checksum([]byte(s), crc64Table) 95 | rand := rand.New(rand.NewSource(int64(sum))) 96 | 97 | color := colors[rand.Intn(len(colors))] 98 | return Prefix(env, w, color, s) 99 | } 100 | --------------------------------------------------------------------------------