fork of https://github.com/sourcegraph/zoekt
1package tenanttest
2
3import (
4 "context"
5 "testing"
6
7 "go.uber.org/atomic"
8
9 "github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
10 "github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype"
11)
12
13func MockEnforce(t *testing.T) {
14 // prevent parallel tests from interfering with each other
15 t.Setenv("mockEnforce", "true")
16
17 old := enforcement.EnforcementMode.Load()
18 t.Cleanup(func() {
19 enforcement.EnforcementMode.Store(old)
20 ResetTestTenants()
21 })
22
23 enforcement.EnforcementMode.Store("strict")
24}
25
26// TestTenantCounter is a counter that is tracks tenants created from NewTestContext().
27var TestTenantCounter atomic.Int64
28
29func NewTestContext() context.Context {
30 return tenanttype.WithTenant(context.Background(), mustTenantFromID(int(TestTenantCounter.Inc())))
31}
32
33// ResetTestTenants resets the test tenant counter that tracks the tenants
34// created from NewTestContext().
35func ResetTestTenants() {
36 TestTenantCounter.Store(0)
37}
38
39func mustTenantFromID(id int) *tenanttype.Tenant {
40 tenant, err := tenanttype.FromID(id)
41 if err != nil {
42 panic(err)
43 }
44 return tenant
45}