<commit_before>package gorethink

import (
	"time"

	test "gopkg.in/check.v1"
)

func (s *RethinkSuite) TestTimeTime(c *test.C) {
	var response time.Time
	res, err := Time(1986, 11, 3, 12, 30, 15, "Z").Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
	c.Assert(response.Equal(time.Date(1986, 11, 3, 12, 30, 15, 0, time.UTC)), test.Equals, true)
}

func (s *RethinkSuite) TestTimeTimeMillisecond(c *test.C) {
	var response time.Time
	res, err := Time(1986, 11, 3, 12, 30, 15.679, "Z").Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
	c.Assert(response.Equal(time.Date(1986, 11, 3, 12, 30, 15, 679.00002*1000*1000, time.UTC)), test.Equals, true)
}

func (s *RethinkSuite) TestTimeEpochTime(c *test.C) {
	var response time.Time
	res, err := EpochTime(531360000).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
	c.Assert(response.Equal(time.Date(1986, 11, 3, 0, 0, 0, 0, time.UTC)), test.Equals, true)
}

func (s *RethinkSuite) TestTimeExpr(c *test.C) {
	var response time.Time
	t := time.Unix(531360000, 0)
	res, err := Expr(Expr(t)).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
}

func (s *RethinkSuite) TestTimeExprMillisecond(c *test.C) {
	var response time.Time
	t := time.Unix(531360000, 679000000)
	res, err := Expr(t).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
	c.Assert(float64(response.UnixNano()), test.Equals, float64(t.UnixNano()))
}

func (s *RethinkSuite) TestTimeISO8601(c *test.C) {
	var t1, t2 time.Time
	t2, _ = time.Parse("2006-01-02T15:04:05-07:00", "1986-11-03T08:30:00-07:00")
	res, err := ISO8601("1986-11-03T08:30:00-07:00").Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&t1)
	c.Assert(err, test.IsNil)
	c.Assert(t1.Equal(t2), test.Equals, true)
}

func (s *RethinkSuite) TestTimeInTimezone(c *test.C) {
	loc, err := time.LoadLocation("MST")
	c.Assert(err, test.IsNil)
	var response []time.Time
	res, err2 := Expr([]interface{}{Now(), Now().InTimezone("-07:00")}).Run(session)
	c.Assert(err2, test.IsNil)

	err = res.All(&response)
	c.Assert(err, test.IsNil)
	c.Assert(response[1].Equal(response[0].In(loc)), test.Equals, true)
}

func (s *RethinkSuite) TestTimeBetween(c *test.C) {
	var response interface{}

	times := Expr([]interface{}{
		Time(1986, 9, 3, 12, 30, 15, "Z"),
		Time(1986, 10, 3, 12, 30, 15, "Z"),
		Time(1986, 11, 3, 12, 30, 15, "Z"),
		Time(1986, 12, 3, 12, 30, 15, "Z"),
	})
	res, err := times.Filter(func(row Term) Term {
		return row.During(Time(1986, 9, 3, 12, 30, 15, "Z"), Time(1986, 11, 3, 12, 30, 15, "Z"))
	}).Count().Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)

	c.Assert(err, test.IsNil)
	c.Assert(int(response.(float64)), test.Equals, 2)
}

func (s *RethinkSuite) TestTimeYear(c *test.C) {
	var response interface{}

	res, err := Time(1986, 12, 3, 12, 30, 15, "Z").Year().Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)

	c.Assert(err, test.IsNil)
	c.Assert(int(response.(float64)), test.Equals, 1986)
}

func (s *RethinkSuite) TestTimeMonth(c *test.C) {
	var response interface{}

	res, err := Time(1986, 12, 3, 12, 30, 15, "Z").Month().Eq(December).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)

	c.Assert(err, test.IsNil)
	c.Assert(response.(bool), test.Equals, true)
}

func (s *RethinkSuite) TestTimeDay(c *test.C) {
	var response interface{}

	res, err := Time(1986, 12, 3, 12, 30, 15, "Z").Day().Eq(Wednesday).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)

	c.Assert(err, test.IsNil)
	c.Assert(response.(bool), test.Equals, true)
}
<commit_msg>Update query_time_test.go<commit_after>package gorethink

import (
	"time"

	test "gopkg.in/check.v1"
)

func (s *RethinkSuite) TestTimeTime(c *test.C) {
	var response time.Time
	res, err := Time(1986, 11, 3, 12, 30, 15, "Z").Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
	c.Assert(response.Equal(time.Date(1986, 11, 3, 12, 30, 15, 0, time.UTC)), test.Equals, true)
}

func (s *RethinkSuite) TestTimeTimeMillisecond(c *test.C) {
	var response time.Time
	res, err := Time(1986, 11, 3, 12, 30, 15.6790123, "Z").Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
	c.Assert(response.Equal(time.Date(1986, 11, 3, 12, 30, 15, 679*1000*1000, time.UTC)), test.Equals, true)
}

func (s *RethinkSuite) TestTimeEpochTime(c *test.C) {
	var response time.Time
	res, err := EpochTime(531360000).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
	c.Assert(response.Equal(time.Date(1986, 11, 3, 0, 0, 0, 0, time.UTC)), test.Equals, true)
}

func (s *RethinkSuite) TestTimeExpr(c *test.C) {
	var response time.Time
	t := time.Unix(531360000, 0)
	res, err := Expr(Expr(t)).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
}

func (s *RethinkSuite) TestTimeExprMillisecond(c *test.C) {
	var response time.Time
	t := time.Unix(531360000, 679000000)
	res, err := Expr(t).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)
	c.Assert(err, test.IsNil)
	c.Assert(float64(response.UnixNano()), test.Equals, float64(t.UnixNano()))
}

func (s *RethinkSuite) TestTimeISO8601(c *test.C) {
	var t1, t2 time.Time
	t2, _ = time.Parse("2006-01-02T15:04:05-07:00", "1986-11-03T08:30:00-07:00")
	res, err := ISO8601("1986-11-03T08:30:00-07:00").Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&t1)
	c.Assert(err, test.IsNil)
	c.Assert(t1.Equal(t2), test.Equals, true)
}

func (s *RethinkSuite) TestTimeInTimezone(c *test.C) {
	loc, err := time.LoadLocation("MST")
	c.Assert(err, test.IsNil)
	var response []time.Time
	res, err2 := Expr([]interface{}{Now(), Now().InTimezone("-07:00")}).Run(session)
	c.Assert(err2, test.IsNil)

	err = res.All(&response)
	c.Assert(err, test.IsNil)
	c.Assert(response[1].Equal(response[0].In(loc)), test.Equals, true)
}

func (s *RethinkSuite) TestTimeBetween(c *test.C) {
	var response interface{}

	times := Expr([]interface{}{
		Time(1986, 9, 3, 12, 30, 15, "Z"),
		Time(1986, 10, 3, 12, 30, 15, "Z"),
		Time(1986, 11, 3, 12, 30, 15, "Z"),
		Time(1986, 12, 3, 12, 30, 15, "Z"),
	})
	res, err := times.Filter(func(row Term) Term {
		return row.During(Time(1986, 9, 3, 12, 30, 15, "Z"), Time(1986, 11, 3, 12, 30, 15, "Z"))
	}).Count().Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)

	c.Assert(err, test.IsNil)
	c.Assert(int(response.(float64)), test.Equals, 2)
}

func (s *RethinkSuite) TestTimeYear(c *test.C) {
	var response interface{}

	res, err := Time(1986, 12, 3, 12, 30, 15, "Z").Year().Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)

	c.Assert(err, test.IsNil)
	c.Assert(int(response.(float64)), test.Equals, 1986)
}

func (s *RethinkSuite) TestTimeMonth(c *test.C) {
	var response interface{}

	res, err := Time(1986, 12, 3, 12, 30, 15, "Z").Month().Eq(December).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)

	c.Assert(err, test.IsNil)
	c.Assert(response.(bool), test.Equals, true)
}

func (s *RethinkSuite) TestTimeDay(c *test.C) {
	var response interface{}

	res, err := Time(1986, 12, 3, 12, 30, 15, "Z").Day().Eq(Wednesday).Run(session)
	c.Assert(err, test.IsNil)

	err = res.One(&response)

	c.Assert(err, test.IsNil)
	c.Assert(response.(bool), test.Equals, true)
}
<|endoftext|>


<commit_before>package capnp

import (
	"testing"
)

func TestRawStructPointer(t *testing.T) {
	tests := []struct {
		ptr    rawPointer
		offset pointerOffset
		size   ObjectSize
	}{
		{0x0000000000000000, 0, ObjectSize{}},
		{0x0000000000000004, 1, ObjectSize{}},
		{0x0000000000000008, 2, ObjectSize{}},
		{0x000000000000000c, 3, ObjectSize{}},
		{0x0000000000000010, 4, ObjectSize{}},
		{0x0403020100000000, 0, ObjectSize{DataSize: 0x0201 * 8, PointerCount: 0x0403}},
		{0xffffffff00000000, 0, ObjectSize{DataSize: 0xffff * 8, PointerCount: 0xffff}},
		{0xfffffffffffffff8, -2, ObjectSize{DataSize: 0xffff * 8, PointerCount: 0xffff}},
		{0x00000000fffffffc, -1, ObjectSize{}},
		{0x04030201fffffffc, -1, ObjectSize{DataSize: 0x0201 * 8, PointerCount: 0x0403}},
		{0xfffffffffffffffc, -1, ObjectSize{DataSize: 0xffff * 8, PointerCount: 0xffff}},
	}
	for _, test := range tests {
		if typ := test.ptr.pointerType(); typ != structPointer {
			t.Errorf("rawPointer(%#016x).pointerType() = %d; want %d", uint64(test.ptr), typ, structPointer)
		}
		if offset := test.ptr.offset(); offset != test.offset {
			t.Errorf("rawPointer(%#016x).offset() = %d; want %d", uint64(test.ptr), offset, test.offset)
		}
		if size := test.ptr.structSize(); size != test.size {
			t.Errorf("rawPointer(%#016x).structSize() = %d; want %d", uint64(test.ptr), size, test.size)
		}
	}
}

func TestRawListPointer(t *testing.T) {
	tests := []struct {
		ptr    rawPointer
		offset pointerOffset
		lt     int
		n      int32
	}{
		{0x0000000000000001, 0, voidList, 0},
		{0x0000000000000005, 1, voidList, 0},
		{0x0000000000000009, 2, voidList, 0},
		{0x000000000000000d, 3, voidList, 0},
		{0x0000000100000001, 0, bit1List, 0},
		{0x0000000200000001, 0, byte1List, 0},
		{0x0000000300000001, 0, byte2List, 0},
		{0x0000000400000001, 0, byte4List, 0},
		{0x0000000500000001, 0, byte8List, 0},
		{0x0000000600000001, 0, pointerList, 0},
		{0x0000000700000001, 0, compositeList, 0},
		{0xfffffff8fffffffd, -1, voidList, 0x1fffffff},
		{0xfffffff9fffffffd, -1, bit1List, 0x1fffffff},
		{0xfffffffafffffffd, -1, byte1List, 0x1fffffff},
		{0xfffffffbfffffffd, -1, byte2List, 0x1fffffff},
		{0xfffffffcfffffffd, -1, byte4List, 0x1fffffff},
		{0xfffffffdfffffffd, -1, byte8List, 0x1fffffff},
		{0xfffffffefffffffd, -1, pointerList, 0x1fffffff},
		{0xfffffffffffffff9, -2, compositeList, 0x1fffffff},
		{0xfffffffffffffffd, -1, compositeList, 0x1fffffff},
	}
	for _, test := range tests {
		if typ := test.ptr.pointerType(); typ != listPointer {
			t.Errorf("rawPointer(%#016x).pointerType() = %d; want %d", uint64(test.ptr), typ, listPointer)
		}
		if offset := test.ptr.offset(); offset != test.offset {
			t.Errorf("rawPointer(%#016x).offset() = %d; want %d", uint64(test.ptr), offset, test.offset)
		}
		if lt := test.ptr.listType(); lt != test.lt {
			t.Errorf("rawPointer(%#016x).listType() = %d; want %d", uint64(test.ptr), lt, test.lt)
		}
		if n := test.ptr.numListElements(); n != test.n {
			t.Errorf("rawPointer(%#016x).numListElements() = %d; want %d", uint64(test.ptr), n, test.n)
		}
	}
}

