package fs

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"linkfog.com/public/lib/cgroup/types"
)

func TestSetIORWLimit(t *testing.T) {
	tmpCgroupDir, err := ioutil.TempDir("/tmp", "cgroup-")
	if err != nil {
		t.Fatal("Cannot create temporary file", err)
	}
	defer os.RemoveAll(tmpCgroupDir)

	tmpIORdLimitFile := filepath.Join(tmpCgroupDir, "blkio.throttle.read_bps_device")
	tmpIOWrLimitFile := filepath.Join(tmpCgroupDir, "blkio.throttle.write_bps_device")
	devInfos := []types.DevInfo{{Name: "test-sda", Major: "1", Minor: "8"}}

	type ioTest struct {
		limit    uint64
		expected string
	}
	ioTests := []ioTest{
		{limit: 10 * 1024 * 1024, expected: "1:8 10485760"},
		{limit: 0, expected: "1:8 0"},
	}
	for i, test := range ioTests {
		// 清除测试环境
		os.Remove(tmpIORdLimitFile)
		os.Remove(tmpIOWrLimitFile)

		devLmts, err := SetIORWLimit(tmpCgroupDir, devInfos, test.limit, test.limit)
		if err != nil {
			t.Fatal(i, err)
		}
		if devLmt, ok := devLmts["1:8"]; !ok {
			t.Fatal(i, "device 1:8 not exist")
		} else if !devLmt.RdLimit || !devLmt.WrLimit {
			t.Fatal(i, "device 1:8 rw limit error")
		}

		ioRdLimitContent, err := ioutil.ReadFile(tmpIORdLimitFile)
		if err != nil {
			t.Fatal(i, err)
		}
		if string(ioRdLimitContent) != test.expected {
			t.Logf("%d io read limit error, content: %s", i, string(ioRdLimitContent))
		}

		ioWrLimitContent, err := ioutil.ReadFile(tmpIOWrLimitFile)
		if err != nil {
			t.Fatal(i, err)
		}
		if string(ioWrLimitContent) != test.expected {
			t.Fatalf("%d io write limit error, content: %s", i, string(ioWrLimitContent))
		}
	}
}