fork of https://github.com/sourcegraph/zoekt
1package main
2
3import (
4 "context"
5 "path/filepath"
6 "testing"
7
8 "github.com/stretchr/testify/require"
9
10 "github.com/sourcegraph/zoekt"
11 "github.com/sourcegraph/zoekt/internal/tenant/tenanttest"
12)
13
14func TestPurgeTenantShards(t *testing.T) {
15 // TestPurgeTenantShards verifies both the basic shard purging functionality
16 // and proper isolation between tenants. It ensures that:
17 // 1. Shards are only purged when a valid tenant context is provided
18 // 2. Only shards belonging to the specified tenant are purged
19 // 3. Compound shards are preserved regardless of tenant
20 // 4. Other tenants' shards remain untouched
21 dir := t.TempDir()
22
23 // Create test shards for different tenants
24 tenant1Ctx := tenanttest.NewTestContext()
25 tenant2Ctx := tenanttest.NewTestContext()
26
27 // Helper to set tenant ID for test shards
28 setTenantID := func(id int) func(in *zoekt.Repository) {
29 return func(in *zoekt.Repository) {
30 in.TenantID = id
31 }
32 }
33
34 // Create test shards for tenant 1
35 tenant1Shard1 := filepath.Join(dir, "tenant1_repo1.zoekt")
36 tenant1Shard2 := filepath.Join(dir, "tenant1_repo2.zoekt")
37 createTestShard(t, "tenant1_repo1", 1, tenant1Shard1, setTenantID(1))
38 createTestShard(t, "tenant1_repo2", 2, tenant1Shard2, setTenantID(1))
39
40 // Create test shards for tenant 2
41 tenant2Shard := filepath.Join(dir, "tenant2_repo1.zoekt")
42 createTestShard(t, "tenant2_repo1", 3, tenant2Shard, setTenantID(2))
43
44 // Create a compound shard (should be skipped)
45 compoundShard := filepath.Join(dir, "compound-1234.zoekt")
46 createTestShard(t, "compound_repo", 4, compoundShard, setTenantID(1))
47
48 // Test cases
49 tests := []struct {
50 name string
51 ctx context.Context
52 wantErr bool
53 check func(t *testing.T, dir string)
54 }{
55 {
56 name: "no tenant in context",
57 ctx: context.Background(),
58 wantErr: true,
59 check: func(t *testing.T, dir string) {
60 // All files should still exist
61 require.FileExists(t, tenant1Shard1)
62 require.FileExists(t, tenant1Shard2)
63 require.FileExists(t, tenant2Shard)
64 require.FileExists(t, compoundShard)
65 },
66 },
67 {
68 name: "purge tenant 1 shards",
69 ctx: tenant1Ctx,
70 check: func(t *testing.T, dir string) {
71 // Tenant 1 shards should be deleted
72 require.NoFileExists(t, tenant1Shard1)
73 require.NoFileExists(t, tenant1Shard2)
74 // Other shards should still exist
75 require.FileExists(t, tenant2Shard)
76 require.FileExists(t, compoundShard)
77 },
78 },
79 {
80 name: "purge tenant 2 shards",
81 ctx: tenant2Ctx,
82 check: func(t *testing.T, dir string) {
83 // Tenant 2 shard should be deleted
84 require.NoFileExists(t, tenant2Shard)
85 // Compound shard should still exist
86 require.FileExists(t, compoundShard)
87 },
88 },
89 }
90
91 for _, tt := range tests {
92 t.Run(tt.name, func(t *testing.T) {
93 err := purgeTenantShards(tt.ctx, dir)
94 if tt.wantErr {
95 require.Error(t, err)
96 return
97 }
98 require.NoError(t, err)
99 tt.check(t, dir)
100 })
101 }
102}