func TestRawOtherPointer(t *testing.T) {
	tests := []struct {
		ptr rawPointer
		typ uint32
		cap CapabilityID
	}{
		{0x0000000000000003, 0, 0},
		{0x0000000000000007, 1, 0},
		{0x000000000000000b, 2, 0},
		{0x000000000000000f, 3, 0},
		{0xfffffffffffffffb, 0x3ffffffe, 0xffffffff},
		{0xffffffffffffffff, 0x3fffffff, 0xffffffff},
	}
	for _, test := range tests {
		if typ := test.ptr.pointerType(); typ != otherPointer {
			t.Errorf("rawPointer(%#016x).pointerType() = %d; want %d", uint64(test.ptr), typ, otherPointer)
		}
		if typ := test.ptr.otherPointerType(); typ != test.typ {
			t.Errorf("rawPointer(%#016x).otherPointerType() = %d; want %d", uint64(test.ptr), typ, test.typ)
		}
		if cap := test.ptr.capabilityIndex(); cap != test.cap {
			t.Errorf("rawPointer(%#016x).capabilityIndex() = %d; want %d", uint64(test.ptr), cap, test.cap)
		}
	}
}

func TestRawFarPointer(t *testing.T) {
	tests := []struct {
		ptr  rawPointer
		typ  int
		addr Address
		seg  SegmentID
	}{
		{0x0000000000000002, farPointer, 0, 0},
		{0x0000000000000006, doubleFarPointer, 0, 0},
		{0x000000000000000a, farPointer, 8, 0},
		{0x000000000000000e, doubleFarPointer, 8, 0},
		{0xfffffffffffffffa, farPointer, 0xfffffff8, 0xffffffff},
		{0xfffffffffffffffe, doubleFarPointer, 0xfffffff8, 0xffffffff},
	}
	for _, test := range tests {
		if typ := test.ptr.pointerType(); typ != test.typ {
			t.Errorf("rawPointer(%#016x).pointerType() = %d; want %d", uint64(test.ptr), typ, test.typ)
		}
		if addr := test.ptr.farAddress(); addr != test.addr {
			t.Errorf("rawPointer(%#016x).farAddress() = %v; want %v", uint64(test.ptr), addr, test.addr)
		}
		if seg := test.ptr.farSegment(); seg != test.seg {
			t.Errorf("rawPointer(%#016x).farSegment() = %d; want %d", uint64(test.ptr), seg, test.seg)
		}
	}
}
<commit_msg>capnp: add inverse rawPointer tests<commit_after>package capnp

import (
	"testing"
)

func TestRawStructPointer(t *testing.T) {
	tests := []struct {
		ptr    rawPointer
		offset pointerOffset
		size   ObjectSize
	}{
		{0x0000000000000000, 0, ObjectSize{}},
		{0x0000000000000004, 1, ObjectSize{}},
		{0x0000000000000008, 2, ObjectSize{}},
		{0x000000000000000c, 3, ObjectSize{}},
		{0x0000000000000010, 4, ObjectSize{}},
		{0x0403020100000000, 0, ObjectSize{DataSize: 0x0201 * 8, PointerCount: 0x0403}},
		{0xffffffff00000000, 0, ObjectSize{DataSize: 0xffff * 8, PointerCount: 0xffff}},
		{0xfffffffffffffff8, -2, ObjectSize{DataSize: 0xffff * 8, PointerCount: 0xffff}},
		{0x00000000fffffffc, -1, ObjectSize{}},
		{0x04030201fffffffc, -1, ObjectSize{DataSize: 0x0201 * 8, PointerCount: 0x0403}},
		{0xfffffffffffffffc, -1, ObjectSize{DataSize: 0xffff * 8, PointerCount: 0xffff}},
	}
	for _, test := range tests {
		if typ := test.ptr.pointerType(); typ != structPointer {
			t.Errorf("rawPointer(%#016x).pointerType() = %d; want %d", uint64(test.ptr), typ, structPointer)
		}
		if offset := test.ptr.offset(); offset != test.offset {
			t.Errorf("rawPointer(%#016x).offset() = %d; want %d", uint64(test.ptr), offset, test.offset)
		}
		if size := test.ptr.structSize(); size != test.size {
			t.Errorf("rawPointer(%#016x).structSize() = %d; want %d", uint64(test.ptr), size, test.size)
		}
	}
	for _, test := range tests {
		ptr := rawStructPointer(test.offset, test.size)
		if ptr != test.ptr {
			t.Errorf("rawStructPointer(%d, %d) = rawPointer(%#016x); want rawPointer(%#016x)", test.offset, test.size, ptr, test.ptr)
		}
	}
}

func TestRawListPointer(t *testing.T) {
	tests := []struct {
		ptr    rawPointer
		offset pointerOffset
		lt     int
		n      int32
	}{
		{0x0000000000000001, 0, voidList, 0},
		{0x0000000000000005, 1, voidList, 0},
		{0x0000000000000009, 2, voidList, 0},
		{0x000000000000000d, 3, voidList, 0},
		{0x0000000100000001, 0, bit1List, 0},
		{0x0000000200000001, 0, byte1List, 0},
		{0x0000000300000001, 0, byte2List, 0},
		{0x0000000400000001, 0, byte4List, 0},
		{0x0000000500000001, 0, byte8List, 0},
		{0x0000000600000001, 0, pointerList, 0},
		{0x0000000700000001, 0, compositeList, 0},
		{0xfffffff8fffffffd, -1, voidList, 0x1fffffff},
		{0xfffffff9fffffffd, -1, bit1List, 0x1fffffff},
		{0xfffffffafffffffd, -1, byte1List, 0x1fffffff},
		{0xfffffffbfffffffd, -1, byte2List, 0x1fffffff},
		{0xfffffffcfffffffd, -1, byte4List, 0x1fffffff},
		{0xfffffffdfffffffd, -1, byte8List, 0x1fffffff},
		{0xfffffffefffffffd, -1, pointerList, 0x1fffffff},
		{0xfffffffffffffff9, -2, compositeList, 0x1fffffff},
		{0xfffffffffffffffd, -1, compositeList, 0x1fffffff},
	}
	for _, test := range tests {
		if typ := test.ptr.pointerType(); typ != listPointer {
			t.Errorf("rawPointer(%#016x).pointerType() = %d; want %d", uint64(test.ptr), typ, listPointer)
		}
		if offset := test.ptr.offset(); offset != test.offset {
			t.Errorf("rawPointer(%#016x).offset() = %d; want %d", uint64(test.ptr), offset, test.offset)
		}
		if lt := test.ptr.listType(); lt != test.lt {
			t.Errorf("rawPointer(%#016x).listType() = %d; want %d", uint64(test.ptr), lt, test.lt)
		}
		if n := test.ptr.numListElements(); n != test.n {
			t.Errorf("rawPointer(%#016x).numListElements() = %d; want %d", uint64(test.ptr), n, test.n)
		}
	}
	for _, test := range tests {
		ptr := rawListPointer(test.offset, test.lt, test.n)
		if ptr != test.ptr {
			t.Errorf("rawListPointer(%d, %d, %d) = rawPointer(%#016x); want rawPointer(%#016x)", test.offset, test.lt, test.n, ptr, test.ptr)
		}
	}
}

func TestRawOtherPointer(t *testing.T) {
	tests := []struct {
		ptr rawPointer
		typ uint32
		cap CapabilityID
	}{
		{0x0000000000000003, 0, 0},
		{0x0000000000000007, 1, 0},
		{0x000000000000000b, 2, 0},
		{0x000000000000000f, 3, 0},
		{0xffffffff00000003, 0, 0xffffffff},
		{0xfffffffffffffffb, 0x3ffffffe, 0xffffffff},
		{0xffffffffffffffff, 0x3fffffff, 0xffffffff},
	}
	for _, test := range tests {
		if typ := test.ptr.pointerType(); typ != otherPointer {
			t.Errorf("rawPointer(%#016x).pointerType() = %d; want %d", uint64(test.ptr), typ, otherPointer)
		}
		if typ := test.ptr.otherPointerType(); typ != test.typ {
			t.Errorf("rawPointer(%#016x).otherPointerType() = %d; want %d", uint64(test.ptr), typ, test.typ)
		}
		if cap := test.ptr.capabilityIndex(); cap != test.cap {
			t.Errorf("rawPointer(%#016x).capabilityIndex() = %d; want %d", uint64(test.ptr), cap, test.cap)
		}
	}
	for _, test := range tests {
		if test.typ != 0 {
			continue
		}
		ptr := rawInterfacePointer(test.cap)
		if ptr != test.ptr {
			t.Errorf("rawInterfacePointer(%d) = rawPointer(%#016x); want rawPointer(%#016x)", test.cap, ptr, test.ptr)
		}
	}
}

