package workloads

import (
	"github.com/pingcap/test-infra/sdk/core"

	"github.com/pingcap/endless/testcase/tiflash/pkg/workload"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/ch"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/partition360"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/ssb"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/tpch"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/zto"
)

type WorkloadType string

const (
	CHWorkloadType           WorkloadType = "ch"
	ZtoWorkloadType                       = "zto"
	SsbWorkloadType                       = "ssb"
	Partition360WorkloadType              = "partition360"
	TpchWorkloadType                      = "tpch"
)

var HtapWorkloads map[WorkloadType]func(ctx core.TestContext) workload.Workload

func init() {
	HtapWorkloads = make(map[WorkloadType]func(ctx core.TestContext) workload.Workload)
	HtapWorkloads[CHWorkloadType] = ch.NewWorkload
	HtapWorkloads[ZtoWorkloadType] = zto.NewWorkload
	HtapWorkloads[SsbWorkloadType] = ssb.NewWorkload
	HtapWorkloads[Partition360WorkloadType] = partition360.NewWorkload
	HtapWorkloads[TpchWorkloadType] = tpch.NewWorkload
}