teldrive/internal/reader/tg_reader_test.go

package reader

import (
	"context"
	"crypto/rand"
	"io"
	"testing"
	"time"

	"github.com/divyam234/teldrive/internal/config"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
)

type testChunkSource struct {
	buffer []byte
}

func (m *testChunkSource) Chunk(ctx context.Context, offset int64, limit int64) ([]byte, error) {
	return m.buffer[offset : offset+limit], nil
}

func (m *testChunkSource) ChunkSize(start, end int64) int64 {
	return 1
}

type testChunkSourceTimeout struct {
	buffer []byte
}

func (m *testChunkSourceTimeout) Chunk(ctx context.Context, offset int64, limit int64) ([]byte, error) {
	if offset == 8 {
		time.Sleep(2 * time.Second)
	}
	return m.buffer[offset : offset+limit], nil
}

func (m *testChunkSourceTimeout) ChunkSize(start, end int64) int64 {
	return 1
}

type TestSuite struct {
	suite.Suite
	config *config.TGConfig
}

func (suite *TestSuite) SetupTest() {
	suite.config = &config.TGConfig{Stream: struct {
		MultiThreads int
		Buffers      int
		ChunkTimeout time.Duration
	}{MultiThreads: 8, Buffers: 10, ChunkTimeout: 1 * time.Second}}
}

func (suite *TestSuite) TestFullRead() {
	ctx := context.Background()
	start := int64(0)
	end := int64(99)
	data := make([]byte, 100)
	rand.Read(data)
	chunkSrc := &testChunkSource{buffer: data}
	reader, err := newTGReader(ctx, start, end, suite.config, chunkSrc)
	assert.NoError(suite.T(), err)
	test_data, err := io.ReadAll(reader)
	assert.Equal(suite.T(), nil, err)
	assert.Equal(suite.T(), data[start:end+1], test_data)
}

func (suite *TestSuite) TestPartialRead() {
	ctx := context.Background()
	start := int64(0)
	end := int64(65)
	data := make([]byte, 100)
	rand.Read(data)
	chunkSrc := &testChunkSource{buffer: data}
	reader, err := newTGReader(ctx, start, end, suite.config, chunkSrc)
	assert.NoError(suite.T(), err)
	test_data, err := io.ReadAll(reader)
	assert.NoError(suite.T(), err)
	assert.Equal(suite.T(), data[start:end+1], test_data)
}

func (suite *TestSuite) TestTimeout() {
	ctx := context.Background()
	start := int64(0)
	end := int64(65)
	data := make([]byte, 100)
	rand.Read(data)
	chunkSrc := &testChunkSourceTimeout{buffer: data}
	reader, err := newTGReader(ctx, start, end, suite.config, chunkSrc)
	assert.NoError(suite.T(), err)
	test_data, err := io.ReadAll(reader)
	assert.Greater(suite.T(), len(test_data), 0)
	assert.Equal(suite.T(), err, ErrorStreamAbandoned)
}

func (suite *TestSuite) TestClose() {
	ctx := context.Background()
	start := int64(0)
	end := int64(65)
	data := make([]byte, 100)
	rand.Read(data)
	chunkSrc := &testChunkSource{buffer: data}
	reader, err := newTGReader(ctx, start, end, suite.config, chunkSrc)
	assert.NoError(suite.T(), err)
	_, err = io.ReadAll(reader)
	assert.NoError(suite.T(), err)
	assert.NoError(suite.T(), reader.Close())
}

func (suite *TestSuite) TestCancellation() {
	ctx, cancel := context.WithCancel(context.Background())
	start := int64(0)
	end := int64(65)
	data := make([]byte, 100)
	rand.Read(data)
	chunkSrc := &testChunkSource{buffer: data}
	reader, err := newTGReader(ctx, start, end, suite.config, chunkSrc)
	assert.NoError(suite.T(), err)
	cancel()
	_, err = io.ReadAll(reader)
	assert.Equal(suite.T(), err, context.Canceled)
	assert.Equal(suite.T(), len(reader.bufferChan), 0)
}

func (suite *TestSuite) TestCancellationWithTimeout() {
	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	_ = cancel
	start := int64(0)
	end := int64(65)
	data := make([]byte, 100)
	rand.Read(data)
	chunkSrc := &testChunkSourceTimeout{buffer: data}
	reader, err := newTGReader(ctx, start, end, suite.config, chunkSrc)
	assert.NoError(suite.T(), err)
	_, err = io.ReadAll(reader)
	assert.Equal(suite.T(), err, context.DeadlineExceeded)
	assert.Equal(suite.T(), len(reader.bufferChan), 0)
}
func Test(t *testing.T) {
	suite.Run(t, new(TestSuite))
}