func TestRawFarPointer(t *testing.T) {
	tests := []struct {
		ptr  rawPointer
		typ  int
		addr Address
		seg  SegmentID
	}{
		{0x0000000000000002, farPointer, 0, 0},
		{0x0000000000000006, doubleFarPointer, 0, 0},
		{0x000000000000000a, farPointer, 8, 0},
		{0x000000000000000e, doubleFarPointer, 8, 0},
		{0xfffffffffffffffa, farPointer, 0xfffffff8, 0xffffffff},
		{0xfffffffffffffffe, doubleFarPointer, 0xfffffff8, 0xffffffff},
	}
	for _, test := range tests {
		if typ := test.ptr.pointerType(); typ != test.typ {
			t.Errorf("rawPointer(%#016x).pointerType() = %d; want %d", uint64(test.ptr), typ, test.typ)
		}
		if addr := test.ptr.farAddress(); addr != test.addr {
			t.Errorf("rawPointer(%#016x).farAddress() = %v; want %v", uint64(test.ptr), addr, test.addr)
		}
		if seg := test.ptr.farSegment(); seg != test.seg {
			t.Errorf("rawPointer(%#016x).farSegment() = %d; want %d", uint64(test.ptr), seg, test.seg)
		}
	}
	for _, test := range tests {
		if test.typ == farPointer {
			ptr := rawFarPointer(test.seg, test.addr)
			if ptr != test.ptr {
				t.Errorf("rawFarPointer(%d, %v) = rawPointer(%#016x); want rawPointer(%#016x)", test.seg, test.addr, ptr, test.ptr)
			}
		} else {
			ptr := rawDoubleFarPointer(test.seg, test.addr)
			if ptr != test.ptr {
				t.Errorf("rawDoubleFarPointer(%d, %v) = rawPointer(%#016x); want rawPointer(%#016x)", test.seg, test.addr, ptr, test.ptr)
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 CloudAwan LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package rbac

import (
	"encoding/hex"
	"errors"
	"strings"
)

type Permission struct {
	Name      string
	Component string
	Method    string
	Path      string // Path is hierarchy
}

func CreatePermission(component string, method string, path string) (*Permission, error) {
	name, err := GetPermissionName(component, method, path)
	if err != nil {
		log.Error(err)
		return nil, err
	}

	return &Permission{
		name,
		component,
		method,
		path,
	}, nil
}

func GetPermissionName(component string, method string, path string) (string, error) {
	if component == "" {
		log.Error("Component couldn't be empty")
		return "", errors.New("Component couldn't be empty")
	}
	if method == "" {
		log.Error("Method couldn't be empty")
		return "", errors.New("Method couldn't be empty")
	}
	if path == "" {
		log.Error("Path couldn't be empty")
		return "", errors.New("Path couldn't be empty")
	}
	if strings.Contains(path, "/") == false {
		log.Error("Path format is invalid")
		return "", errors.New("Path format is invalid")
	}

	return hex.EncodeToString([]byte(component + " " + method + " " + path)), nil
}

func (permission *Permission) HasPermission(component string, method string, path string) bool {
	// * means all
	if permission.Component == "*" {
		return true
	} else if permission.Component == component {
		// * means all
		if permission.Method == "*" {
			// * means all
			if permission.Path == "*" {
				return true
			} else {
				// Prefix for hierarchy authorization
				return strings.HasPrefix(path, permission.Path)
			}
		} else if permission.Method == method {
			// * means all
			if permission.Path == "*" {
				return true
			} else {
				// Prefix for hierarchy authorization
				return strings.HasPrefix(path, permission.Path)
			}
		} else {
			// Different method won't apply path hierarchy authorization
			return false
		}
	} else {
		return false
	}
}

// Check whether user has the target permission node or child permission node of the target permission node along the tree
func (permission *Permission) HasChildPermission(component string, method string, path string) bool {
	// * means all
	if permission.Component == "*" {
		return true
	} else if permission.Component == component {
		// * means all
		if permission.Method == "*" {
			// Prefix for hierarchy authorization. Unlike HasPermission, here it check whether target permission is the same or the child of the permissions owned
			return strings.HasPrefix(permission.Path, path)
		} else if permission.Method == method {
			// Prefix for hierarchy authorization. Unlike HasPermission, here it check whether target permission is the same or the child of the permissions owned
			return strings.HasPrefix(permission.Path, path)
		} else {
			// Different method won't apply path hierarchy authorization
			return false
		}
	} else {
		return false
	}
}
<commit_msg>Remove unnecessary check<commit_after>// Copyright 2015 CloudAwan LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package rbac

import (
	"encoding/hex"
	"errors"
	"strings"
)

type Permission struct {
	Name      string
	Component string
	Method    string
	Path      string // Path is hierarchy
}

func CreatePermission(component string, method string, path string) (*Permission, error) {
	name, err := GetPermissionName(component, method, path)
	if err != nil {
		log.Error(err)
		return nil, err
	}

	return &Permission{
		name,
		component,
		method,
		path,
	}, nil
}

func GetPermissionName(component string, method string, path string) (string, error) {
	if component == "" {
		log.Error("Component couldn't be empty")
		return "", errors.New("Component couldn't be empty")
	}
	if method == "" {
		log.Error("Method couldn't be empty")
		return "", errors.New("Method couldn't be empty")
	}
	if path == "" {
		log.Error("Path couldn't be empty")
		return "", errors.New("Path couldn't be empty")
	}

	return hex.EncodeToString([]byte(component + " " + method + " " + path)), nil
}

func (permission *Permission) HasPermission(component string, method string, path string) bool {
	// * means all
	if permission.Component == "*" {
		return true
	} else if permission.Component == component {
		// * means all
		if permission.Method == "*" {
			// * means all
			if permission.Path == "*" {
				return true
			} else {
				// Prefix for hierarchy authorization
				return strings.HasPrefix(path, permission.Path)
			}
		} else if permission.Method == method {
			// * means all
			if permission.Path == "*" {
				return true
			} else {
				// Prefix for hierarchy authorization
				return strings.HasPrefix(path, permission.Path)
			}
		} else {
			// Different method won't apply path hierarchy authorization
			return false
		}
	} else {
		return false
	}
}

// Check whether user has the target permission node or child permission node of the target permission node along the tree
func (permission *Permission) HasChildPermission(component string, method string, path string) bool {
	// * means all
	if permission.Component == "*" {
		return true
	} else if permission.Component == component {
		// * means all
		if permission.Method == "*" {
			// Prefix for hierarchy authorization. Unlike HasPermission, here it check whether target permission is the same or the child of the permissions owned
			return strings.HasPrefix(permission.Path, path)
		} else if permission.Method == method {
			// Prefix for hierarchy authorization. Unlike HasPermission, here it check whether target permission is the same or the child of the permissions owned
			return strings.HasPrefix(permission.Path, path)
		} else {
			// Different method won't apply path hierarchy authorization
			return false
		}
	} else {
		return false
	}
}
<|endoftext|>


<commit_before>package xmpp

import (
	"crypto/tls"
	"encoding/xml"
	"errors"
	"fmt"
	"log"
	"net"
)

const (
	nsStream = "http://etherx.jabber.org/streams"
	nsTLS = "urn:ietf:params:xml:ns:xmpp-tls"
)

type Stream struct {
	conn net.Conn
	dec *xml.Decoder
}

func NewStream(addr string) (*Stream, error) {

	log.Println("Connecting to", addr)

	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	if _, err := conn.Write([]byte("<?xml version='1.0'?>")); err != nil {
		return nil, err
	}

	dec := xml.NewDecoder(conn)
	return &Stream{conn, dec}, nil
}

func (stream *Stream) UpgradeTLS(config *tls.Config) error {

	log.Println("Upgrading to TLS")

	if err := stream.Send("<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>"); err != nil {
		return err
	}

	p := tlsProceed{}
	if err := stream.Decode(&p); err != nil {
		return err
	}

	conn := tls.Client(stream.conn, &tls.Config{InsecureSkipVerify: true})
	if err := conn.Handshake(); err != nil {
		return err
	}

	stream.conn = conn
	stream.dec = xml.NewDecoder(stream.conn)

	return nil
}

func (stream *Stream) Send(s string) error {
	if _, err := stream.conn.Write([]byte(s)); err != nil {
		return err
	}
	return nil
}

func (stream *Stream) Next(match *xml.Name) (*xml.StartElement, error) {
	for {
		t, err := stream.dec.Token()
		if err != nil {
			return nil, err
		}
		if e, ok := t.(xml.StartElement); ok {
			if match != nil && e.Name != *match {
				return nil, errors.New(fmt.Sprintf("Expected %s, got %s", *match, e.Name))
			}
			return &e, nil
		}
	}
	panic("Unreachable")
}

func (stream *Stream) Decode(i interface{}) error {
	return stream.dec.Decode(i)
}

func (stream *Stream) DecodeElement(i interface{}, se *xml.StartElement) error {
	return stream.dec.DecodeElement(i, se)
}

type tlsProceed struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls proceed"`
}
<commit_msg>Add bug about UpgradeTLS doing XMPP-specific stuff.<commit_after>package xmpp

import (
	"crypto/tls"
	"encoding/xml"
	"errors"
	"fmt"
	"log"
	"net"
)

const (
	nsStream = "http://etherx.jabber.org/streams"
	nsTLS = "urn:ietf:params:xml:ns:xmpp-tls"
)

type Stream struct {
	conn net.Conn
	dec *xml.Decoder
}

func NewStream(addr string) (*Stream, error) {

	log.Println("Connecting to", addr)

	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	if _, err := conn.Write([]byte("<?xml version='1.0'?>")); err != nil {
		return nil, err
	}

	dec := xml.NewDecoder(conn)
	return &Stream{conn, dec}, nil
}

func (stream *Stream) UpgradeTLS(config *tls.Config) error {

	log.Println("Upgrading to TLS")

	if err := stream.Send("<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>"); err != nil {
		return err
	}

	p := tlsProceed{}
	if err := stream.Decode(&p); err != nil {
		return err
	}

	conn := tls.Client(stream.conn, &tls.Config{InsecureSkipVerify: true})
	if err := conn.Handshake(); err != nil {
		return err
	}

	stream.conn = conn
	stream.dec = xml.NewDecoder(stream.conn)

	return nil
}

func (stream *Stream) Send(s string) error {
	if _, err := stream.conn.Write([]byte(s)); err != nil {
		return err
	}
	return nil
}

func (stream *Stream) Next(match *xml.Name) (*xml.StartElement, error) {
	for {
		t, err := stream.dec.Token()
		if err != nil {
			return nil, err
		}
		if e, ok := t.(xml.StartElement); ok {
			if match != nil && e.Name != *match {
				return nil, errors.New(fmt.Sprintf("Expected %s, got %s", *match, e.Name))
			}
			return &e, nil
		}
	}
	panic("Unreachable")
}

func (stream *Stream) Decode(i interface{}) error {
	return stream.dec.Decode(i)
}

func (stream *Stream) DecodeElement(i interface{}, se *xml.StartElement) error {
	return stream.dec.DecodeElement(i, se)
}

type tlsProceed struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls proceed"`
}

// BUG(matt): UpgradeTLS shoudln't be doing anything specific to XMPP.
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package state

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/sys"
	"os"
	"sync"
	"time"
)

// A map from stat info for a file to a set of scores that represented that
// file's contents at the time the stat info was collected. (This of course is
// not atomic, so it's more like "around the time that the stat info was
// collected".)
//
// All methods are safe for concurrent calling.
type ScoreMap interface {
	// Set a list of scores for a particular key.
	Set(key ScoreMapKey, scores []blob.Score)

	// Get the list of scores previously set for a key, or nil if no list has
	// been set.
	Get(key ScoreMapKey) (scores []blob.Score)
}

// Create an empty map.
func NewScoreMap() ScoreMap {
	return &scoreMap{
		elems: make(map[ScoreMapKey][]blob.Score),
	}
}

// Contains fields used by git for a similar purpose according to racy-git.txt.
type ScoreMapKey struct {
	Path        string
	Permissions os.FileMode
	Uid         sys.UserId
	Gid         sys.GroupId
	MTime       time.Time
	Inode       uint64
	Size        uint64
}

////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////

type scoreMap struct {
	mutex sync.RWMutex
	elems map[ScoreMapKey][]blob.Score // Protected by mutex
}

func (s *scoreMap) Set(key ScoreMapKey, scores []blob.Score) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.elems[key] = scores
}

func (s *scoreMap) Get(key ScoreMapKey) (scores []blob.Score) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	scores, _ = s.elems[key]
	return
}

////////////////////////////////////////////////////////////////////////
// Gob encoding
////////////////////////////////////////////////////////////////////////

func init() {
	// Make sure that scoreMaps can be encoded where ScoreMap interface variables
	// are expected.
	gob.Register(&scoreMap{})
}

func (s *scoreMap) GobDecode(b []byte) (err error) {
	// Decode the map of elements.
	buf := bytes.NewBuffer(b)
	decoder := gob.NewDecoder(buf)

	elems := map[ScoreMapKey][]blob.Score{}
	if err = decoder.Decode(&elems); err != nil {
		err = fmt.Errorf("Decoding map: %v", err)
		return
	}

	// Overwrite our map.
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.elems = elems

	return
}

func (s *scoreMap) GobEncode() (b []byte, err error) {
	// Encode our map.
	buf := new(bytes.Buffer)
	encoder := gob.NewEncoder(buf)

	s.mutex.RLock()
	defer s.mutex.RUnlock()

	if err = encoder.Encode(s.elems); err != nil {
		err = fmt.Errorf("Encoding map: %v", err)
		return
	}

	b = buf.Bytes()
	return
}
<commit_msg>Use a cache for storing scores in the score map.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package state

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/cache"
	"github.com/jacobsa/comeback/sys"
	"os"
	"time"
)

// A map from stat info for a file to a set of scores that represented that
// file's contents at the time the stat info was collected. (This of course is
// not atomic, so it's more like "around the time that the stat info was
// collected".)
//
// All methods are safe for concurrent calling.
type ScoreMap interface {
	// Set a list of scores for a particular key.
	Set(key ScoreMapKey, scores []blob.Score)

	// Get the list of scores previously set for a key, or nil if no list has
	// been set.
	Get(key ScoreMapKey) (scores []blob.Score)
}

// Create an empty map.
func NewScoreMap() ScoreMap {
	return &scoreMap{
		ScoreCache: cache.NewLruCache(1e6),
	}
}

// Contains fields used by git for a similar purpose according to racy-git.txt.
type ScoreMapKey struct {
	Path        string
	Permissions os.FileMode
	Uid         sys.UserId
	Gid         sys.GroupId
	MTime       time.Time
	Inode       uint64
	Size        uint64
}

////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////

func init() {
	// Make sure that scoreMaps can be encoded where ScoreMap interface variables
	// are expected.
	gob.Register(&scoreMap{})
}

type scoreMap struct {
	ScoreCache cache.Cache
}

func toCacheKey(k ScoreMapKey) string {
	buf := new(bytes.Buffer)
	encoder := gob.NewEncoder(buf)

	if err := encoder.Encode(k); err != nil {
		panic(fmt.Sprintf("Error encoding ScoreMapKey: %v", err))
	}

	return buf.String()
}

func (s *scoreMap) Set(key ScoreMapKey, scores []blob.Score) {
	s.ScoreCache.Insert(toCacheKey(key), scores)
}

func (s *scoreMap) Get(key ScoreMapKey) (scores []blob.Score) {
	v := s.ScoreCache.LookUp(toCacheKey(key))
	if v == nil {
		return
	}

	scores = v.([]blob.Score)
	return
}
<|endoftext|>


