You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
413 lines
9.9 KiB
Go
413 lines
9.9 KiB
Go
package chunking
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/gzip"
|
|
"crypto/rand"
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/rqlite/rqlite/v8/command/proto"
|
|
)
|
|
|
|
func Test_SingleChunk(t *testing.T) {
|
|
data := []byte("Hello, World!")
|
|
chunk := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 1,
|
|
IsLast: true,
|
|
Data: mustCompressData(data),
|
|
}
|
|
|
|
dir, err := os.MkdirTemp("", "dechunker-test")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(dir)
|
|
|
|
dechunker, err := NewDechunker(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create Dechunker: %v", err)
|
|
}
|
|
|
|
isLast, err := dechunker.WriteChunk(chunk)
|
|
if err != nil {
|
|
t.Fatalf("failed to write chunk: %v", err)
|
|
}
|
|
if !isLast {
|
|
t.Errorf("WriteChunk did not return true for isLast")
|
|
}
|
|
|
|
filePath, err := dechunker.Close()
|
|
if err != nil {
|
|
t.Fatalf("failed to close Dechunker: %v", err)
|
|
}
|
|
|
|
// Check the contents of the output file.
|
|
got, err := os.ReadFile(filePath)
|
|
if err != nil {
|
|
t.Fatalf("failed to read output file: %v", err)
|
|
}
|
|
if string(got) != string(data) {
|
|
t.Errorf("output file data = %q; want %q", got, data)
|
|
}
|
|
}
|
|
|
|
func Test_MultiChunk(t *testing.T) {
|
|
data1 := []byte("Hello, World!")
|
|
chunk1 := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 1,
|
|
IsLast: false,
|
|
Data: mustCompressData(data1),
|
|
}
|
|
data2 := []byte("I'm OK")
|
|
chunk2 := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 2,
|
|
IsLast: true,
|
|
Data: mustCompressData(data2),
|
|
}
|
|
|
|
dir, err := os.MkdirTemp("", "dechunker-test")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(dir)
|
|
|
|
dechunker, err := NewDechunker(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create Dechunker: %v", err)
|
|
}
|
|
|
|
for _, chunk := range []*proto.LoadChunkRequest{chunk1, chunk2} {
|
|
isLast, err := dechunker.WriteChunk(chunk)
|
|
if err != nil {
|
|
t.Fatalf("failed to write chunk: %v", err)
|
|
}
|
|
if isLast != chunk.IsLast {
|
|
t.Errorf("WriteChunk returned wrong value for isLast")
|
|
}
|
|
}
|
|
|
|
filePath, err := dechunker.Close()
|
|
if err != nil {
|
|
t.Fatalf("failed to close Dechunker: %v", err)
|
|
}
|
|
|
|
// Check the contents of the output file.
|
|
got, err := os.ReadFile(filePath)
|
|
if err != nil {
|
|
t.Fatalf("failed to read output file: %v", err)
|
|
}
|
|
if string(got) != string(data1)+string(data2) {
|
|
t.Errorf("output file data = %q; want %q", got, string(data1)+string(data2))
|
|
}
|
|
}
|
|
|
|
func Test_MultiChunkNilData(t *testing.T) {
|
|
data1 := []byte("Hello, World!")
|
|
chunk1 := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 1,
|
|
IsLast: false,
|
|
Data: mustCompressData(data1),
|
|
}
|
|
data2 := []byte("I'm OK")
|
|
chunk2 := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 2,
|
|
IsLast: false,
|
|
Data: mustCompressData(data2),
|
|
}
|
|
chunk3 := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 3,
|
|
IsLast: true,
|
|
Data: nil,
|
|
}
|
|
|
|
dir, err := os.MkdirTemp("", "dechunker-test")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(dir)
|
|
|
|
dechunker, err := NewDechunker(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create Dechunker: %v", err)
|
|
}
|
|
|
|
for _, chunk := range []*proto.LoadChunkRequest{chunk1, chunk2, chunk3} {
|
|
isLast, err := dechunker.WriteChunk(chunk)
|
|
if err != nil {
|
|
t.Fatalf("failed to write chunk: %v", err)
|
|
}
|
|
if isLast != chunk.IsLast {
|
|
t.Errorf("WriteChunk returned wrong value for isLast")
|
|
}
|
|
}
|
|
|
|
filePath, err := dechunker.Close()
|
|
if err != nil {
|
|
t.Fatalf("failed to close Dechunker: %v", err)
|
|
}
|
|
|
|
// Check the contents of the output file.
|
|
got, err := os.ReadFile(filePath)
|
|
if err != nil {
|
|
t.Fatalf("failed to read output file: %v", err)
|
|
}
|
|
if string(got) != string(data1)+string(data2) {
|
|
t.Errorf("output file data = %q; want %q", got, string(data1)+string(data2))
|
|
}
|
|
}
|
|
|
|
func Test_UnexpectedStreamID(t *testing.T) {
|
|
originalData := []byte("Hello, World!")
|
|
compressedData := mustCompressData(originalData)
|
|
|
|
chunk1 := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 1,
|
|
IsLast: false,
|
|
Data: compressedData,
|
|
}
|
|
|
|
chunk2 := &proto.LoadChunkRequest{
|
|
StreamId: "456",
|
|
SequenceNum: 2,
|
|
IsLast: true,
|
|
Data: compressedData,
|
|
}
|
|
|
|
dir, err := os.MkdirTemp("", "dechunker-test")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(dir)
|
|
|
|
dechunker, err := NewDechunker(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create Dechunker: %v", err)
|
|
}
|
|
|
|
_, err = dechunker.WriteChunk(chunk1)
|
|
if err != nil {
|
|
t.Fatalf("failed to write first chunk: %v", err)
|
|
}
|
|
|
|
// Write the second chunk with a different stream ID, which should result in an error.
|
|
_, err = dechunker.WriteChunk(chunk2)
|
|
if err == nil {
|
|
t.Fatalf("expected error when writing chunk with different stream ID, but got nil")
|
|
}
|
|
if !strings.Contains(err.Error(), "unexpected stream ID") {
|
|
t.Errorf("error = %v; want an error about unexpected stream ID", err)
|
|
}
|
|
}
|
|
|
|
func Test_ChunksOutOfOrder(t *testing.T) {
|
|
originalData := []byte("Hello, World!")
|
|
compressedData := mustCompressData(originalData)
|
|
|
|
chunk1 := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 1,
|
|
IsLast: false,
|
|
Data: compressedData,
|
|
}
|
|
|
|
chunk3 := &proto.LoadChunkRequest{
|
|
StreamId: "123",
|
|
SequenceNum: 3,
|
|
IsLast: true,
|
|
Data: compressedData,
|
|
}
|
|
|
|
dir, err := os.MkdirTemp("", "dechunker-test")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(dir)
|
|
|
|
dechunker, err := NewDechunker(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create Dechunker: %v", err)
|
|
}
|
|
|
|
// Write the first chunk to the Dechunker.
|
|
_, err = dechunker.WriteChunk(chunk1)
|
|
if err != nil {
|
|
t.Fatalf("failed to write first chunk: %v", err)
|
|
}
|
|
|
|
// Write the third chunk (skipping the second), which should result in an error.
|
|
_, err = dechunker.WriteChunk(chunk3)
|
|
if err == nil {
|
|
t.Fatalf("expected error when writing chunks out of order, but got nil")
|
|
}
|
|
if !strings.Contains(err.Error(), "chunks received out of order") {
|
|
t.Errorf("error = %v; want an error about chunks received out of order", err)
|
|
}
|
|
}
|
|
|
|
func Test_ReassemblyOfLargeData(t *testing.T) {
|
|
dir, err := os.MkdirTemp("", "test-*")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
|
|
d, err := NewDechunker(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create Dechunker: %v", err)
|
|
}
|
|
defer os.Remove(dir)
|
|
|
|
// Create a large random dataset.
|
|
largeData := make([]byte, 2*1024*1024) // 2 MB of data.
|
|
if _, err := rand.Read(largeData); err != nil {
|
|
t.Fatalf("failed to generate large data: %v", err)
|
|
}
|
|
|
|
// Split the large data into chunks.
|
|
numChunks := 16
|
|
chunkSize := len(largeData) / numChunks
|
|
|
|
// Write the chunks to the Dechunker.
|
|
for i := 0; i < numChunks; i++ {
|
|
start := i * chunkSize
|
|
end := start + chunkSize
|
|
|
|
isLast := i == numChunks-1
|
|
if _, err := d.WriteChunk(&proto.LoadChunkRequest{
|
|
StreamId: "1",
|
|
SequenceNum: int64(i + 1),
|
|
IsLast: isLast,
|
|
Data: mustCompressData(largeData[start:end]),
|
|
}); err != nil {
|
|
t.Fatalf("failed to write chunk: %v", err)
|
|
}
|
|
}
|
|
|
|
outFilePath, err := d.Close()
|
|
if err != nil {
|
|
t.Fatalf("failed to close Dechunker: %v", err)
|
|
}
|
|
defer os.Remove(outFilePath)
|
|
|
|
// The output data should be the same as the original largeData.
|
|
outData, err := os.ReadFile(outFilePath)
|
|
if err != nil {
|
|
t.Fatalf("failed to read output file data: %v", err)
|
|
}
|
|
if !bytes.Equal(outData, largeData) {
|
|
t.Fatalf("output file data does not match original data")
|
|
}
|
|
}
|
|
|
|
func Test_CreateDechunkerManager(t *testing.T) {
|
|
dir := os.TempDir()
|
|
manager, err := NewDechunkerManager(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create DechunkerManager: %v", err)
|
|
}
|
|
defer manager.Close()
|
|
if manager == nil {
|
|
t.Fatalf("expected DechunkerManager instance, got nil")
|
|
}
|
|
if manager.dir != dir {
|
|
t.Errorf("unexpected dir in manager: expected %s, got %s", dir, manager.dir)
|
|
}
|
|
if manager.m != nil {
|
|
t.Errorf("expected manager map to be nil, got: %v", manager.m)
|
|
}
|
|
}
|
|
|
|
func Test_CreateDechunkerManagerFail(t *testing.T) {
|
|
manager, err := NewDechunkerManager("/does/not/exist/surely")
|
|
if err == nil {
|
|
t.Fatalf("expected an error but got none")
|
|
}
|
|
if manager != nil {
|
|
t.Fatalf("expected nil manager but got one")
|
|
}
|
|
}
|
|
|
|
func Test_GetDechunker(t *testing.T) {
|
|
dir := os.TempDir()
|
|
manager, err := NewDechunkerManager(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create DechunkerManager: %v", err)
|
|
}
|
|
defer manager.Close()
|
|
|
|
dechunker1, err := manager.Get("test1")
|
|
if err != nil {
|
|
t.Fatalf("failed to get dechunker: %v", err)
|
|
}
|
|
if dechunker1 == nil {
|
|
t.Fatalf("expected Dechunker instance, got nil")
|
|
}
|
|
|
|
// Make sure we get the same dechunker.
|
|
dechunker2, err := manager.Get("test1")
|
|
if err != nil {
|
|
t.Fatalf("failed to get dechunker: %v", err)
|
|
}
|
|
if dechunker2 == nil {
|
|
t.Fatalf("expected Dechunker instance, got nil")
|
|
}
|
|
if dechunker1 != dechunker2 {
|
|
t.Errorf("expected same dechunker instance, got different instances")
|
|
}
|
|
}
|
|
|
|
func Test_DeleteDechunker(t *testing.T) {
|
|
dir := os.TempDir()
|
|
manager, err := NewDechunkerManager(dir)
|
|
if err != nil {
|
|
t.Fatalf("failed to create DechunkerManager: %v", err)
|
|
}
|
|
defer manager.Close()
|
|
|
|
dechunker1, err := manager.Get("test1")
|
|
if err != nil {
|
|
t.Fatalf("failed to get dechunker: %v", err)
|
|
}
|
|
if dechunker1 == nil {
|
|
t.Fatalf("expected Dechunker instance, got nil")
|
|
}
|
|
|
|
manager.Delete("test1")
|
|
|
|
// Attempt to get the Dechunker for the ID "test1"
|
|
dechunker2, err := manager.Get("test1")
|
|
if err != nil {
|
|
t.Fatalf("failed to get dechunker: %v", err)
|
|
}
|
|
if dechunker2 == nil {
|
|
t.Fatalf("expected Dechunker instance, got nil")
|
|
}
|
|
if dechunker1 == dechunker2 {
|
|
t.Errorf("expected different dechunker instances, got same instance")
|
|
}
|
|
}
|
|
|
|
func mustCompressData(data []byte) []byte {
|
|
var buf bytes.Buffer
|
|
gz := gzip.NewWriter(&buf)
|
|
|
|
if _, err := gz.Write(data); err != nil {
|
|
panic(fmt.Sprintf("failed to write compressed data: %v", err))
|
|
}
|
|
|
|
if err := gz.Close(); err != nil {
|
|
panic(fmt.Sprintf("failed to close gzip writer: %v", err))
|
|
}
|
|
|
|
return buf.Bytes()
|
|
}
|