mirror of
https://github.com/siderolabs/talos.git
synced 2025-09-14 10:21:11 +02:00
3257751bc0
Because of the bug Drainer never worked properly, as `shutdown` channel wasn't initialied. Also add unit-tests and add some small clean-ups which don't affect functionality. Signed-off-by: Andrey Smirnov <andrey.smirnov@talos-systems.com>
100 lines
1.8 KiB
Go
100 lines
1.8 KiB
Go
// This Source Code Form is subject to the terms of the Mozilla Public
|
|
// License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
|
|
package runtime_test
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/talos-systems/talos/internal/app/machined/pkg/runtime"
|
|
)
|
|
|
|
//nolint:gocyclo
|
|
func TestDrainer(t *testing.T) {
|
|
drainer := runtime.NewDrainer()
|
|
|
|
sub1 := drainer.Subscribe()
|
|
sub2 := drainer.Subscribe()
|
|
|
|
errCh := make(chan error)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
go func() {
|
|
errCh <- drainer.Drain(ctx)
|
|
}()
|
|
|
|
select {
|
|
case <-sub1.EventCh():
|
|
case <-time.After(time.Second):
|
|
require.Fail(t, "should be notified")
|
|
}
|
|
|
|
select {
|
|
case <-sub2.EventCh():
|
|
case <-time.After(time.Second):
|
|
require.Fail(t, "should be notified")
|
|
}
|
|
|
|
select {
|
|
case <-errCh:
|
|
require.Fail(t, "shouldn't be drained now")
|
|
default:
|
|
}
|
|
|
|
sub1.Cancel()
|
|
|
|
select {
|
|
case <-errCh:
|
|
require.Fail(t, "shouldn't be drained now")
|
|
default:
|
|
}
|
|
|
|
sub3 := drainer.Subscribe()
|
|
|
|
select {
|
|
case <-sub3.EventCh():
|
|
case <-time.After(time.Second):
|
|
require.Fail(t, "should be notified")
|
|
}
|
|
|
|
sub3.Cancel()
|
|
sub2.Cancel()
|
|
|
|
select {
|
|
case err := <-errCh:
|
|
assert.NoError(t, err)
|
|
case <-time.After(time.Second):
|
|
require.Fail(t, "should be drained now")
|
|
}
|
|
}
|
|
|
|
func TestDrainTimeout(t *testing.T) {
|
|
drainer := runtime.NewDrainer()
|
|
|
|
drainer.Subscribe()
|
|
|
|
errCh := make(chan error)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
|
|
go func() {
|
|
errCh <- drainer.Drain(ctx)
|
|
}()
|
|
|
|
select {
|
|
case err := <-errCh:
|
|
assert.ErrorIs(t, err, context.DeadlineExceeded)
|
|
case <-time.After(5 * time.Second):
|
|
require.Fail(t, "should be drained now")
|
|
}
|
|
}
|