github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/libraries/doltcore/sqle/cluster/progress_notifier_test.go (about) 1 // Copyright 2023 Dolthub, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package cluster 16 17 import ( 18 "context" 19 "testing" 20 "time" 21 22 "github.com/stretchr/testify/assert" 23 ) 24 25 func TestProgressNotifier(t *testing.T) { 26 t.Run("WaitBeforeBeginAttempt", func(t *testing.T) { 27 p := new(ProgressNotifier) 28 f := p.Wait() 29 a := p.BeginAttempt() 30 p.RecordSuccess(a) 31 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) 32 t.Cleanup(cancel) 33 assert.NoError(t, f(ctx)) 34 }) 35 36 t.Run("WaitAfterBeginAttempt", func(t *testing.T) { 37 p := new(ProgressNotifier) 38 a := p.BeginAttempt() 39 f := p.Wait() 40 p.RecordSuccess(a) 41 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) 42 t.Cleanup(cancel) 43 assert.ErrorIs(t, f(ctx), context.DeadlineExceeded) 44 45 a = p.BeginAttempt() 46 p.RecordSuccess(a) 47 ctx, cancel = context.WithTimeout(context.Background(), 10*time.Millisecond) 48 t.Cleanup(cancel) 49 assert.NoError(t, f(ctx)) 50 }) 51 52 t.Run("WaitBeforeAttemptFailure", func(t *testing.T) { 53 p := new(ProgressNotifier) 54 f := p.Wait() 55 a := p.BeginAttempt() 56 p.RecordFailure(a) 57 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) 58 t.Cleanup(cancel) 59 assert.ErrorIs(t, f(ctx), context.DeadlineExceeded) 60 61 a = p.BeginAttempt() 62 p.RecordSuccess(a) 63 ctx, cancel = context.WithTimeout(context.Background(), 10*time.Millisecond) 64 t.Cleanup(cancel) 65 assert.NoError(t, f(ctx)) 66 }) 67 }