<commit_before>package server

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/privacylab/talek/common"
	_ "github.com/privacylab/talek/pir/pircpu"
)

import "testing"

func fromEnvOrDefault(envKey string, defaultVal int) int {
	if os.Getenv(envKey) != "" {
		val, _ := strconv.Atoi(os.Getenv(envKey))
		return val
	}
	return defaultVal
}

func testConf() Config {
	return Config{
		Config: &common.Config{
			NumBuckets:         uint64(fromEnvOrDefault("NUM_BUCKETS", 512)),
			BucketDepth:        uint64(fromEnvOrDefault("BUCKET_DEPTH", 4)),
			DataSize:           uint64(fromEnvOrDefault("DATA_SIZE", 512)),
			BloomFalsePositive: 0.95,
			MaxLoadFactor:      0.95,
			LoadFactorStep:     0.02,
		},
		ReadBatch:        fromEnvOrDefault("BATCH_SIZE", 8),
		WriteInterval:    time.Second,
		ReadInterval:     time.Second,
		TrustDomain:      nil,
		TrustDomainIndex: 0,
	}
}

func TestShardSanity(t *testing.T) {
	shard := NewShard("Test Shard", "cpu.0", testConf())
	if shard == nil {
		t.Error("Failed to create shard.")
		return
	}

	writeReplyChan := make(chan *common.WriteReply)

	data := make([]byte, conf.Config.DataSize)
	copy(data, bytes.NewBufferString("Magic").Bytes())
	shard.Write(&common.ReplicaWriteArgs{
		WriteArgs: common.WriteArgs{
			Bucket1:        0,
			Bucket2:        1,
			Data:           data,
			InterestVector: []byte{},
			ReplyChan:      writeReplyChan,
		},
		EpochFlag: false,
	})
<commit_msg>fix size issue in shard_test<commit_after>package server

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/privacylab/talek/common"
	_ "github.com/privacylab/talek/pir/pircpu"
)

import "testing"

func fromEnvOrDefault(envKey string, defaultVal int) int {
	if os.Getenv(envKey) != "" {
		val, _ := strconv.Atoi(os.Getenv(envKey))
		return val
	}
	return defaultVal
}

func testConf() Config {
	return Config{
		Config: &common.Config{
			NumBuckets:         uint64(fromEnvOrDefault("NUM_BUCKETS", 512)),
			BucketDepth:        uint64(fromEnvOrDefault("BUCKET_DEPTH", 4)),
			DataSize:           uint64(fromEnvOrDefault("DATA_SIZE", 512)),
			BloomFalsePositive: 0.95,
			MaxLoadFactor:      0.95,
			LoadFactorStep:     0.02,
		},
		ReadBatch:        fromEnvOrDefault("BATCH_SIZE", 8),
		WriteInterval:    time.Second,
		ReadInterval:     time.Second,
		TrustDomain:      nil,
		TrustDomainIndex: 0,
	}
}

func TestShardSanity(t *testing.T) {
	conf := testConf()
	shard := NewShard("Test Shard", "cpu.0", conf)
	if shard == nil {
		t.Error("Failed to create shard.")
		return
	}

	writeReplyChan := make(chan *common.WriteReply)

	data := make([]byte, conf.Config.DataSize)
	copy(data, bytes.NewBufferString("Magic").Bytes())
	shard.Write(&common.ReplicaWriteArgs{
		WriteArgs: common.WriteArgs{
			Bucket1:        0,
			Bucket2:        1,
			Data:           data,
			InterestVector: []byte{},
			ReplyChan:      writeReplyChan,
		},
		EpochFlag: false,
	})

{'fraction_non_alphanumeric': 0.13863404689092762, 'fraction_numerical': 0.019707781175671082, 'mean_word_length': 3.453857791225416, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7017741', 'n_tokens_mistral': 1210, 'n_tokens_neox': 1081, 'n_words': 204}