mirror of
https://github.com/siderolabs/talos.git
synced 2025-08-23 07:31:13 +02:00
e35a910da2
Stream chunker should be cancellable at any point of execution, plus it should be stop chunking on EOF. Signed-off-by: Andrey Smirnov <smirnov.andrey@gmail.com>
136 lines
3.1 KiB
Go
136 lines
3.1 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 stream_test
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/talos-systems/talos/internal/pkg/chunker/stream"
|
|
)
|
|
|
|
type StreamChunkerSuite struct {
|
|
suite.Suite
|
|
|
|
reader *io.PipeReader
|
|
writer *io.PipeWriter
|
|
}
|
|
|
|
func (suite *StreamChunkerSuite) SetupTest() {
|
|
suite.reader, suite.writer = io.Pipe()
|
|
}
|
|
|
|
func (suite *StreamChunkerSuite) TearDownTest() {
|
|
suite.Require().NoError(suite.writer.Close())
|
|
suite.Require().NoError(suite.reader.Close())
|
|
}
|
|
|
|
func collectChunks(chunksCh <-chan []byte) <-chan []byte {
|
|
combinedCh := make(chan []byte)
|
|
|
|
go func() {
|
|
res := []byte(nil)
|
|
|
|
for chunk := range chunksCh {
|
|
res = append(res, chunk...)
|
|
}
|
|
|
|
combinedCh <- res
|
|
}()
|
|
|
|
return combinedCh
|
|
}
|
|
|
|
func (suite *StreamChunkerSuite) TestStreaming() {
|
|
chunker := stream.NewChunker(suite.reader)
|
|
|
|
ctx, ctxCancel := context.WithCancel(context.Background())
|
|
defer ctxCancel()
|
|
|
|
chunksCh := chunker.Read(ctx)
|
|
combinedCh := collectChunks(chunksCh)
|
|
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("abc"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("def"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("ghi"))
|
|
time.Sleep(50 * time.Millisecond)
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("jkl"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("mno"))
|
|
|
|
suite.Require().NoError(suite.writer.Close())
|
|
|
|
suite.Require().Equal([]byte("abcdefghijklmno"), <-combinedCh)
|
|
}
|
|
|
|
func (suite *StreamChunkerSuite) TestStreamingSmallBuf() {
|
|
chunker := stream.NewChunker(suite.reader, stream.Size(1))
|
|
|
|
ctx, ctxCancel := context.WithCancel(context.Background())
|
|
defer ctxCancel()
|
|
|
|
chunksCh := chunker.Read(ctx)
|
|
combinedCh := collectChunks(chunksCh)
|
|
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("abc"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("def"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("ghi"))
|
|
time.Sleep(50 * time.Millisecond)
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("jkl"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("mno"))
|
|
|
|
suite.Require().NoError(suite.writer.Close())
|
|
|
|
suite.Require().Equal([]byte("abcdefghijklmno"), <-combinedCh)
|
|
}
|
|
|
|
func (suite *StreamChunkerSuite) TestStreamingCancel() {
|
|
chunker := stream.NewChunker(suite.reader)
|
|
|
|
ctx, ctxCancel := context.WithCancel(context.Background())
|
|
defer ctxCancel()
|
|
|
|
chunksCh := chunker.Read(ctx)
|
|
combinedCh := collectChunks(chunksCh)
|
|
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("abc"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("def"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("ghi"))
|
|
time.Sleep(50 * time.Millisecond)
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("jkl"))
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte("mno"))
|
|
time.Sleep(50 * time.Millisecond)
|
|
|
|
ctxCancel()
|
|
|
|
// need any I/O for chunker to notice that context got canceled
|
|
// nolint: errcheck
|
|
suite.writer.Write([]byte(""))
|
|
|
|
suite.Require().Equal([]byte("abcdefghijklmno"), <-combinedCh)
|
|
}
|
|
|
|
func TestStreamChunkerSuite(t *testing.T) {
|
|
suite.Run(t, new(StreamChunkerSuite))
|
|
}
|