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