<commit_before>/*
 *
 * Copyright 2017 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package grpclb

import (
	"context"
	"fmt"
	"io"
	"net"
	"reflect"
	"time"

	timestamppb "github.com/golang/protobuf/ptypes/timestamp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer"
	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/internal"
	"google.golang.org/grpc/internal/channelz"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/resolver"
)

// processServerList updates balaner's internal state, create/remove SubConns
// and regenerates picker using the received serverList.
func (lb *lbBalancer) processServerList(l *lbpb.ServerList) {
	if grpclog.V(2) {
		grpclog.Infof("lbBalancer: processing server list: %+v", l)
	}
	lb.mu.Lock()
	defer lb.mu.Unlock()

	// Set serverListReceived to true so fallback will not take effect if it has
	// not hit timeout.
	lb.serverListReceived = true

	// If the new server list == old server list, do nothing.
	if reflect.DeepEqual(lb.fullServerList, l.Servers) {
		if grpclog.V(2) {
			grpclog.Infof("lbBalancer: new serverlist same as the previous one, ignoring")
		}
		return
	}
	lb.fullServerList = l.Servers

	var backendAddrs []resolver.Address
	for i, s := range l.Servers {
		if s.Drop {
			continue
		}

		md := metadata.Pairs(lbTokenKey, s.LoadBalanceToken)
		ip := net.IP(s.IpAddress)
		ipStr := ip.String()
		if ip.To4() == nil {
			// Add square brackets to ipv6 addresses, otherwise net.Dial() and
			// net.SplitHostPort() will return too many colons error.
			ipStr = fmt.Sprintf("[%s]", ipStr)
		}
		addr := resolver.Address{
			Addr:     fmt.Sprintf("%s:%d", ipStr, s.Port),
			Metadata: &md,
		}
		if grpclog.V(2) {
			grpclog.Infof("lbBalancer: server list entry[%d]: ipStr:|%s|, port:|%d|, load balancer token:|%v|",
				i, ipStr, s.Port, s.LoadBalanceToken)
		}
		backendAddrs = append(backendAddrs, addr)
	}

	// Call refreshSubConns to create/remove SubConns.  If we are in fallback,
	// this is also exiting fallback.
	lb.refreshSubConns(backendAddrs, false, lb.usePickFirst)
}

// refreshSubConns creates/removes SubConns with backendAddrs, and refreshes
// balancer state and picker.
//
// Caller must hold lb.mu.
func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback bool, pickFirst bool) {
	opts := balancer.NewSubConnOptions{}
	if !fallback {
		opts.CredsBundle = lb.grpclbBackendCreds
	}

	lb.backendAddrs = backendAddrs
	lb.backendAddrsWithoutMetadata = nil

	fallbackModeChanged := lb.inFallback != fallback
	lb.inFallback = fallback

	balancingPolicyChanged := lb.usePickFirst != pickFirst
	oldUsePickFirst := lb.usePickFirst
	lb.usePickFirst = pickFirst

	if fallbackModeChanged || balancingPolicyChanged {
		// Remove all SubConns when switching balancing policy or switching
		// fallback mode.
		//
		// For fallback mode switching with pickfirst, we want to recreate the
		// SubConn because the creds could be different.
		for a, sc := range lb.subConns {
			if oldUsePickFirst {
				// If old SubConn were created for pickfirst, bypass cache and
				// remove directly.
				lb.cc.cc.RemoveSubConn(sc)
			} else {
				lb.cc.RemoveSubConn(sc)
			}
			delete(lb.subConns, a)
		}
	}

	if lb.usePickFirst {
		var sc balancer.SubConn
		for _, sc = range lb.subConns {
			break
		}
		if sc != nil {
			sc.UpdateAddresses(backendAddrs)
			sc.Connect()
			return
		}
		// This bypasses the cc wrapper with SubConn cache.
		sc, err := lb.cc.cc.NewSubConn(backendAddrs, opts)
		if err != nil {
			grpclog.Warningf("grpclb: failed to create new SubConn: %v", err)
			return
		}
		sc.Connect()
		lb.subConns[backendAddrs[0]] = sc
		lb.scStates[sc] = connectivity.Idle
		return
	}

	// addrsSet is the set converted from backendAddrsWithoutMetadata, it's used to quick
	// lookup for an address.
	addrsSet := make(map[resolver.Address]struct{})
	// Create new SubConns.
	for _, addr := range backendAddrs {
		addrWithoutMD := addr
		addrWithoutMD.Metadata = nil
		addrsSet[addrWithoutMD] = struct{}{}
		lb.backendAddrsWithoutMetadata = append(lb.backendAddrsWithoutMetadata, addrWithoutMD)

		if _, ok := lb.subConns[addrWithoutMD]; !ok {
			// Use addrWithMD to create the SubConn.
			sc, err := lb.cc.NewSubConn([]resolver.Address{addr}, opts)
			if err != nil {
				grpclog.Warningf("grpclb: failed to create new SubConn: %v", err)
				continue
			}
			lb.subConns[addrWithoutMD] = sc // Use the addr without MD as key for the map.
			if _, ok := lb.scStates[sc]; !ok {
				// Only set state of new sc to IDLE. The state could already be
				// READY for cached SubConns.
				lb.scStates[sc] = connectivity.Idle
			}
			sc.Connect()
		}
	}

	for a, sc := range lb.subConns {
		// a was removed by resolver.
		if _, ok := addrsSet[a]; !ok {
			lb.cc.RemoveSubConn(sc)
			delete(lb.subConns, a)
			// Keep the state of this sc in b.scStates until sc's state becomes Shutdown.
			// The entry will be deleted in HandleSubConnStateChange.
		}
	}

	// Regenerate and update picker after refreshing subconns because with
	// cache, even if SubConn was newed/removed, there might be no state
	// changes (the subconn will be kept in cache, not actually
	// newed/removed).
	lb.updateStateAndPicker(true, true)
}

func (lb *lbBalancer) readServerList(s *balanceLoadClientStream) error {
	for {
		reply, err := s.Recv()
		if err != nil {
			if err == io.EOF {
				return errServerTerminatedConnection
			}
			return fmt.Errorf("grpclb: failed to recv server list: %v", err)
		}
		if serverList := reply.GetServerList(); serverList != nil {
			lb.processServerList(serverList)
		}
	}
}

func (lb *lbBalancer) sendLoadReport(s *balanceLoadClientStream, interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
		case <-s.Context().Done():
			return
		}
		stats := lb.clientStats.toClientStats()
		t := time.Now()
		stats.Timestamp = &timestamppb.Timestamp{
			Seconds: t.Unix(),
			Nanos:   int32(t.Nanosecond()),
		}
		if err := s.Send(&lbpb.LoadBalanceRequest{
			LoadBalanceRequestType: &lbpb.LoadBalanceRequest_ClientStats{
				ClientStats: stats,
			},
		}); err != nil {
			return
		}
	}
}

func (lb *lbBalancer) callRemoteBalancer() (backoff bool, _ error) {
	lbClient := &loadBalancerClient{cc: lb.ccRemoteLB}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	stream, err := lbClient.BalanceLoad(ctx, grpc.WaitForReady(true))
	if err != nil {
		return true, fmt.Errorf("grpclb: failed to perform RPC to the remote balancer %v", err)
	}
	lb.mu.Lock()
	lb.remoteBalancerConnected = true
	lb.mu.Unlock()

	// grpclb handshake on the stream.
	initReq := &lbpb.LoadBalanceRequest{
		LoadBalanceRequestType: &lbpb.LoadBalanceRequest_InitialRequest{
			InitialRequest: &lbpb.InitialLoadBalanceRequest{
				Name: lb.target,
			},
		},
	}
	if err := stream.Send(initReq); err != nil {
		return true, fmt.Errorf("grpclb: failed to send init request: %v", err)
	}
	reply, err := stream.Recv()
	if err != nil {
		return true, fmt.Errorf("grpclb: failed to recv init response: %v", err)
	}
	initResp := reply.GetInitialResponse()
	if initResp == nil {
		return true, fmt.Errorf("grpclb: reply from remote balancer did not include initial response")
	}
	if initResp.LoadBalancerDelegate != "" {
		return true, fmt.Errorf("grpclb: Delegation is not supported")
	}

	go func() {
		if d := convertDuration(initResp.ClientStatsReportInterval); d > 0 {
			lb.sendLoadReport(stream, d)
		}
	}()
	// No backoff if init req/resp handshake was successful.
	return false, lb.readServerList(stream)
}

func (lb *lbBalancer) watchRemoteBalancer() {
	var retryCount int
	for {
		doBackoff, err := lb.callRemoteBalancer()
		select {
		case <-lb.doneCh:
			return
		default:
			if err != nil {
				if err == errServerTerminatedConnection {
					grpclog.Info(err)
				} else {
					grpclog.Warning(err)
				}
			}
		}
		// Trigger a re-resolve when the stream errors.
		lb.cc.cc.ResolveNow(resolver.ResolveNowOption{})

		lb.mu.Lock()
		lb.remoteBalancerConnected = false
		lb.fullServerList = nil
		// Enter fallback when connection to remote balancer is lost, and the
		// aggregated state is not Ready.
		if !lb.inFallback && lb.state != connectivity.Ready {
			// Entering fallback.
			lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
		}
		lb.mu.Unlock()

		if !doBackoff {
			retryCount = 0
			continue
		}

		timer := time.NewTimer(lb.backoff.Backoff(retryCount))
		select {
		case <-timer.C:
		case <-lb.doneCh:
			timer.Stop()
			return
		}
		retryCount++
	}
}

func (lb *lbBalancer) dialRemoteLB(remoteLBName string) {
	var dopts []grpc.DialOption
	if creds := lb.opt.DialCreds; creds != nil {
		if err := creds.OverrideServerName(remoteLBName); err == nil {
			dopts = append(dopts, grpc.WithTransportCredentials(creds))
		} else {
			grpclog.Warningf("grpclb: failed to override the server name in the credentials: %v, using Insecure", err)
			dopts = append(dopts, grpc.WithInsecure())
		}
	} else if bundle := lb.grpclbClientConnCreds; bundle != nil {
		dopts = append(dopts, grpc.WithCredentialsBundle(bundle))
	} else {
		dopts = append(dopts, grpc.WithInsecure())
	}
	if lb.opt.Dialer != nil {
		dopts = append(dopts, grpc.WithContextDialer(lb.opt.Dialer))
	}
	// Explicitly set pickfirst as the balancer.
	dopts = append(dopts, grpc.WithBalancerName(grpc.PickFirstBalancerName))
	wrb := internal.WithResolverBuilder.(func(resolver.Builder) grpc.DialOption)
	dopts = append(dopts, wrb(lb.manualResolver))
	if channelz.IsOn() {
		dopts = append(dopts, grpc.WithChannelzParentID(lb.opt.ChannelzParentID))
	}

	// Enable Keepalive for grpclb client.
	dopts = append(dopts, grpc.WithKeepaliveParams(keepalive.ClientParameters{
		Time:                20 * time.Second,
		Timeout:             10 * time.Second,
		PermitWithoutStream: true,
	}))

	// DialContext using manualResolver.Scheme, which is a random scheme
	// generated when init grpclb. The target scheme here is not important.
	//
	// The grpc dial target will be used by the creds (ALTS) as the authority,
	// so it has to be set to remoteLBName that comes from resolver.
	cc, err := grpc.DialContext(context.Background(), remoteLBName, dopts...)
	if err != nil {
		grpclog.Fatalf("failed to dial: %v", err)
	}
	lb.ccRemoteLB = cc
	go lb.watchRemoteBalancer()
}
<commit_msg>grpclb: avoid use of reflect.DeepEqual on proto.Message types (#3101)<commit_after>/*
 *
 * Copyright 2017 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package grpclb

import (
	"context"
	"fmt"
	"io"
	"net"
	"time"

	"github.com/golang/protobuf/proto"
	timestamppb "github.com/golang/protobuf/ptypes/timestamp"
	"github.com/google/go-cmp/cmp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer"
	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/internal"
	"google.golang.org/grpc/internal/channelz"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/resolver"
)

// processServerList updates balaner's internal state, create/remove SubConns
// and regenerates picker using the received serverList.
func (lb *lbBalancer) processServerList(l *lbpb.ServerList) {
	if grpclog.V(2) {
		grpclog.Infof("lbBalancer: processing server list: %+v", l)
	}
	lb.mu.Lock()
	defer lb.mu.Unlock()

	// Set serverListReceived to true so fallback will not take effect if it has
	// not hit timeout.
	lb.serverListReceived = true

	// If the new server list == old server list, do nothing.
	if cmp.Equal(lb.fullServerList, l.Servers, cmp.Comparer(proto.Equal)) {
		if grpclog.V(2) {
			grpclog.Infof("lbBalancer: new serverlist same as the previous one, ignoring")
		}
		return
	}
	lb.fullServerList = l.Servers

	var backendAddrs []resolver.Address
	for i, s := range l.Servers {
		if s.Drop {
			continue
		}

		md := metadata.Pairs(lbTokenKey, s.LoadBalanceToken)
		ip := net.IP(s.IpAddress)
		ipStr := ip.String()
		if ip.To4() == nil {
			// Add square brackets to ipv6 addresses, otherwise net.Dial() and
			// net.SplitHostPort() will return too many colons error.
			ipStr = fmt.Sprintf("[%s]", ipStr)
		}
		addr := resolver.Address{
			Addr:     fmt.Sprintf("%s:%d", ipStr, s.Port),
			Metadata: &md,
		}
		if grpclog.V(2) {
			grpclog.Infof("lbBalancer: server list entry[%d]: ipStr:|%s|, port:|%d|, load balancer token:|%v|",
				i, ipStr, s.Port, s.LoadBalanceToken)
		}
		backendAddrs = append(backendAddrs, addr)
	}

	// Call refreshSubConns to create/remove SubConns.  If we are in fallback,
	// this is also exiting fallback.
	lb.refreshSubConns(backendAddrs, false, lb.usePickFirst)
}

// refreshSubConns creates/removes SubConns with backendAddrs, and refreshes
// balancer state and picker.
//
// Caller must hold lb.mu.
func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback bool, pickFirst bool) {
	opts := balancer.NewSubConnOptions{}
	if !fallback {
		opts.CredsBundle = lb.grpclbBackendCreds
	}

	lb.backendAddrs = backendAddrs
	lb.backendAddrsWithoutMetadata = nil

	fallbackModeChanged := lb.inFallback != fallback
	lb.inFallback = fallback

	balancingPolicyChanged := lb.usePickFirst != pickFirst
	oldUsePickFirst := lb.usePickFirst
	lb.usePickFirst = pickFirst

	if fallbackModeChanged || balancingPolicyChanged {
		// Remove all SubConns when switching balancing policy or switching
		// fallback mode.
		//
		// For fallback mode switching with pickfirst, we want to recreate the
		// SubConn because the creds could be different.
		for a, sc := range lb.subConns {
			if oldUsePickFirst {
				// If old SubConn were created for pickfirst, bypass cache and
				// remove directly.
				lb.cc.cc.RemoveSubConn(sc)
			} else {
				lb.cc.RemoveSubConn(sc)
			}
			delete(lb.subConns, a)
		}
	}

	if lb.usePickFirst {
		var sc balancer.SubConn
		for _, sc = range lb.subConns {
			break
		}
		if sc != nil {
			sc.UpdateAddresses(backendAddrs)
			sc.Connect()
			return
		}
		// This bypasses the cc wrapper with SubConn cache.
		sc, err := lb.cc.cc.NewSubConn(backendAddrs, opts)
		if err != nil {
			grpclog.Warningf("grpclb: failed to create new SubConn: %v", err)
			return
		}
		sc.Connect()
		lb.subConns[backendAddrs[0]] = sc
		lb.scStates[sc] = connectivity.Idle
		return
	}

	// addrsSet is the set converted from backendAddrsWithoutMetadata, it's used to quick
	// lookup for an address.
	addrsSet := make(map[resolver.Address]struct{})
	// Create new SubConns.
	for _, addr := range backendAddrs {
		addrWithoutMD := addr
		addrWithoutMD.Metadata = nil
		addrsSet[addrWithoutMD] = struct{}{}
		lb.backendAddrsWithoutMetadata = append(lb.backendAddrsWithoutMetadata, addrWithoutMD)

		if _, ok := lb.subConns[addrWithoutMD]; !ok {
			// Use addrWithMD to create the SubConn.
			sc, err := lb.cc.NewSubConn([]resolver.Address{addr}, opts)
			if err != nil {
				grpclog.Warningf("grpclb: failed to create new SubConn: %v", err)
				continue
			}
			lb.subConns[addrWithoutMD] = sc // Use the addr without MD as key for the map.
			if _, ok := lb.scStates[sc]; !ok {
				// Only set state of new sc to IDLE. The state could already be
				// READY for cached SubConns.
				lb.scStates[sc] = connectivity.Idle
			}
			sc.Connect()
		}
	}

	for a, sc := range lb.subConns {
		// a was removed by resolver.
		if _, ok := addrsSet[a]; !ok {
			lb.cc.RemoveSubConn(sc)
			delete(lb.subConns, a)
			// Keep the state of this sc in b.scStates until sc's state becomes Shutdown.
			// The entry will be deleted in HandleSubConnStateChange.
		}
	}

	// Regenerate and update picker after refreshing subconns because with
	// cache, even if SubConn was newed/removed, there might be no state
	// changes (the subconn will be kept in cache, not actually
	// newed/removed).
	lb.updateStateAndPicker(true, true)
}

func (lb *lbBalancer) readServerList(s *balanceLoadClientStream) error {
	for {
		reply, err := s.Recv()
		if err != nil {
			if err == io.EOF {
				return errServerTerminatedConnection
			}
			return fmt.Errorf("grpclb: failed to recv server list: %v", err)
		}
		if serverList := reply.GetServerList(); serverList != nil {
			lb.processServerList(serverList)
		}
	}
}

func (lb *lbBalancer) sendLoadReport(s *balanceLoadClientStream, interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
		case <-s.Context().Done():
			return
		}
		stats := lb.clientStats.toClientStats()
		t := time.Now()
		stats.Timestamp = &timestamppb.Timestamp{
			Seconds: t.Unix(),
			Nanos:   int32(t.Nanosecond()),
		}
		if err := s.Send(&lbpb.LoadBalanceRequest{
			LoadBalanceRequestType: &lbpb.LoadBalanceRequest_ClientStats{
				ClientStats: stats,
			},
		}); err != nil {
			return
		}
	}
}

func (lb *lbBalancer) callRemoteBalancer() (backoff bool, _ error) {
	lbClient := &loadBalancerClient{cc: lb.ccRemoteLB}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	stream, err := lbClient.BalanceLoad(ctx, grpc.WaitForReady(true))
	if err != nil {
		return true, fmt.Errorf("grpclb: failed to perform RPC to the remote balancer %v", err)
	}
	lb.mu.Lock()
	lb.remoteBalancerConnected = true
	lb.mu.Unlock()

	// grpclb handshake on the stream.
	initReq := &lbpb.LoadBalanceRequest{
		LoadBalanceRequestType: &lbpb.LoadBalanceRequest_InitialRequest{
			InitialRequest: &lbpb.InitialLoadBalanceRequest{
				Name: lb.target,
			},
		},
	}
	if err := stream.Send(initReq); err != nil {
		return true, fmt.Errorf("grpclb: failed to send init request: %v", err)
	}
	reply, err := stream.Recv()
	if err != nil {
		return true, fmt.Errorf("grpclb: failed to recv init response: %v", err)
	}
	initResp := reply.GetInitialResponse()
	if initResp == nil {
		return true, fmt.Errorf("grpclb: reply from remote balancer did not include initial response")
	}
	if initResp.LoadBalancerDelegate != "" {
		return true, fmt.Errorf("grpclb: Delegation is not supported")
	}

	go func() {
		if d := convertDuration(initResp.ClientStatsReportInterval); d > 0 {
			lb.sendLoadReport(stream, d)
		}
	}()
	// No backoff if init req/resp handshake was successful.
	return false, lb.readServerList(stream)
}

func (lb *lbBalancer) watchRemoteBalancer() {
	var retryCount int
	for {
		doBackoff, err := lb.callRemoteBalancer()
		select {
		case <-lb.doneCh:
			return
		default:
			if err != nil {
				if err == errServerTerminatedConnection {
					grpclog.Info(err)
				} else {
					grpclog.Warning(err)
				}
			}
		}
		// Trigger a re-resolve when the stream errors.
		lb.cc.cc.ResolveNow(resolver.ResolveNowOption{})

		lb.mu.Lock()
		lb.remoteBalancerConnected = false
		lb.fullServerList = nil
		// Enter fallback when connection to remote balancer is lost, and the
		// aggregated state is not Ready.
		if !lb.inFallback && lb.state != connectivity.Ready {
			// Entering fallback.
			lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
		}
		lb.mu.Unlock()

		if !doBackoff {
			retryCount = 0
			continue
		}

		timer := time.NewTimer(lb.backoff.Backoff(retryCount))
		select {
		case <-timer.C:
		case <-lb.doneCh:
			timer.Stop()
			return
		}
		retryCount++
	}
}

func (lb *lbBalancer) dialRemoteLB(remoteLBName string) {
	var dopts []grpc.DialOption
	if creds := lb.opt.DialCreds; creds != nil {
		if err := creds.OverrideServerName(remoteLBName); err == nil {
			dopts = append(dopts, grpc.WithTransportCredentials(creds))
		} else {
			grpclog.Warningf("grpclb: failed to override the server name in the credentials: %v, using Insecure", err)
			dopts = append(dopts, grpc.WithInsecure())
		}
	} else if bundle := lb.grpclbClientConnCreds; bundle != nil {
		dopts = append(dopts, grpc.WithCredentialsBundle(bundle))
	} else {
		dopts = append(dopts, grpc.WithInsecure())
	}
	if lb.opt.Dialer != nil {
		dopts = append(dopts, grpc.WithContextDialer(lb.opt.Dialer))
	}
	// Explicitly set pickfirst as the balancer.
	dopts = append(dopts, grpc.WithBalancerName(grpc.PickFirstBalancerName))
	wrb := internal.WithResolverBuilder.(func(resolver.Builder) grpc.DialOption)
	dopts = append(dopts, wrb(lb.manualResolver))
	if channelz.IsOn() {
		dopts = append(dopts, grpc.WithChannelzParentID(lb.opt.ChannelzParentID))
	}

	// Enable Keepalive for grpclb client.
	dopts = append(dopts, grpc.WithKeepaliveParams(keepalive.ClientParameters{
		Time:                20 * time.Second,
		Timeout:             10 * time.Second,
		PermitWithoutStream: true,
	}))

	// DialContext using manualResolver.Scheme, which is a random scheme
	// generated when init grpclb. The target scheme here is not important.
	//
	// The grpc dial target will be used by the creds (ALTS) as the authority,
	// so it has to be set to remoteLBName that comes from resolver.
	cc, err := grpc.DialContext(context.Background(), remoteLBName, dopts...)
	if err != nil {
		grpclog.Fatalf("failed to dial: %v", err)
	}
	lb.ccRemoteLB = cc
	go lb.watchRemoteBalancer()
}
<|endoftext|>


<commit_before>// Copyright 2011 Mick Killianey.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strings

import (
	"fmt"
	"github.com/rdrdr/hamcrest/base"
	"regexp"
	"strings"
)

// Applies the given matcher to the result of writing the input object's
// to a string by using fmt.Sprintf("%v", object).
func ToString(matcher *base.Matcher) *base.Matcher {
	match := func(actual interface{}) *base.Result {
		if stringer, ok := actual.(fmt.Stringer); ok {
			s := stringer.String()
			result := matcher.Match(s)
			return base.NewResultf(
				result.Matched(), "String() returned %v", s).
				WithCauses(result)
		}
		s :=  fmt.Sprintf("%v", actual)
		result := matcher.Match(s)
		return base.NewResultf(result.Matched(),
			"Not a fmt.Stringer, but prints as %v", s).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToString(%v)", matcher)
}

// Applies the given matcher to the result of writing the input object's
// to a string by using fmt.Sprintf("%#v", object).
func ToGoString(matcher *base.Matcher) *base.Matcher {
	match := func(actual interface{}) *base.Result {
		if gostringer, ok := actual.(fmt.GoStringer); ok {
			s := gostringer.GoString()
			result := matcher.Match(s)
			return base.NewResultf(result.Matched(),
				"GoString() returned %v", s).
				WithCauses(result)
		}
		s := fmt.Sprintf("%#v", actual)
		result := matcher.Match(s)
		return base.NewResultf(result.Matched(),
			"Not a fmt.GoStringer, but prints as %v", s).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToGoString(%v)", matcher)
}

// Creates a new matcher that applies the given matcher to the result of
// converting an input string to lowercase (using strings.ToLower).
// If the input value is not a string, the matcher fails to match.
func ToLower(matcher *base.Matcher) *base.Matcher {
	match := func(s string) *base.Result {
		lower := strings.ToLower(s)
		result := matcher.Match(lower)
		return base.NewResultf(result.Matched(),
			"ToLower is %v", lower).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToLower(%v)", matcher)
}

// Creates a new matcher that applies the given matcher to the result of
// converting an input string to uppercase (using strings.ToUpper).
// If the input value is not a string, the matcher fails to match.
func ToUpper(matcher *base.Matcher) *base.Matcher {
	match := func(s string) *base.Result {
		upper := strings.ToUpper(s)
		result := matcher.Match(upper)
		return base.NewResultf(result.Matched(),
			"ToUpper is %v", upper).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToUpper(%v)", matcher)
}

func EqualToIgnoringCase(expected string) *base.Matcher {
	expectedToLower := strings.ToLower(expected)
	match := func(actual string) *base.Result {
		actualToLower := strings.ToLower(actual)
		if actualToLower == expectedToLower {
			return base.NewResultf(true,
				"\"%v\" matches \"%v\" (ignoring case)",
				actual, expected)
		}
		return base.NewResultf(false,
			"\"%v\" differs from \"%v\" (ignoring case)",
			actual, expected)
	}
	return base.NewMatcherf(match, "EqualToIgnoringCase(\"%v\")", expected)
}

// Creates a new matcher that applies the given matcher to the result of
// converting an input string its length. (using the `len()` builtin).
// If the input value is not a string, the matcher fails to match.
func ToLen(matcher *base.Matcher) *base.Matcher {
	match := func(s string) *base.Result {
		length := len(s)
		result := matcher.Match(length)
		return base.NewResultf(result.Matched(),
			"length is %v", length).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToLen(%v)", matcher)
}

// Matches strings that begin with the given prefix.
func HasPrefix(prefix string) *base.Matcher {
	maxLength := len(prefix) + 8 // arbitrary extra amount
	match := func (s string) *base.Result {
		continued := ""
		if len(s) > maxLength {
			s, continued = s[:maxLength], "..."
		}
		if strings.HasPrefix(s, prefix) {
			return base.NewResultf(true,
				"\"%v%v\" starts with \"%v\"", s, continued, prefix)
		}
		return base.NewResultf(false,
			"\"%v%v\" does not start with \"%v\"", s, continued, prefix)
	}
	return base.NewMatcherf(match, "HasPrefix(\"%v\")", prefix)
}

// Matches strings that end with the given prefix.
func HasSuffix(suffix string) *base.Matcher {
	maxLength := len(suffix) + 8 // arbitrary extra amount
	match := func (s string) *base.Result {
		continued := ""
		if len(s) > maxLength {
			continued, s = "...", s[len(s) - maxLength:]
		}
		if strings.HasSuffix(s, suffix) {
			return base.NewResultf(true,
				"\"%v%v\" ends with \"%v\"", s, continued, suffix)
		}
		return base.NewResultf(false,
			"\"%v%v\" does not end with \"%v\"", s, continued, suffix)
	}
	return base.NewMatcherf(match, "HasSuffix(\"%v\")", suffix)
}

// Matches strings that contain the given substring.
func Contains(substring string) *base.Matcher {
	match := func (s string) *base.Result {
		extra := 8
		if foundStart := strings.Index(s, substring); foundStart >= 0 {
			foundEnd := foundStart + len(substring)
			start, end := foundStart - extra, foundEnd + extra
			prefix, suffix := "", ""
			if start <= 0 {
				start = 0
			} else {
				prefix = "..."
			}
			if end >= len(s) {
				end = len(s)
			} else {
				suffix = "..."
			}
			return base.NewResultf(true,
				"substring \"%v\" appears in \"%v%v[%v]%v%v\"", substring,
				prefix, s[start:foundStart], substring, s[foundEnd:end], suffix)
		}
		return base.NewResultf(false,
			"substring \"%v\" does not appear in \"%v\"",
			substring, s)
	}
	return base.NewMatcherf(match, "Contains(\"%v\")", substring)
}

// Matches strings that contain the given regexp pattern, using
// the same syntax as the standard regexp package.
func HasPattern(pattern string) *base.Matcher {
	re := regexp.MustCompile(pattern)
	match := func (s string) *base.Result {
		if found := re.FindStringIndex(s); found != nil {
			start, end := found[0], found[1]
			return base.NewResultf(true,
				"pattern \"%v\" matched substring[%v:%v]=\"%v\"",
				pattern, start, end, s[start:end])
		}
		return base.NewResultf(false,
			"pattern \"%v\" not found in \"%v\"", pattern, s)
	}
	return base.NewMatcherf(match, "HasPattern[\"%v\"]", pattern)
}

type WithPatternClause struct {
	re *regexp.Regexp
	group int
}

func WithPattern(pattern string) *WithPatternClause {
	return &WithPatternClause{re: regexp.MustCompile(pattern), group: 0 }
}

//
//
//
func (self *WithPatternClause) Group(index int) *WithPatternClause {
	numGroups := self.re.NumSubexp()
	if index < 0 || index > numGroups {
		panic(fmt.Sprintf("Group %v doesn't exist: only %v groups in re %v",
			index, numGroups, self.re))
	}
	return &WithPatternClause{re: self.re, group: index }
}

// Completes a matcher that finds every occurrence of a pattern in the
// given input and applies the matcher to it, only matching if every
// occurrence matches.  For example:
//    i_before_e_except := ToLower(WithPattern(".ei").Each(StartsWith("c")))
// will match:
//    "ceiling receipt"
// but not:
//    "deceiver seizure"
func (self *WithPatternClause) EachMatch(matcher *base.Matcher) *base.Matcher {
	re := self.re
	match := func (s string) *base.Result {
		matches := re.FindAllStringIndex(s, -1)
		if matches == nil {
			return base.NewResultf(true,
				"No occurrences of pattern \"%v\"", re)
		}
		for _, loc := range matches {
			start, end := loc[0], loc[1]
			substring := s[start:end]
			result := matcher.Match(substring)
			if !result.Matched() {
				return base.NewResultf(false,
					"did not match substring[%v:%v]=\"%v\" for pattern \"%v\"",
					start, end, substring, re)
			}
		}
		return base.NewResultf(true,
			"Matched all occurrences of pattern \"%v\"", re)
	}
	return base.NewMatcherf(match, "EachMatch[\"%v\"][%v]", re, matcher)
}

// Completes a matcher that finds every occurrence of a pattern in the
// given input and applies the matcher to it, only matching if at least
// one occurrence matches.  For example:
//    here_kitty := WithPattern(".at").AnyMatch(StartsWith("c"))
// will match:
//    "that cat is phat"
// but not:
//    "Matt spat at a rat"
func (self *WithPatternClause) AnyMatch(matcher *base.Matcher) *base.Matcher {
	re := self.re
	match := func (s string) *base.Result {
		matches := re.FindAllStringIndex(s, -1)
		if matches == nil {
			return base.NewResultf(false, "No occurrences of pattern \"%v\"", re)
		}
		occurrences := 0
		for _, loc := range matches {
			occurrences += 1
			start, end := loc[0], loc[1]
			substring := s[start:end]
			result := matcher.Match(substring)
			if result.Matched() {
				return base.NewResultf(true,
					"matched substring[%v:%v]=\"%v\" on pattern \"%v\"",
					start, end, substring, re)
			}
		}
		return base.NewResultf(false,
			"Matched none of the %v occurrences of pattern \"%v\"",
			occurrences, re)
	}
	return base.NewMatcherf(match, "AnyMatch[\"%v\"][%v]", re, matcher)
}

<commit_msg>Use group number in regex matching<commit_after>// Copyright 2011 Mick Killianey.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strings

import (
	"fmt"
	"github.com/rdrdr/hamcrest/base"
	"regexp"
	"strings"
)

// Applies the given matcher to the result of writing the input object's
// to a string by using fmt.Sprintf("%v", object).
func ToString(matcher *base.Matcher) *base.Matcher {
	match := func(actual interface{}) *base.Result {
		if stringer, ok := actual.(fmt.Stringer); ok {
			s := stringer.String()
			result := matcher.Match(s)
			return base.NewResultf(
				result.Matched(), "String() returned %v", s).
				WithCauses(result)
		}
		s :=  fmt.Sprintf("%v", actual)
		result := matcher.Match(s)
		return base.NewResultf(result.Matched(),
			"Not a fmt.Stringer, but prints as %v", s).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToString(%v)", matcher)
}

// Applies the given matcher to the result of writing the input object's
// to a string by using fmt.Sprintf("%#v", object).
func ToGoString(matcher *base.Matcher) *base.Matcher {
	match := func(actual interface{}) *base.Result {
		if gostringer, ok := actual.(fmt.GoStringer); ok {
			s := gostringer.GoString()
			result := matcher.Match(s)
			return base.NewResultf(result.Matched(),
				"GoString() returned %v", s).
				WithCauses(result)
		}
		s := fmt.Sprintf("%#v", actual)
		result := matcher.Match(s)
		return base.NewResultf(result.Matched(),
			"Not a fmt.GoStringer, but prints as %v", s).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToGoString(%v)", matcher)
}

// Creates a new matcher that applies the given matcher to the result of
// converting an input string to lowercase (using strings.ToLower).
// If the input value is not a string, the matcher fails to match.
func ToLower(matcher *base.Matcher) *base.Matcher {
	match := func(s string) *base.Result {
		lower := strings.ToLower(s)
		result := matcher.Match(lower)
		return base.NewResultf(result.Matched(),
			"ToLower is %v", lower).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToLower(%v)", matcher)
}

// Creates a new matcher that applies the given matcher to the result of
// converting an input string to uppercase (using strings.ToUpper).
// If the input value is not a string, the matcher fails to match.
func ToUpper(matcher *base.Matcher) *base.Matcher {
	match := func(s string) *base.Result {
		upper := strings.ToUpper(s)
		result := matcher.Match(upper)
		return base.NewResultf(result.Matched(),
			"ToUpper is %v", upper).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToUpper(%v)", matcher)
}

func EqualToIgnoringCase(expected string) *base.Matcher {
	expectedToLower := strings.ToLower(expected)
	match := func(actual string) *base.Result {
		actualToLower := strings.ToLower(actual)
		if actualToLower == expectedToLower {
			return base.NewResultf(true,
				"\"%v\" matches \"%v\" (ignoring case)",
				actual, expected)
		}
		return base.NewResultf(false,
			"\"%v\" differs from \"%v\" (ignoring case)",
			actual, expected)
	}
	return base.NewMatcherf(match, "EqualToIgnoringCase(\"%v\")", expected)
}

// Creates a new matcher that applies the given matcher to the result of
// converting an input string its length. (using the `len()` builtin).
// If the input value is not a string, the matcher fails to match.
func ToLen(matcher *base.Matcher) *base.Matcher {
	match := func(s string) *base.Result {
		length := len(s)
		result := matcher.Match(length)
		return base.NewResultf(result.Matched(),
			"length is %v", length).
			WithCauses(result)
	}
	return base.NewMatcherf(match, "ToLen(%v)", matcher)
}

// Matches strings that begin with the given prefix.
func HasPrefix(prefix string) *base.Matcher {
	maxLength := len(prefix) + 8 // arbitrary extra amount
	match := func (s string) *base.Result {
		continued := ""
		if len(s) > maxLength {
			s, continued = s[:maxLength], "..."
		}
		if strings.HasPrefix(s, prefix) {
			return base.NewResultf(true,
				"\"%v%v\" starts with \"%v\"", s, continued, prefix)
		}
		return base.NewResultf(false,
			"\"%v%v\" does not start with \"%v\"", s, continued, prefix)
	}
	return base.NewMatcherf(match, "HasPrefix(\"%v\")", prefix)
}

// Matches strings that end with the given prefix.
func HasSuffix(suffix string) *base.Matcher {
	maxLength := len(suffix) + 8 // arbitrary extra amount
	match := func (s string) *base.Result {
		continued := ""
		if len(s) > maxLength {
			continued, s = "...", s[len(s) - maxLength:]
		}
		if strings.HasSuffix(s, suffix) {
			return base.NewResultf(true,
				"\"%v%v\" ends with \"%v\"", s, continued, suffix)
		}
		return base.NewResultf(false,
			"\"%v%v\" does not end with \"%v\"", s, continued, suffix)
	}
	return base.NewMatcherf(match, "HasSuffix(\"%v\")", suffix)
}

// Matches strings that contain the given substring.
func Contains(substring string) *base.Matcher {
	match := func (s string) *base.Result {
		extra := 8
		if foundStart := strings.Index(s, substring); foundStart >= 0 {
			foundEnd := foundStart + len(substring)
			start, end := foundStart - extra, foundEnd + extra
			prefix, suffix := "", ""
			if start <= 0 {
				start = 0
			} else {
				prefix = "..."
			}
			if end >= len(s) {
				end = len(s)
			} else {
				suffix = "..."
			}
			return base.NewResultf(true,
				"substring \"%v\" appears in \"%v%v[%v]%v%v\"", substring,
				prefix, s[start:foundStart], substring, s[foundEnd:end], suffix)
		}
		return base.NewResultf(false,
			"substring \"%v\" does not appear in \"%v\"",
			substring, s)
	}
	return base.NewMatcherf(match, "Contains(\"%v\")", substring)
}

// Matches strings that contain the given regexp pattern, using
// the same syntax as the standard regexp package.
func HasPattern(pattern string) *base.Matcher {
	re := regexp.MustCompile(pattern)
	match := func (s string) *base.Result {
		if found := re.FindStringIndex(s); found != nil {
			start, end := found[0], found[1]
			return base.NewResultf(true,
				"pattern \"%v\" matched substring[%v:%v]=\"%v\"",
				pattern, start, end, s[start:end])
		}
		return base.NewResultf(false,
			"pattern \"%v\" not found in \"%v\"", pattern, s)
	}
	return base.NewMatcherf(match, "HasPattern[\"%v\"]", pattern)
}

type WithPatternClause struct {
	re *regexp.Regexp
	group int
}

func WithPattern(pattern string) *WithPatternClause {
	return &WithPatternClause{re: regexp.MustCompile(pattern), group: 0 }
}

//
//
//
func (self *WithPatternClause) Group(index int) *WithPatternClause {
	numGroups := self.re.NumSubexp()
	if index < 0 || index > numGroups {
		panic(fmt.Sprintf("Group %v doesn't exist: only %v groups in re %v",
			index, numGroups, self.re))
	}
	return &WithPatternClause{re: self.re, group: index }
}

// Completes a matcher that finds every occurrence of a pattern in the
// given input and applies the matcher to it, only matching if every
// occurrence matches.  For example:
//    i_before_e_except := ToLower(WithPattern(".ei").Each(StartsWith("c")))
// will match:
//    "ceiling receipt"
// but not:
//    "deceiver seizure"
func (self *WithPatternClause) EachMatch(matcher *base.Matcher) *base.Matcher {
	re := self.re
	match := func (s string) *base.Result {
		matches := re.FindAllStringIndex(s, -1)
		if matches == nil {
			return base.NewResultf(true,
				"No occurrences of pattern \"%v\"", re)
		}
		groupOffset := 2 * self.group
		for _, loc := range matches {
			start, end := loc[groupOffset], loc[groupOffset + 1]
			substring := s[start:end]
			result := matcher.Match(substring)
			if !result.Matched() {
				return base.NewResultf(false,
					"did not match substring[%v:%v]=\"%v\" for pattern \"%v\" group %v",
					start, end, substring, re, self.group)
			}
		}
		return base.NewResultf(true,
			"Matched all occurrences of pattern \"%v\"", re)
	}
	return base.NewMatcherf(match, "EachMatch[\"%v\"][%v]", re, matcher)
}

// Completes a matcher that finds every occurrence of a pattern in the
// given input and applies the matcher to it, only matching if at least
// one occurrence matches.  For example:
//    here_kitty := WithPattern(".at").AnyMatch(StartsWith("c"))
// will match:
//    "that cat is phat"
// but not:
//    "Matt spat at a rat"
func (self *WithPatternClause) AnyMatch(matcher *base.Matcher) *base.Matcher {
	re := self.re
	match := func (s string) *base.Result {
		matches := re.FindAllStringIndex(s, -1)
		if matches == nil {
			return base.NewResultf(false, "No occurrences of pattern \"%v\"", re)
		}
		occurrences := 0
		groupOffset := 2 * self.group
		for _, loc := range matches {
			occurrences += 1
			start, end := loc[groupOffset], loc[groupOffset + 1]
			substring := s[start:end]
			result := matcher.Match(substring)
			if result.Matched() {
				return base.NewResultf(true,
					"matched substring[%v:%v]=\"%v\" on pattern \"%v\" group %v",
					start, end, substring, re, self.group)
			}
		}
		return base.NewResultf(false,
			"Matched none of the %v occurrences of pattern \"%v\"",
			occurrences, re)
	}
	return base.NewMatcherf(match, "AnyMatch[\"%v\"][%v]", re, matcher)
}

// Completes a matcher that finds every occurrence of a pattern in
// the given input and applies the matcher to it, only matching if
// there is exactly one occurrence, and the provided matcher matches
// that occurrence.  For example:
//     treasure := WithPattern("x+").TheMatch(ToLen(Equals(1)))
// will match:
//     "..x.."
// but not:
//     "..y.." because there are no occurrences of the pattern
//     "..x..x.." because there are more than one occurrences of the pattern
//     "..xx.." because the ToLen matcher does not match
func (self *WithPatternClause) TheMatch(matcher *base.Matcher) *base.Matcher {
	re := self.re
	match := func (s string) *base.Result {
		matches := re.FindAllStringIndex(s, 2)
		if matches == nil {
			return base.NewResultf(false, "No occurrences of pattern \"%v\"", re)
		}
		if len(matches) > 1 {
			return base.NewResultf(false, "Multiple occurrences of pattern \"%v\"", re)
		}
		loc := matches[0]
		groupOffset := 2 * self.group
		start, end := loc[groupOffset], loc[groupOffset + 1]
		substring := s[start:end]
		result := matcher.Match(substring)
		return base.NewResultf(result.Matched(),
			"Matched substring[%v:%v]=\"%v\" on pattern \"%v\" group %v",
				start, end, substring, re, self.group)
	}
	return base.NewMatcherf(match, "TheMatch[\"%v\"][%v]", re, matcher)
}

<|endoftext|>


<commit_before>// +build !race
// TODO(https://github.com/kubernetes/kubernetes/issues/90952): Remove the above.

/*
Copyright 2020 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package leaderelection

import (
	"context"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	fakekube "k8s.io/client-go/kubernetes/fake"
	ktesting "k8s.io/client-go/testing"
	"knative.dev/pkg/reconciler"
	_ "knative.dev/pkg/system/testing"
)

func TestWithBuilder(t *testing.T) {
	const buckets = 3
	cc := ComponentConfig{
		Component:     "the-component",
		Buckets:       buckets,
		LeaseDuration: 15 * time.Second,
		RenewDeadline: 10 * time.Second,
		RetryPeriod:   2 * time.Second,
	}
	kc := fakekube.NewSimpleClientset()
	ctx := context.Background()

	gotNames := make(sets.String, buckets)
	promoted := make(chan string)
	demoted := make(chan struct{})
	laf := &reconciler.LeaderAwareFuncs{
		PromoteFunc: func(bkt reconciler.Bucket, enq func(reconciler.Bucket, types.NamespacedName)) error {
			promoted <- bkt.Name()
			return nil
		},
		DemoteFunc: func(bkt reconciler.Bucket) {
			demoted <- struct{}{}
		},
	}
	enq := func(reconciler.Bucket, types.NamespacedName) {}

	created := make(chan struct{})
	kc.PrependReactor("create", "leases",
		func(action ktesting.Action) (bool, runtime.Object, error) {
			created <- struct{}{}
			return false, nil, nil
		},
	)

	updated := make(chan struct{})
	kc.PrependReactor("update", "leases",
		func(action ktesting.Action) (bool, runtime.Object, error) {
			// Only close updated once.
			select {
			case <-updated:
			default:
				close(updated)
			}
			return false, nil, nil
		},
	)

	if HasLeaderElection(ctx) {
		t.Error("HasLeaderElection() = true, wanted false")
	}
	if le, err := BuildElector(ctx, laf, "name", enq); err != nil {
		t.Errorf("BuildElector() = %v, wanted an unopposedElector", err)
	} else if _, ok := le.(*unopposedElector); !ok {
		t.Errorf("BuildElector() = %T, wanted an unopposedElector", le)
	}

	ctx = WithDynamicLeaderElectorBuilder(ctx, kc, cc)
	if !HasLeaderElection(ctx) {
		t.Error("HasLeaderElection() = false, wanted true")
	}

	le, err := BuildElector(ctx, laf, "name", enq)
	if err != nil {
		t.Fatal("BuildElector() =", err)
	}

	// We shouldn't see leases until we Run the elector.
	select {
	case s := <-promoted:
		gotNames.Insert(s)
		t.Error("Got promoted, want no actions.")
	case <-demoted:
		t.Error("Got demoted, want no actions.")
	case <-created:
		t.Error("Got created, want no actions.")
	case <-updated:
		t.Error("Got updated, want no actions.")
	default:
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go le.Run(ctx)

	// We expect 3 lease to be created.
	for i := 0; i < buckets; i++ {
		select {
		case <-created:
		case <-time.After(1 * time.Second):
			t.Fatal("Timed out waiting for lease creation.")
		}
	}
	// We expect to have been promoted 3 times.
	for i := 0; i < buckets; i++ {
		select {
		case s := <-promoted:
			gotNames.Insert(s)
		case <-time.After(time.Second):
			t.Fatal("Timed out waiting for promotion.")
		}
	}

	// Cancelling the context should case us to give up leadership.
	cancel()

	select {
	case <-updated:
		// We expect the lease to be updated.
	case <-time.After(time.Second):
		t.Fatal("Timed out waiting for lease update.")
	}
	// We expect to have been demoted 3 times.
	for i := 0; i < buckets; i++ {
		select {
		case <-demoted:
		case <-time.After(time.Second):
			t.Fatal("Timed out waiting for demotion.")
		}
	}

	want := sets.NewString(
		"the-component.name.00-of-03",
		"the-component.name.01-of-03",
		"the-component.name.02-of-03",
	)
	if !gotNames.Equal(want) {
		t.Errorf("BucketSet.BucketList() = %q, want: %q", gotNames, want)
	}
}

func TestBuilderWithCustomizedLeaseName(t *testing.T) {
	const buckets = 3
	cc := ComponentConfig{
		Component:     "the-component",
		Buckets:       buckets,
		LeaseDuration: 15 * time.Second,
		RenewDeadline: 10 * time.Second,
		RetryPeriod:   2 * time.Second,
		LeaseName: func(i uint32) string {
			return fmt.Sprintf("bucket-%02d", i)
		},
	}
	kc := fakekube.NewSimpleClientset()
	ctx := context.Background()

	gotNames := make(sets.String, buckets)
	promoted := make(chan string)
	laf := &reconciler.LeaderAwareFuncs{
		PromoteFunc: func(bkt reconciler.Bucket, enq func(reconciler.Bucket, types.NamespacedName)) error {
			promoted <- bkt.Name()
			return nil
		},
	}
	enq := func(reconciler.Bucket, types.NamespacedName) {}
	ctx = WithDynamicLeaderElectorBuilder(ctx, kc, cc)
	le, err := BuildElector(ctx, laf, "name", enq)
	if err != nil {
		t.Fatal("BuildElector() =", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go le.Run(ctx)

	// We expect to have been promoted 3 times.
	for i := 0; i < buckets; i++ {
		select {
		case s := <-promoted:
			gotNames.Insert(s)
		case <-time.After(time.Second):
			t.Fatal("Timed out waiting for promotion.")
		}
	}

	want := sets.NewString(
		"bucket-00",
		"bucket-01",
		"bucket-02",
	)
	if !gotNames.Equal(want) {
		t.Errorf("BucketSet.BucketList() = %q, want: %q", gotNames, want)
	}
}

func TestNewStatefulSetBucketAndSet(t *testing.T) {
	wantNames := []string{
		"http://as-0.autoscaler.knative-testing.svc.cluster.local:80",
		"http://as-1.autoscaler.knative-testing.svc.cluster.local:80",
		"http://as-2.autoscaler.knative-testing.svc.cluster.local:80",
	}

	os.Setenv(controllerOrdinalEnv, "as-2")
	os.Setenv(serviceNameEnv, "autoscaler")
	t.Cleanup(func() {
		os.Unsetenv(controllerOrdinalEnv)
		os.Unsetenv(serviceNameEnv)
	})

	_, _, err := NewStatefulSetBucketAndSet(2)
	if err == nil {
		// Ordinal 2 should be range [0, 2)
		t.Fatal("Expected error from NewStatefulSetBucketAndSet but got nil")
	}

	bkt, bs, err := NewStatefulSetBucketAndSet(3)
	if err != nil {
		// Ordinal 2 should be range [0, 2)
		t.Fatal("NewStatefulSetBucketAndSet() = ", err)
	}

	if got, want := bkt.Name(), wantNames[2]; got != want {
		t.Errorf("Bucket.Name() = %s, want = %s", got, want)
	}

	gotNames := bs.BucketList()
	if !cmp.Equal(gotNames, wantNames) {
		t.Errorf("BucketSet.BucketList() = %q, want: %q", gotNames, wantNames)
	}
}

func TestWithStatefulSetBuilder(t *testing.T) {
	cc := ComponentConfig{
		Component: "the-component",
		Buckets:   3,
	}
	const podDNS = "http://as-2.autoscaler.knative-testing.svc.cluster.local:80"
	ctx := context.Background()

	promoted := make(chan struct{})
	laf := &reconciler.LeaderAwareFuncs{
		PromoteFunc: func(bkt reconciler.Bucket, enq func(reconciler.Bucket, types.NamespacedName)) error {
			close(promoted)
			return nil
		},
	}
	enq := func(reconciler.Bucket, types.NamespacedName) {}

	if os.Setenv(controllerOrdinalEnv, "as-2") != nil {
		t.Fatalf("Failed to set env var %s=%s", controllerOrdinalEnv, "as-2")
	}
	if os.Setenv(serviceNameEnv, "autoscaler") != nil {
		t.Fatalf("Failed to set env var %s=%s", serviceNameEnv, "autoscaler")
	}
	t.Cleanup(func() {
		os.Unsetenv(controllerOrdinalEnv)
		os.Unsetenv(serviceNameEnv)
	})

	ctx = WithDynamicLeaderElectorBuilder(ctx, nil, cc)
	if !HasLeaderElection(ctx) {
		t.Error("HasLeaderElection() = false, wanted true")
	}

	b := ctx.Value(builderKey{})
	ssb, ok := b.(*statefulSetBuilder)
	if !ok || ssb == nil {
		t.Fatal("StatefulSetBuilder not found on context")
	}

	le, err := BuildElector(ctx, laf, "name", enq)
	if err != nil {
		t.Fatal("BuildElector() =", err)
	}

	ule, ok := le.(*unopposedElector)
	if !ok {
		t.Fatalf("BuildElector() = %T, wanted an unopposedElector", le)
	}
	if got, want := ule.bkt.Name(), podDNS; got != want {
		t.Errorf("bkt.Name() = %s, wanted %s", got, want)
	}

	// Shouldn't be promoted until we Run the elector.
	select {
	case <-promoted:
		t.Error("Got promoted, want no actions.")
	default:
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go le.Run(ctx)

	select {
	case <-promoted:
		// We expect to have been promoted.
	case <-time.After(time.Second):
		t.Fatal("Timed out waiting for promotion.")
	}
}
<commit_msg>race was fixed in k8s 1.18 (#1864)<commit_after>/*
Copyright 2020 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package leaderelection

import (
	"context"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	fakekube "k8s.io/client-go/kubernetes/fake"
	ktesting "k8s.io/client-go/testing"
	"knative.dev/pkg/reconciler"
	_ "knative.dev/pkg/system/testing"
)

func TestWithBuilder(t *testing.T) {
	const buckets = 3
	cc := ComponentConfig{
		Component:     "the-component",
		Buckets:       buckets,
		LeaseDuration: 15 * time.Second,
		RenewDeadline: 10 * time.Second,
		RetryPeriod:   2 * time.Second,
	}
	kc := fakekube.NewSimpleClientset()
	ctx := context.Background()

	gotNames := make(sets.String, buckets)
	promoted := make(chan string)
	demoted := make(chan struct{})
	laf := &reconciler.LeaderAwareFuncs{
		PromoteFunc: func(bkt reconciler.Bucket, enq func(reconciler.Bucket, types.NamespacedName)) error {
			promoted <- bkt.Name()
			return nil
		},
		DemoteFunc: func(bkt reconciler.Bucket) {
			demoted <- struct{}{}
		},
	}
	enq := func(reconciler.Bucket, types.NamespacedName) {}

	created := make(chan struct{})
	kc.PrependReactor("create", "leases",
		func(action ktesting.Action) (bool, runtime.Object, error) {
			created <- struct{}{}
			return false, nil, nil
		},
	)

	updated := make(chan struct{})
	kc.PrependReactor("update", "leases",
		func(action ktesting.Action) (bool, runtime.Object, error) {
			// Only close updated once.
			select {
			case <-updated:
			default:
				close(updated)
			}
			return false, nil, nil
		},
	)

	if HasLeaderElection(ctx) {
		t.Error("HasLeaderElection() = true, wanted false")
	}
	if le, err := BuildElector(ctx, laf, "name", enq); err != nil {
		t.Errorf("BuildElector() = %v, wanted an unopposedElector", err)
	} else if _, ok := le.(*unopposedElector); !ok {
		t.Errorf("BuildElector() = %T, wanted an unopposedElector", le)
	}

	ctx = WithDynamicLeaderElectorBuilder(ctx, kc, cc)
	if !HasLeaderElection(ctx) {
		t.Error("HasLeaderElection() = false, wanted true")
	}

	le, err := BuildElector(ctx, laf, "name", enq)
	if err != nil {
		t.Fatal("BuildElector() =", err)
	}

	// We shouldn't see leases until we Run the elector.
	select {
	case s := <-promoted:
		gotNames.Insert(s)
		t.Error("Got promoted, want no actions.")
	case <-demoted:
		t.Error("Got demoted, want no actions.")
	case <-created:
		t.Error("Got created, want no actions.")
	case <-updated:
		t.Error("Got updated, want no actions.")
	default:
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go le.Run(ctx)

	// We expect 3 lease to be created.
	for i := 0; i < buckets; i++ {
		select {
		case <-created:
		case <-time.After(1 * time.Second):
			t.Fatal("Timed out waiting for lease creation.")
		}
	}
	// We expect to have been promoted 3 times.
	for i := 0; i < buckets; i++ {
		select {
		case s := <-promoted:
			gotNames.Insert(s)
		case <-time.After(time.Second):
			t.Fatal("Timed out waiting for promotion.")
		}
	}

	// Cancelling the context should case us to give up leadership.
	cancel()

	select {
	case <-updated:
		// We expect the lease to be updated.
	case <-time.After(time.Second):
		t.Fatal("Timed out waiting for lease update.")
	}
	// We expect to have been demoted 3 times.
	for i := 0; i < buckets; i++ {
		select {
		case <-demoted:
		case <-time.After(time.Second):
			t.Fatal("Timed out waiting for demotion.")
		}
	}

	want := sets.NewString(
		"the-component.name.00-of-03",
		"the-component.name.01-of-03",
		"the-component.name.02-of-03",
	)
	if !gotNames.Equal(want) {
		t.Errorf("BucketSet.BucketList() = %q, want: %q", gotNames, want)
	}
}

func TestBuilderWithCustomizedLeaseName(t *testing.T) {
	const buckets = 3
	cc := ComponentConfig{
		Component:     "the-component",
		Buckets:       buckets,
		LeaseDuration: 15 * time.Second,
		RenewDeadline: 10 * time.Second,
		RetryPeriod:   2 * time.Second,
		LeaseName: func(i uint32) string {
			return fmt.Sprintf("bucket-%02d", i)
		},
	}
	kc := fakekube.NewSimpleClientset()
	ctx := context.Background()

	gotNames := make(sets.String, buckets)
	promoted := make(chan string)
	laf := &reconciler.LeaderAwareFuncs{
		PromoteFunc: func(bkt reconciler.Bucket, enq func(reconciler.Bucket, types.NamespacedName)) error {
			promoted <- bkt.Name()
			return nil
		},
	}
	enq := func(reconciler.Bucket, types.NamespacedName) {}
	ctx = WithDynamicLeaderElectorBuilder(ctx, kc, cc)
	le, err := BuildElector(ctx, laf, "name", enq)
	if err != nil {
		t.Fatal("BuildElector() =", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go le.Run(ctx)

	// We expect to have been promoted 3 times.
	for i := 0; i < buckets; i++ {
		select {
		case s := <-promoted:
			gotNames.Insert(s)
		case <-time.After(time.Second):
			t.Fatal("Timed out waiting for promotion.")
		}
	}

	want := sets.NewString(
		"bucket-00",
		"bucket-01",
		"bucket-02",
	)
	if !gotNames.Equal(want) {
		t.Errorf("BucketSet.BucketList() = %q, want: %q", gotNames, want)
	}
}

func TestNewStatefulSetBucketAndSet(t *testing.T) {
	wantNames := []string{
		"http://as-0.autoscaler.knative-testing.svc.cluster.local:80",
		"http://as-1.autoscaler.knative-testing.svc.cluster.local:80",
		"http://as-2.autoscaler.knative-testing.svc.cluster.local:80",
	}

	os.Setenv(controllerOrdinalEnv, "as-2")
	os.Setenv(serviceNameEnv, "autoscaler")
	t.Cleanup(func() {
		os.Unsetenv(controllerOrdinalEnv)
		os.Unsetenv(serviceNameEnv)
	})

	_, _, err := NewStatefulSetBucketAndSet(2)
	if err == nil {
		// Ordinal 2 should be range [0, 2)
		t.Fatal("Expected error from NewStatefulSetBucketAndSet but got nil")
	}

	bkt, bs, err := NewStatefulSetBucketAndSet(3)
	if err != nil {
		// Ordinal 2 should be range [0, 2)
		t.Fatal("NewStatefulSetBucketAndSet() = ", err)
	}

	if got, want := bkt.Name(), wantNames[2]; got != want {
		t.Errorf("Bucket.Name() = %s, want = %s", got, want)
	}

	gotNames := bs.BucketList()
	if !cmp.Equal(gotNames, wantNames) {
		t.Errorf("BucketSet.BucketList() = %q, want: %q", gotNames, wantNames)
	}
}

func TestWithStatefulSetBuilder(t *testing.T) {
	cc := ComponentConfig{
		Component: "the-component",
		Buckets:   3,
	}
	const podDNS = "http://as-2.autoscaler.knative-testing.svc.cluster.local:80"
	ctx := context.Background()

	promoted := make(chan struct{})
	laf := &reconciler.LeaderAwareFuncs{
		PromoteFunc: func(bkt reconciler.Bucket, enq func(reconciler.Bucket, types.NamespacedName)) error {
			close(promoted)
			return nil
		},
	}
	enq := func(reconciler.Bucket, types.NamespacedName) {}

	if os.Setenv(controllerOrdinalEnv, "as-2") != nil {
		t.Fatalf("Failed to set env var %s=%s", controllerOrdinalEnv, "as-2")
	}
	if os.Setenv(serviceNameEnv, "autoscaler") != nil {
		t.Fatalf("Failed to set env var %s=%s", serviceNameEnv, "autoscaler")
	}
	t.Cleanup(func() {
		os.Unsetenv(controllerOrdinalEnv)
		os.Unsetenv(serviceNameEnv)
	})

	ctx = WithDynamicLeaderElectorBuilder(ctx, nil, cc)
	if !HasLeaderElection(ctx) {
		t.Error("HasLeaderElection() = false, wanted true")
	}

	b := ctx.Value(builderKey{})
	ssb, ok := b.(*statefulSetBuilder)
	if !ok || ssb == nil {
		t.Fatal("StatefulSetBuilder not found on context")
	}

	le, err := BuildElector(ctx, laf, "name", enq)
	if err != nil {
		t.Fatal("BuildElector() =", err)
	}

	ule, ok := le.(*unopposedElector)
	if !ok {
		t.Fatalf("BuildElector() = %T, wanted an unopposedElector", le)
	}
	if got, want := ule.bkt.Name(), podDNS; got != want {
		t.Errorf("bkt.Name() = %s, wanted %s", got, want)
	}

	// Shouldn't be promoted until we Run the elector.
	select {
	case <-promoted:
		t.Error("Got promoted, want no actions.")
	default:
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go le.Run(ctx)

	select {
	case <-promoted:
		// We expect to have been promoted.
	case <-time.After(time.Second):
		t.Fatal("Timed out waiting for promotion.")
	}
}
<|endoftext|>