Files
hse-2026/homework/03-http-server/tests/query_gen.go
T
2026-09-24 21:30:37 +03:00

147 lines
4.2 KiB
Go

package hw3test
import (
"math/rand"
"strings"
"github.com/stretchr/testify/require"
)
// LargePutPolicy keeps a small number of guaranteed successful writes above
// the solution's memory limit, while bounding the remaining PUT bodies.
type LargePutPolicy struct {
RunSeed int64
LargeSizes []int64
OtherMaxBytes int64
}
// QueriesGen contains config for generating queries.
type QueriesGen struct {
// Number of queries to generate.
Count int
GetFile bool // Allow GET requests for files.
GetFileNoErrors bool // Disable requests for non-existing files.
GetDirectory bool // Allow GET requests for directories.
Compression bool // Allow compression in GET requests.
Post bool // Allow POST requests.
Put bool // Allow PUT requests.
Delete bool // Allow DELETE requests.
AllHeaders bool // Allow extra headers.
LargePuts *LargePutPolicy
}
func (g *QueriesGen) Generate(t *TC, env *Env, gen *EnvGen, seeds []int64, opts RunOpts, runSeed int64) []Query {
env = env.Clone()
largePuts := 0
stats := &Stats{}
env.RootDir.Stats("", stats)
stats.Normalize()
var methods []string
if g.GetFile || g.GetDirectory {
methods = append(methods, "GET")
}
if g.Post {
methods = append(methods, "POST")
}
if g.Put {
methods = append(methods, "PUT")
}
if g.Delete {
methods = append(methods, "DELETE")
}
var queries []Query
for _, seed := range seeds {
r := rand.New(rand.NewSource(seed))
method := methods[r.Intn(len(methods))]
var genPaths []string
if len(stats.DirPaths) > 0 && !(method == "GET" && !g.GetDirectory) {
genPaths = append(genPaths, stats.DirPaths[r.Intn(len(stats.DirPaths))])
genPaths = append(genPaths, stats.DirPaths[r.Intn(len(stats.DirPaths))])
}
if len(stats.FilePaths) > 0 && !(method == "GET" && !g.GetFile) {
genPaths = append(genPaths, stats.FilePaths[r.Intn(len(stats.FilePaths))])
genPaths = append(genPaths, stats.FilePaths[r.Intn(len(stats.FilePaths))])
}
if !(method == "GET" && g.GetFileNoErrors) {
// TODO: better non-existing path generator
randomPath := gen.FilenameGen(r) + "/" + gen.FilenameGen(r)
genPaths = append(genPaths, randomPath)
if len(stats.DirPaths) > 0 {
randomDir := stats.DirPaths[r.Intn(len(stats.DirPaths))]
num := 2
if method == "POST" {
num = 4
}
for i := 0; i < num; i++ {
genPaths = append(genPaths, randomDir+"/"+gen.FilenameGen(r))
}
}
}
if len(genPaths) == 0 {
continue
}
hostHeader := opts.ServerDomain
if g.AllHeaders && r.Intn(5) == 0 {
// to get 400
hostHeader = "hse.ru"
}
path := strings.TrimPrefix(genPaths[r.Intn(len(genPaths))], "/")
if path == "" && method != "GET" {
// Keep the existing large-file scenarios and RNG sequence.
// Requests targeting the root are covered on a small tree in contract.go.
path = "MISSING_ROOT_TARGET"
}
gzipRequested := g.Compression && r.Intn(2) == 1
query := Query{
Seed: seed,
Method: method,
Path: path,
Gzip: method == "GET" && gzipRequested,
CreateDirectory: method == "POST" && (r.Intn(2) == 1),
RemoveDirectory: method == "DELETE" && (r.Intn(2) == 1),
HostHeader: hostHeader,
VerifyHeaders: g.AllHeaders,
}
if (method == "POST" && !query.CreateDirectory) || method == "PUT" {
query.FileContent = gen.GenerateFile(r)
}
action := query.Action(env, &opts)
if query.Method == "PUT" && g.LargePuts != nil {
policy := g.LargePuts
if _, ok := action.(ReplaceFileAction); ok && runSeed == policy.RunSeed && largePuts < len(policy.LargeSizes) {
query.FileContent.Size = policy.LargeSizes[largePuts]
largePuts++
} else if query.FileContent.Size > policy.OtherMaxBytes {
query.FileContent.Size = policy.OtherMaxBytes
}
}
queries = append(queries, query)
if action != nil {
changed := action.ApplyEnv(t, env)
if changed {
stats = &Stats{}
env.RootDir.Stats("", stats)
stats.Normalize()
}
}
}
if g.LargePuts != nil && runSeed == g.LargePuts.RunSeed {
require.Equal(t, len(g.LargePuts.LargeSizes), largePuts, "not enough successful PUT queries for large-file coverage")
}
return queries
}