From 479c4d8c444255d2ef2729aff60984fc2c14a306 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Leszczy=C5=84ski?= <2000michal@wp.pl> Date: Thu, 3 Oct 2024 17:36:12 +0200 Subject: [PATCH] feat(restore): batch, add test for simple batching scenario --- pkg/service/restore/batch_test.go | 128 ++++++++++++++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 pkg/service/restore/batch_test.go diff --git a/pkg/service/restore/batch_test.go b/pkg/service/restore/batch_test.go new file mode 100644 index 000000000..413b1baa7 --- /dev/null +++ b/pkg/service/restore/batch_test.go @@ -0,0 +1,128 @@ +// Copyright (C) 2024 ScyllaDB + +package restore + +import ( + "testing" + + "github.com/scylladb/scylla-manager/v3/pkg/service/backup/backupspec" +) + +func TestBatchDispatcher(t *testing.T) { + l1 := backupspec.Location{ + Provider: "s3", + Path: "l1", + } + l2 := backupspec.Location{ + Provider: "s3", + Path: "l2", + } + workload := []LocationWorkload{ + { + Location: l1, + Size: 160, + Tables: []TableWorkload{ + { + Size: 50, + RemoteDirs: []RemoteDirWorkload{ + { + RemoteSSTableDir: "a", + Size: 20, + SSTables: []RemoteSSTable{ + {Size: 5}, + {Size: 15}, + }, + }, + { + RemoteSSTableDir: "b", + Size: 30, + SSTables: []RemoteSSTable{ + {Size: 10}, + {Size: 20}, + }, + }, + }, + }, + { + Size: 110, + RemoteDirs: []RemoteDirWorkload{ + { + RemoteSSTableDir: "c", + Size: 110, + SSTables: []RemoteSSTable{ + {Size: 50}, + {Size: 60}, + }, + }, + }, + }, + }, + }, + { + Location: l2, + Size: 200, + Tables: []TableWorkload{ + { + Size: 200, + RemoteDirs: []RemoteDirWorkload{ + { + RemoteSSTableDir: "d", + Size: 200, + SSTables: []RemoteSSTable{ + {Size: 110}, + {Size: 90}, + }, + }, + }, + }, + }, + }, + } + locationHosts := map[backupspec.Location][]string{ + l1: {"h1", "h2"}, + l2: {"h3"}, + } + hostToShard := map[string]uint{ + "h1": 1, + "h2": 2, + "h3": 3, + } + + bd := newBatchDispatcher(workload, 1, hostToShard, locationHosts) + + scenario := []struct { + host string + ok bool + dir string + size int64 + count int + }{ + {host: "h1", ok: true, dir: "c", size: 60, count: 1}, + {host: "h1", ok: true, dir: "c", size: 50, count: 1}, + {host: "h2", ok: true, dir: "b", size: 30, count: 2}, + {host: "h3", ok: true, dir: "d", size: 200, count: 2}, + {host: "h3", ok: false}, + {host: "h2", ok: true, dir: "a", size: 20, count: 2}, + {host: "h1", ok: false}, + {host: "h2", ok: false}, + } + + for _, step := range scenario { + b, ok := bd.DispatchBatch(step.host) + if ok != step.ok { + t.Fatalf("Step: %+v, expected ok=%v, got ok=%v", step, step.ok, ok) + } + if ok == false { + continue + } + if b.RemoteSSTableDir != step.dir { + t.Fatalf("Step: %+v, expected dir=%v, got dir=%v", step, step.dir, b.RemoteSSTableDir) + } + if b.Size != step.size { + t.Fatalf("Step: %+v, expected size=%v, got size=%v", step, step.size, b.Size) + } + if len(b.SSTables) != step.count { + t.Fatalf("Step: %+v, expected count=%v, got count=%v", step, step.count, len(b.SSTables)) + } + } +}