cloudflared-mirror/features/selector_test.go

129 lines
3.2 KiB
Go
Raw Normal View History

package features
import (
"context"
"encoding/json"
"fmt"
"testing"
"time"
"github.com/rs/zerolog"
"github.com/stretchr/testify/require"
)
func TestUnmarshalFeaturesRecord(t *testing.T) {
tests := []struct {
record []byte
expectedPercentage int32
expectedErr bool
}{
{
record: []byte(`{"pq":0}`),
expectedPercentage: 0,
},
{
record: []byte(`{"pq":39}`),
expectedPercentage: 39,
},
{
record: []byte(`{"pq":100}`),
expectedPercentage: 100,
},
{
record: []byte(`{}`),
expectedPercentage: 0, // Unmarshal to default struct if key is not present
},
{
record: []byte(`{"kyber":768}`),
expectedPercentage: 0, // Unmarshal to default struct if key is not present
},
{
record: []byte(`{"pq":"kyber768"}`),
expectedErr: true,
},
}
for _, test := range tests {
var features featuresRecord
err := json.Unmarshal(test.record, &features)
if test.expectedErr {
require.Error(t, err, test)
} else {
require.NoError(t, err)
require.Equal(t, test.expectedPercentage, features.PostQuantumPercentage, test)
}
}
}
func TestRefreshFeaturesRecord(t *testing.T) {
// The hash of the accountTag is 82
accountTag := t.Name()
threshold := switchThreshold(accountTag)
percentages := []int32{0, 10, 80, 83, 100}
refreshFreq := time.Millisecond * 10
selector := newTestSelector(t, percentages, nil, refreshFreq)
for _, percentage := range percentages {
if percentage > threshold {
require.Equal(t, PostQuantumPrefer, selector.PostQuantumMode())
} else {
require.Equal(t, PostQuantumDisabled, selector.PostQuantumMode())
}
time.Sleep(refreshFreq + time.Millisecond)
}
// Make sure error doesn't override the last fetched features
require.Equal(t, PostQuantumPrefer, selector.PostQuantumMode())
}
func TestStaticFeatures(t *testing.T) {
percentages := []int32{0}
pqMode := PostQuantumStrict
selector := newTestSelector(t, percentages, &pqMode, time.Millisecond*10)
require.Equal(t, PostQuantumStrict, selector.PostQuantumMode())
}
// Verify that if the first refresh fails, the selector will use default features
func TestFailedRefreshInitToDefault(t *testing.T) {
selector := newTestSelector(t, []int32{}, nil, time.Second)
require.Equal(t, featuresRecord{}, selector.features)
require.Equal(t, PostQuantumDisabled, selector.PostQuantumMode())
}
func newTestSelector(t *testing.T, percentages []int32, pqMode *PostQuantumMode, refreshFreq time.Duration) *FeatureSelector {
accountTag := t.Name()
logger := zerolog.Nop()
resolver := &mockResolver{
percentages: percentages,
}
staticFeatures := StaticFeatures{
PostQuantumMode: pqMode,
}
selector, err := newFeatureSelector(context.Background(), accountTag, &logger, resolver, staticFeatures, refreshFreq)
require.NoError(t, err)
return selector
}
type mockResolver struct {
nextIndex int
percentages []int32
}
func (mr *mockResolver) lookupRecord(ctx context.Context) ([]byte, error) {
if mr.nextIndex >= len(mr.percentages) {
return nil, fmt.Errorf("no more record to lookup")
}
record, err := json.Marshal(featuresRecord{
PostQuantumPercentage: mr.percentages[mr.nextIndex],
})
mr.nextIndex++
return record, err
}