109 lines
3.2 KiB
Go
109 lines
3.2 KiB
Go
package statesim
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/btree"
|
|
"pgregory.net/rapid"
|
|
|
|
"cosmossdk.io/schema"
|
|
"cosmossdk.io/schema/appdata"
|
|
)
|
|
|
|
// App is a collection of simulated module states corresponding to an app's schema for testing purposes.
|
|
type App struct {
|
|
options Options
|
|
moduleStates *btree.Map[string, *Module]
|
|
updateGen *rapid.Generator[appdata.ObjectUpdateData]
|
|
}
|
|
|
|
// NewApp creates a new simulation App for the given app schema. The app schema can be nil
|
|
// if the user desires initializing modules with InitializeModule instead.
|
|
func NewApp(appSchema map[string]schema.ModuleSchema, options Options) *App {
|
|
app := &App{
|
|
options: options,
|
|
moduleStates: &btree.Map[string, *Module]{},
|
|
}
|
|
|
|
for moduleName, moduleSchema := range appSchema {
|
|
moduleState := NewModule(moduleSchema, options)
|
|
app.moduleStates.Set(moduleName, moduleState)
|
|
}
|
|
|
|
moduleNameSelector := rapid.Custom(func(t *rapid.T) string {
|
|
return rapid.SampledFrom(app.moduleStates.Keys()).Draw(t, "moduleName")
|
|
})
|
|
|
|
numUpdatesGen := rapid.IntRange(1, 2)
|
|
app.updateGen = rapid.Custom(func(t *rapid.T) appdata.ObjectUpdateData {
|
|
moduleName := moduleNameSelector.Draw(t, "moduleName")
|
|
moduleState, ok := app.moduleStates.Get(moduleName)
|
|
require.True(t, ok)
|
|
numUpdates := numUpdatesGen.Draw(t, "numUpdates")
|
|
updates := make([]schema.ObjectUpdate, numUpdates)
|
|
for i := 0; i < numUpdates; i++ {
|
|
update := moduleState.UpdateGen().Draw(t, fmt.Sprintf("update[%d]", i))
|
|
updates[i] = update
|
|
}
|
|
return appdata.ObjectUpdateData{
|
|
ModuleName: moduleName,
|
|
Updates: updates,
|
|
}
|
|
})
|
|
|
|
return app
|
|
}
|
|
|
|
// InitializeModule initializes the module with the provided schema. This returns an error if the
|
|
// module is already initialized in state.
|
|
func (a *App) InitializeModule(data appdata.ModuleInitializationData) error {
|
|
if _, ok := a.moduleStates.Get(data.ModuleName); ok {
|
|
return fmt.Errorf("module %s already initialized", data.ModuleName)
|
|
}
|
|
|
|
a.moduleStates.Set(data.ModuleName, NewModule(data.Schema, a.options))
|
|
return nil
|
|
}
|
|
|
|
// ApplyUpdate applies the given object update to the module.
|
|
func (a *App) ApplyUpdate(data appdata.ObjectUpdateData) error {
|
|
moduleState, ok := a.moduleStates.Get(data.ModuleName)
|
|
if !ok {
|
|
// we don't have this module so skip the update
|
|
return nil
|
|
}
|
|
|
|
for _, update := range data.Updates {
|
|
err := moduleState.ApplyUpdate(update)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// UpdateGen is a generator for object update data against the app. It is stateful and includes a certain number of
|
|
// updates and deletions to existing objects.
|
|
func (a *App) UpdateGen() *rapid.Generator[appdata.ObjectUpdateData] {
|
|
return a.updateGen
|
|
}
|
|
|
|
// GetModule returns the module state for the given module name.
|
|
func (a *App) GetModule(moduleName string) (ModuleState, bool) {
|
|
return a.moduleStates.Get(moduleName)
|
|
}
|
|
|
|
// Modules iterates over all the module state instances in the app.
|
|
func (a *App) Modules(f func(moduleName string, modState ModuleState) bool) {
|
|
a.moduleStates.Scan(func(key string, value *Module) bool {
|
|
return f(key, value)
|
|
})
|
|
}
|
|
|
|
// NumModules returns the number of modules in the app.
|
|
func (a *App) NumModules() int {
|
|
return a.moduleStates.Len()
|
|
}
|