<commit_before>package models

import (
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"github.com/astaxie/beego/orm"
)

type App struct {
	Id              int    `orm:"column(id);auto"`
	Type            int8   `orm:"column(type)"`
	ApplicationName string `orm:"column(application_name);size(255)"`
	LifeCycle       string `orm:"column(life_cycle);size(255)"`
	Level           int8   `orm:"column(level)"`
	OwnerId         int    `orm:"column(owner_id)"`
	Default         bool   `orm:"column(default);null"`
}

func (t *App) TableName() string {
	return "app"
}

func init() {
	orm.RegisterModel(new(App))
}

// AddApp insert a new App into database and returns
// last inserted Id on success.
func AddApp(m *App) (id int64, err error) {
	o := orm.NewOrm()
	id, err = o.Insert(m)
	return
}

func AddDefApp(userId int) {
	if ExistByName("资源池") {
		return
	}
	defApp := new(App)
	defApp.Type = 0
	defApp.Level = 3
	defApp.ApplicationName = "资源池"
	defApp.LifeCycle = "公测"
	defApp.OwnerId = userId
	defApp.Default = true

	o := orm.NewOrm()
	if id, err := o.Insert(defApp); err == nil {
		s := new(Set)
		s.ApplicationID = int(id)
		s.SetName = "空闲机池"
		s.EnviType = 3
		s.ServiceStatus = 1
		s.Default = true
		s.Owner = userId
		if setId, err := AddSet(s); err == nil {
			m := new(Module)
			m.ApplicationId = int(id)
			m.SetId = int(setId)
			m.ModuleName = "空闲机"
			m.Owner = userId
			AddModule(m)
		}
	}
	return
}

func GetDefAppByUserId(userId int) (info map[string]interface{}, err error) {
	var app App
	var set Set
	var mod Module

	info = make(map[string]interface{})
	o := orm.NewOrm()

	if err = o.QueryTable("app").Filter("OwnerId", userId).Filter("Default", true).One(&app); err != nil {
		return
	}

	if err = o.QueryTable("set").Filter("ApplicationID", app.Id).One(&set); err != nil {
		return
	}

	if err = o.QueryTable("module").Filter("ApplicationId", app.Id).Filter("SetId", set.SetID).One(&mod); err != nil {
		return
	}

	fmt.Println("iiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiii app=", app, "set=", set, "mod=", mod)
	fmt.Println("app.id", app.Id)
	info["AppId"] = app.Id
	info["AppName"] = app.ApplicationName
	info["SetId"] = set.SetID
	info["SetName"] = set.SetName
	info["ModuleId"] = mod.Id
	info["ModuleName"] = mod.ModuleName
	return
}

// GetAppById retrieves App by Id. Returns error if
// Id doesn't exist
func GetAppById(id int) (v *App, err error) {
	o := orm.NewOrm()
	v = &App{Id: id}
	if err = o.Read(v); err == nil {
		return v, nil
	}
	return nil, err
}

func ExistByName(name string) bool {
	o := orm.NewOrm()
	return o.QueryTable("app").Filter("ApplicationName", name).Exist()
}

func GetAppTopoById(id int) (ml []interface{}, err error) {
	//	[{"id":"5524","text":"aaa","spriteCssClass":"c-icon icon-group","type":"set","expanded":false,"number":32,"items":[{"id":"7025","spriteCssClass":"c-icon icon-modal","text":"1","operator":"1842605324","bakoperator":"1842605324","type":"module","number":32}]}]
	var fields []string
	var sortby []string
	var order []string
	var query map[string]string = make(map[string]string)
	var limit int64 = 0
	var offset int64 = 0

	query["application_id"] = strconv.Itoa(id)
	app, _ := GetAppById(id)
	if app.Level == 3 {
		query["default"] = "0"
	} else {
		query["default"] = "1"
	}

	if sets, err := GetAllSet(query, fields, sortby, order, offset, limit); err == nil {
		for _, set := range sets {
			s := make(map[string]interface{})
			s["id"] = set.(Set).SetID
			s["text"] = set.(Set).SetName
			if app.Level == 3 {
				s["spriteCssClass"] = "c-icon icon-group"
				s["expanded"] = false
			} else {
				s["spriteCssClass"] = "c-icon icon-group hide"
				s["expanded"] = true
			}
			s["type"] = "set"

			s["number"] = 32
			var items []interface{}
			var fields []string
			var sortby []string
			var order []string
			var query map[string]string = make(map[string]string)
			var limit int64 = 0
			var offset int64 = 0

			query["application_id"] = strconv.Itoa(id)
			query["set_id"] = strconv.Itoa(set.(Set).SetID)
			if mods, err := GetAllModule(query, fields, sortby, order, offset, limit); err == nil {
				for _, mod := range mods {
					if mod.(Module).ModuleName != "空闲机" {
						m := make(map[string]interface{})
						m["id"] = mod.(Module).Id
						m["text"] = mod.(Module).ModuleName
						m["spriteCssClass"] = "c-icon icon-modal"
						m["type"] = "module"
						m["operator"] = mod.(Module).Operator
						m["bakoperator"] = mod.(Module).BakOperator
						m["number"] = 32
						items = append(items, m)
					}
				}
				s["items"] = items
			}
			ml = append(ml, s)
		}
	}

	return
}

func GetEmptyById(id int) (info map[string]interface{}, options map[int]string, err error) {
	info = make(map[string]interface{})
	options = make(map[int]string)
	var topo []interface{}
	var setItems []interface{}
	var emptyItems []interface{}
	var app *App
	var sets []interface{}
	var mods []interface{}

	setItems = append(setItems,1)[:0]
	// 1. 获取业务信息
	if app, err = GetAppById(id); err != nil {
		return
	}

	// 2.获取业务下集群信息，2级结构需要特殊处理
	var fields []string
	var sortby []string
	var order []string
	var query map[string]string = make(map[string]string)
	var limit int64 = 0
	var offset int64 = 0

	query["application_id"] = strconv.Itoa(id)

	if sets, err = GetAllSet(query, fields, sortby, order, offset, limit); err != nil {
		return
	}
	query["application_id"] = strconv.Itoa(id)

	if mods, err = GetAllModule(query, fields, sortby, order, offset, limit); err != nil {
		return
	}

	for _, set := range sets {
		s := set.(Set)
		setItem := make(map[string]interface{})
		setItem["id"] = s.SetID
		setItem["appId"] = app.Id
		setItem["Name"] = s.SetName
		setItem["spriteCssClass"] = "c-icon icon-group fa-hide set"
		setItem["type"] = "set"
		setItem["expanded"] = false
		setItem["number"], _ = GetHostCount(s.SetID, "SetID")
		var modItems []interface{}
		for _, mod := range mods {
			m := mod.(Module)
			if m.SetId == s.SetID {
				modItem := make(map[string]interface{})
				modItem["id"] = m.Id
				modItem["appId"] = app.Id
				modItem["Name"] = m.ModuleName
				modItem["spriteCssClass"] = "c-icon icon-modal module"
				modItem["type"] = "module"
				modItem["number"], _ = GetHostCount(m.Id, "ModuleID")
				modItems = append(modItems, modItem)
				if m.ModuleName == "空闲机" || m.ModuleName == "故障机" {
					emptyItems = append(emptyItems, modItem)
				} else {
					if app.Level == 2 {
						setItems = append(setItems, modItem)
						options[m.Id] = m.ModuleName
					} else {
						options[m.Id] = fmt.Sprintf("%s-%s", s.SetName, m.ModuleName)
					}
				}
			}
		}
		if !s.Default {
			setItem["items"] = modItems
			setItems = append(setItems, setItem)
		}
	}

	topoItem := make(map[string]interface{})
	topoItem["id"] = app.Id
	topoItem["Name"] = app.ApplicationName
	topoItem["type"] = "application"
	topoItem["spriteCssClass"] = "c-icon icon-app application"
	topoItem["expanded"] = true
	topoItem["lvl"] = app.Level
	topoItem["number"], _ = GetHostCount(app.Id, "ApplicationID")
	topoItem["items"] = setItems
	topo = append(topo, topoItem)
	info["topo"] = topo
	info["empty"] = emptyItems

	return
}

// GetAllApp retrieves all App matches certain condition. Returns empty list if
// no records exist
func GetAllApp(query map[string]string, fields []string, sortby []string, order []string,
	offset int64, limit int64) (ml []interface{}, err error) {
	o := orm.NewOrm()
	qs := o.QueryTable(new(App))
	// query k=v
	for k, v := range query {
		// rewrite dot-notation to Object__Attribute
		k = strings.Replace(k, ".", "__", -1)
		qs = qs.Filter(k, v)
	}
	// order by:
	var sortFields []string
	if len(sortby) != 0 {
		if len(sortby) == len(order) {
			// 1) for each sort field, there is an associated order
			for i, v := range sortby {
				orderby := ""
				if order[i] == "desc" {
					orderby = "-" + v
				} else if order[i] == "asc" {
					orderby = v
				} else {
					return nil, errors.New("Error: Invalid order. Must be either [asc|desc]")
				}
				sortFields = append(sortFields, orderby)
			}
			qs = qs.OrderBy(sortFields...)
		} else if len(sortby) != len(order) && len(order) == 1 {
			// 2) there is exactly one order, all the sorted fields will be sorted by this order
			for _, v := range sortby {
				orderby := ""
				if order[0] == "desc" {
					orderby = "-" + v
				} else if order[0] == "asc" {
					orderby = v
				} else {
					return nil, errors.New("Error: Invalid order. Must be either [asc|desc]")
				}
				sortFields = append(sortFields, orderby)
			}
		} else if len(sortby) != len(order) && len(order) != 1 {
			return nil, errors.New("Error: 'sortby', 'order' sizes mismatch or 'order' size is not 1")
		}
	} else {
		if len(order) != 0 {
			return nil, errors.New("Error: unused 'order' fields")
		}
	}

	var l []App
	qs = qs.OrderBy(sortFields...)
	if _, err := qs.Limit(limit, offset).All(&l, fields...); err == nil {
		if len(fields) == 0 {
			for _, v := range l {
				ml = append(ml, v)
			}
		} else {
			// trim unused fields
			for _, v := range l {
				m := make(map[string]interface{})
				val := reflect.ValueOf(v)
				for _, fname := range fields {
					m[fname] = val.FieldByName(fname).Interface()
				}
				ml = append(ml, m)
			}
		}
		return ml, nil
	}
	return nil, err
}

// UpdateApp updates App by Id and returns error if
// the record to be updated doesn't exist
func UpdateAppById(m *App) (err error) {
	o := orm.NewOrm()
	v := App{Id: m.Id}
	// ascertain id exists in the database
	if err = o.Read(&v); err == nil {
		var num int64
		if num, err = o.Update(m); err == nil {
			fmt.Println("Number of records updated in database:", num)
		}
	}
	return
}

// DeleteApp deletes App by Id and returns error if
// the record to be deleted doesn't exist
func DeleteApp(id int) (err error) {
	// TODO: 检测要删除业务下是否有主机
	o := orm.NewOrm()
	v := App{Id: id}
	DeleteSetByAppId(id)
	DeleteModuleByAppId(id)
	// ascertain id exists in the database
	if err = o.Read(&v); err == nil {
		var num int64
		if num, err = o.Delete(&App{Id: id}); err == nil {
			fmt.Println("Number of records deleted in database:", num)
		}
	}
	return
}
<commit_msg>暴露API接口工其他系统调用<commit_after>package models

import (
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"github.com/astaxie/beego/orm"
)

type App struct {
	Id              int    `orm:"column(id);auto"`
	Type            int8   `orm:"column(type)"`
	ApplicationName string `orm:"column(application_name);size(255)"`
	LifeCycle       string `orm:"column(life_cycle);size(255)"`
	Level           int8   `orm:"column(level)"`
	OwnerId         int    `orm:"column(owner_id)"`
	Default         bool   `orm:"column(default);null"`
}

func (t *App) TableName() string {
	return "app"
}

func init() {
	orm.RegisterModel(new(App))
}

// AddApp insert a new App into database and returns
// last inserted Id on success.
func AddApp(m *App) (id int64, err error) {
	o := orm.NewOrm()
	id, err = o.Insert(m)
	return
}

func AddDefApp(userId int) {
	if ExistByName("资源池") {
		return
	}
	defApp := new(App)
	defApp.Type = 0
	defApp.Level = 3
	defApp.ApplicationName = "资源池"
	defApp.LifeCycle = "公测"
	defApp.OwnerId = userId
	defApp.Default = true

	o := orm.NewOrm()
	if id, err := o.Insert(defApp); err == nil {
		s := new(Set)
		s.ApplicationID = int(id)
		s.SetName = "空闲机池"
		s.EnviType = 3
		s.ServiceStatus = 1
		s.Default = true
		s.Owner = userId
		if setId, err := AddSet(s); err == nil {
			m := new(Module)
			m.ApplicationId = int(id)
			m.SetId = int(setId)
			m.ModuleName = "空闲机"
			m.Owner = userId
			AddModule(m)
		}
	}
	return
}

func GetDefAppByUserId(userId int) (info map[string]interface{}, err error) {
	var app App
	var set Set
	var mod Module

	info = make(map[string]interface{})
	o := orm.NewOrm()

	if err = o.QueryTable("app").Filter("OwnerId", userId).Filter("Default", true).One(&app); err != nil {
		return
	}

	if err = o.QueryTable("set").Filter("ApplicationID", app.Id).One(&set); err != nil {
		return
	}

	if err = o.QueryTable("module").Filter("ApplicationId", app.Id).Filter("SetId", set.SetID).One(&mod); err != nil {
		return
	}

	fmt.Println("iiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiii app=", app, "set=", set, "mod=", mod)
	fmt.Println("app.id", app.Id)
	info["AppId"] = app.Id
	info["AppName"] = app.ApplicationName
	info["SetId"] = set.SetID
	info["SetName"] = set.SetName
	info["ModuleId"] = mod.Id
	info["ModuleName"] = mod.ModuleName
	return
}

// GetAppById retrieves App by Id. Returns error if
// Id doesn't exist
func GetAppById(id int) (v *App, err error) {
	o := orm.NewOrm()
	v = &App{Id: id}
	if err = o.Read(v); err == nil {
		return v, nil
	}
	return nil, err
}

func ExistByName(name string) bool {
	o := orm.NewOrm()
	return o.QueryTable("app").Filter("ApplicationName", name).Exist()
}

func GetAppTopoById(id int) (ml []interface{}, err error) {
	//	[{"id":"5524","text":"aaa","spriteCssClass":"c-icon icon-group","type":"set","expanded":false,"number":32,"items":[{"id":"7025","spriteCssClass":"c-icon icon-modal","text":"1","operator":"1842605324","bakoperator":"1842605324","type":"module","number":32}]}]
	var fields []string
	var sortby []string
	var order []string
	var query map[string]string = make(map[string]string)
	var limit int64 = 0
	var offset int64 = 0

	query["application_id"] = strconv.Itoa(id)
	app, _ := GetAppById(id)
	if app.Level == 3 {
		query["default"] = "0"
	} else {
		query["default"] = "1"
	}

	if sets, err := GetAllSet(query, fields, sortby, order, offset, limit); err == nil {
		for _, set := range sets {
			s := make(map[string]interface{})
			s["id"] = set.(Set).SetID
			s["text"] = set.(Set).SetName
			if app.Level == 3 {
				s["spriteCssClass"] = "c-icon icon-group"
				s["expanded"] = false
			} else {
				s["spriteCssClass"] = "c-icon icon-group hide"
				s["expanded"] = true
			}
			s["type"] = "set"

			s["number"] = 32
			var items []interface{}
			var fields []string
			var sortby []string
			var order []string
			var query map[string]string = make(map[string]string)
			var limit int64 = 0
			var offset int64 = 0

			query["application_id"] = strconv.Itoa(id)
			query["set_id"] = strconv.Itoa(set.(Set).SetID)
			if mods, err := GetAllModule(query, fields, sortby, order, offset, limit); err == nil {
				for _, mod := range mods {
					if mod.(Module).ModuleName != "空闲机" {
						m := make(map[string]interface{})
						m["id"] = mod.(Module).Id
						m["text"] = mod.(Module).ModuleName
						m["spriteCssClass"] = "c-icon icon-modal"
						m["type"] = "module"
						m["operator"] = mod.(Module).Operator
						m["bakoperator"] = mod.(Module).BakOperator
						m["number"] = 32
						items = append(items, m)
					}
				}
				s["items"] = items
			}
			ml = append(ml, s)
		}
	}

	return
}

func GetEmptyById(id int) (info map[string]interface{}, options map[int]string, err error) {
	info = make(map[string]interface{})
	options = make(map[int]string)
	var topo []interface{}
	var setItems []interface{}
	var emptyItems []interface{}
	var app *App
	var sets []interface{}
	var mods []interface{}

	setItems = append(setItems,1)[:0]
	// 1. 获取业务信息
	if app, err = GetAppById(id); err != nil {
		return
	}

	// 2.获取业务下集群信息，2级结构需要特殊处理
	var fields []string
	var sortby []string
	var order []string
	var query map[string]string = make(map[string]string)
	var limit int64 = 0
	var offset int64 = 0

	query["application_id"] = strconv.Itoa(id)

	if sets, err = GetAllSet(query, fields, sortby, order, offset, limit); err != nil {
		return
	}
	query["application_id"] = strconv.Itoa(id)

	if mods, err = GetAllModule(query, fields, sortby, order, offset, limit); err != nil {
		return
	}

	for _, set := range sets {
		s := set.(Set)
		setItem := make(map[string]interface{})
		setItem["id"] = s.SetID
		setItem["appId"] = app.Id
		setItem["Name"] = s.SetName
		setItem["spriteCssClass"] = "c-icon icon-group fa-hide set"
		setItem["type"] = "set"
		setItem["expanded"] = false
		setItem["number"], _ = GetHostCount(s.SetID, "SetID")
		var modItems []interface{}
		for _, mod := range mods {
			m := mod.(Module)
			if m.SetId == s.SetID {
				modItem := make(map[string]interface{})
				modItem["id"] = m.Id
				modItem["appId"] = app.Id
				modItem["Name"] = m.ModuleName
				modItem["spriteCssClass"] = "c-icon icon-modal module"
				modItem["type"] = "module"
				modItem["number"], _ = GetHostCount(m.Id, "ModuleID")
				modItems = append(modItems, modItem)
				if m.ModuleName == "空闲机" || m.ModuleName == "故障机" {
					emptyItems = append(emptyItems, modItem)
				} else {
					if app.Level == 2 {
						setItems = append(setItems, modItem)
						options[m.Id] = m.ModuleName
					} else {
						options[m.Id] = fmt.Sprintf("%s-%s", s.SetName, m.ModuleName)
					}
				}
			}
		}
		if !s.Default {
			setItem["items"] = modItems
			setItems = append(setItems, setItem)
		}
	}

	topoItem := make(map[string]interface{})
	topoItem["id"] = app.Id
	topoItem["Name"] = app.ApplicationName
	topoItem["type"] = "application"
	topoItem["spriteCssClass"] = "c-icon icon-app application"
	topoItem["expanded"] = true
	topoItem["lvl"] = app.Level
	topoItem["number"], _ = GetHostCount(app.Id, "ApplicationID")
	topoItem["items"] = setItems
	topo = append(topo, topoItem)
	info["topo"] = topo
	info["empty"] = emptyItems

	return
}

// GetAllApp retrieves all App matches certain condition. Returns empty list if
// no records exist
func GetAllApp(query map[string]string, fields []string, sortby []string, order []string,
	offset int64, limit int64) (ml []interface{}, err error) {
	o := orm.NewOrm()
	qs := o.QueryTable(new(App))
	// query k=v
	for k, v := range query {
		// rewrite dot-notation to Object__Attribute
		k = strings.Replace(k, ".", "__", -1)
		qs = qs.Filter(k, v)
	}
	// order by:
	var sortFields []string
	if len(sortby) != 0 {
		if len(sortby) == len(order) {
			// 1) for each sort field, there is an associated order
			for i, v := range sortby {
				orderby := ""
				if order[i] == "desc" {
					orderby = "-" + v
				} else if order[i] == "asc" {
					orderby = v
				} else {
					return nil, errors.New("Error: Invalid order. Must be either [asc|desc]")
				}
				sortFields = append(sortFields, orderby)
			}
			qs = qs.OrderBy(sortFields...)
		} else if len(sortby) != len(order) && len(order) == 1 {
			// 2) there is exactly one order, all the sorted fields will be sorted by this order
			for _, v := range sortby {
				orderby := ""
				if order[0] == "desc" {
					orderby = "-" + v
				} else if order[0] == "asc" {
					orderby = v
				} else {
					return nil, errors.New("Error: Invalid order. Must be either [asc|desc]")
				}
				sortFields = append(sortFields, orderby)
			}
		} else if len(sortby) != len(order) && len(order) != 1 {
			return nil, errors.New("Error: 'sortby', 'order' sizes mismatch or 'order' size is not 1")
		}
	} else {
		if len(order) != 0 {
			return nil, errors.New("Error: unused 'order' fields")
		}
	}

	var l []App
	qs = qs.OrderBy(sortFields...)
	if _, err := qs.Limit(limit, offset).All(&l, fields...); err == nil {
		if len(fields) == 0 {
			for _, v := range l {
				ml = append(ml, v)
			}
		} else {
			// trim unused fields
			for _, v := range l {
				m := make(map[string]interface{})
				val := reflect.ValueOf(v)
				for _, fname := range fields {
					m[fname] = val.FieldByName(fname).Interface()
				}
				ml = append(ml, m)
			}
		}
		return ml, nil
	}
	return nil, err
}

// UpdateApp updates App by Id and returns error if
// the record to be updated doesn't exist
func UpdateAppById(m *App) (err error) {
	o := orm.NewOrm()
	v := App{Id: m.Id}
	// ascertain id exists in the database
	if err = o.Read(&v); err == nil {
		var num int64
		if num, err = o.Update(m); err == nil {
			fmt.Println("Number of records updated in database:", num)
		}
	}
	return
}

// DeleteApp deletes App by Id and returns error if
// the record to be deleted doesn't exist
func DeleteApp(id int) (err error) {
	// TODO: 检测要删除业务下是否有主机
	o := orm.NewOrm()
	v := App{Id: id}
	DeleteSetByAppId(id)
	DeleteModuleByAppId(id)
	// ascertain id exists in the database
	if err = o.Read(&v); err == nil {
		var num int64
		if num, err = o.Delete(&App{Id: id}); err == nil {
			fmt.Println("Number of records deleted in database:", num)
		}
	}
	return
}

// 暴露的API接口要获取的数据
func GetCCModuleTree(id int) (ml []interface{}, err error) {
	//	[{"id":"5524","text":"aaa","spriteCssClass":"c-icon icon-group","type":"set","expanded":false,"number":32,"items":[{"id":"7025","spriteCssClass":"c-icon icon-modal","text":"1","operator":"1842605324","bakoperator":"1842605324","type":"module","number":32}]}]
	var fields []string
	var sortby []string
	var order []string
	var query map[string]string = make(map[string]string)
	var limit int64 = 0
	var offset int64 = 0

	query["application_id"] = strconv.Itoa(id)
	app, _ := GetAppById(id)
	if app.Level == 3 {
		query["default"] = "0"
	} else {
		query["default"] = "1"
	}

	if sets, err := GetAllSet(query, fields, sortby, order, offset, limit); err == nil {
		for _, set := range sets {
			s := make(map[string]interface{})
			s["setId"] = set.(Set).SetID
			s["text"] = set.(Set).SetName
			if app.Level == 3 {
				s["expanded"] = false
			} else {
				s["expanded"] = true
			}

			s["hostNum"], _ = GetHostCount(set.(Set).SetID, "SetID")
			var items []interface{}
			var fields []string
			var sortby []string
			var order []string
			var query map[string]string = make(map[string]string)
			var limit int64 = 0
			var offset int64 = 0

			query["application_id"] = strconv.Itoa(id)
			query["set_id"] = strconv.Itoa(set.(Set).SetID)
			if mods, err := GetAllModule(query, fields, sortby, order, offset, limit); err == nil {
				for _, mod := range mods {
					if mod.(Module).ModuleName != "空闲机" {
						m := make(map[string]interface{})
						m["moduleId"] = mod.(Module).Id
						m["text"] = mod.(Module).ModuleName
						m["hostNum"], _ = GetHostCount(mod.(Module).Id, "ModuleID")
						items = append(items, m)
					}
				}
				s["items"] = items
			}
			ml = append(ml, s)
		}
	}

	return
}
<|endoftext|>


<commit_before>package charm_test

import (
	"bytes"
	"io"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/goyaml"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/testing"
	stdtesting "testing"
)

func Test(t *stdtesting.T) {
	TestingT(t)
}

type CharmSuite struct{}

var _ = Suite(&CharmSuite{})

func (s *CharmSuite) TestRead(c *C) {
	bPath := testing.Charms.BundlePath(c.MkDir(), "dummy")
	ch, err := charm.Read(bPath)
	c.Assert(err, IsNil)
	c.Assert(ch.Meta().Name, Equals, "dummy")
	dPath := testing.Charms.DirPath("dummy")
	ch, err = charm.Read(dPath)
	c.Assert(err, IsNil)
	c.Assert(ch.Meta().Name, Equals, "dummy")
}

var inferRepoTests = []struct {
	url  string
	path string
}{
	{"cs:precise/wordpress", ""},
	{"local:oneiric/wordpress", "/some/path"},
}

func (s *CharmSuite) TestInferRepository(c *C) {
	defer config.SetJujuHome(config.SetJujuHome(c.MkDir()))
	for i, t := range inferRepoTests {
		c.Logf("test %d", i)
		curl, err := charm.InferURL(t.url, "precise")
		c.Assert(err, IsNil)
		repo, err := charm.InferRepository(curl, "/some/path")
		c.Assert(err, IsNil)
		switch repo := repo.(type) {
		case *charm.LocalRepository:
			c.Assert(repo.Path, Equals, t.path)
		default:
			c.Assert(repo, FitsTypeOf, charm.Store())
		}
	}
	curl, err := charm.InferURL("local:whatever", "precise")
	c.Assert(err, IsNil)
	_, err = charm.InferRepository(curl, "")
	c.Assert(err, ErrorMatches, "path to local repository not specified")
	curl.Schema = "foo"
	_, err = charm.InferRepository(curl, "")
	c.Assert(err, ErrorMatches, "unknown schema for charm URL.*")
}

func checkDummy(c *C, f charm.Charm, path string) {
	c.Assert(f.Revision(), Equals, 1)
	c.Assert(f.Meta().Name, Equals, "dummy")
	c.Assert(f.Config().Options["title"].Default, Equals, "My Title")
	switch f := f.(type) {
	case *charm.Bundle:
		c.Assert(f.Path, Equals, path)

	case *charm.Dir:
		c.Assert(f.Path, Equals, path)
	}
}

type YamlHacker map[interface{}]interface{}

func ReadYaml(r io.Reader) YamlHacker {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		panic(err)
	}
	m := make(map[interface{}]interface{})
	err = goyaml.Unmarshal(data, m)
	if err != nil {
		panic(err)
	}
	return YamlHacker(m)
}

func (yh YamlHacker) Reader() io.Reader {
	data, err := goyaml.Marshal(yh)
	if err != nil {
		panic(err)
	}
	return bytes.NewBuffer(data)
}
<commit_msg>config: merged trunk and again some review changes<commit_after>package charm_test

import (
	"bytes"
	"io"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/goyaml"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/testing"
	stdtesting "testing"
)

func Test(t *stdtesting.T) {
	TestingT(t)
}

type CharmSuite struct{}

var _ = Suite(&CharmSuite{})

func (s *CharmSuite) TestRead(c *C) {
	bPath := testing.Charms.BundlePath(c.MkDir(), "dummy")
	ch, err := charm.Read(bPath)
	c.Assert(err, IsNil)
	c.Assert(ch.Meta().Name, Equals, "dummy")
	dPath := testing.Charms.DirPath("dummy")
	ch, err = charm.Read(dPath)
	c.Assert(err, IsNil)
	c.Assert(ch.Meta().Name, Equals, "dummy")
}

var inferRepoTests = []struct {
	url  string
	path string
}{
	{"cs:precise/wordpress", ""},
	{"local:oneiric/wordpress", "/some/path"},
}

func (s *CharmSuite) TestInferRepository(c *C) {
	for i, t := range inferRepoTests {
		c.Logf("test %d", i)
		curl, err := charm.InferURL(t.url, "precise")
		c.Assert(err, IsNil)
		repo, err := charm.InferRepository(curl, "/some/path")
		c.Assert(err, IsNil)
		switch repo := repo.(type) {
		case *charm.LocalRepository:
			c.Assert(repo.Path, Equals, t.path)
		default:
			c.Assert(repo, FitsTypeOf, charm.Store())
		}
	}
	curl, err := charm.InferURL("local:whatever", "precise")
	c.Assert(err, IsNil)
	_, err = charm.InferRepository(curl, "")
	c.Assert(err, ErrorMatches, "path to local repository not specified")
	curl.Schema = "foo"
	_, err = charm.InferRepository(curl, "")
	c.Assert(err, ErrorMatches, "unknown schema for charm URL.*")
}

func checkDummy(c *C, f charm.Charm, path string) {
	c.Assert(f.Revision(), Equals, 1)
	c.Assert(f.Meta().Name, Equals, "dummy")
	c.Assert(f.Config().Options["title"].Default, Equals, "My Title")
	switch f := f.(type) {
	case *charm.Bundle:
		c.Assert(f.Path, Equals, path)

	case *charm.Dir:
		c.Assert(f.Path, Equals, path)
	}
}

type YamlHacker map[interface{}]interface{}

func ReadYaml(r io.Reader) YamlHacker {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		panic(err)
	}
	m := make(map[interface{}]interface{})
	err = goyaml.Unmarshal(data, m)
	if err != nil {
		panic(err)
	}
	return YamlHacker(m)
}

func (yh YamlHacker) Reader() io.Reader {
	data, err := goyaml.Marshal(yh)
	if err != nil {
		panic(err)
	}
	return bytes.NewBuffer(data)
}
<|endoftext|>


<commit_before>package neustar

import (
	"encoding/json"
	"fmt"
	"net/http"
)

const (
	// MonitorURI is the endpoint for calls to the monitoring API
	MonitorURI = "monitor/1.0"

	// LocationsURI is the endpoint for locations calls
	LocationsURI = "/locations"

	// SummaryURI is the summary endpoint for summary calls
	SummaryURI = "/summary"
)

// MonitorTypes is a slice of valid monitor types
var MonitorTypes = []string{"RealBrowserUser", "VirtualUser", "dns"}

// BrowserTypes is a slice of valid browser types
var BrowserTypes = []string{"FF", "CHROME", "IE"}

// UpdateIntervals is a slice of valid intervals
var UpdateIntervals = []int{1, 2, 3, 4, 5, 10, 15, 20, 30, 60}

// UpdateMonitorParameters holds the allowed options for updating
// a monitor
type UpdateMonitorParameters struct {
	Name        string `json:"name"`
	Description string `json:"name"`
	Interval    int    `json:"interval"`
	TestScript  string `json:"testScript"`
	Locations   string `json:"locations"`
	AlertPolicy string `json:"alertPolicy"`
	Browser     string `json:"browser"`
	Active      string `json:"Active"`
}

// CreateMonitorParameters holds the parameters needed by the create
// monitor endpoint
type CreateMonitorParameters struct {
	// The name of the monitor
	Name string `json:"name"`

	// A description of what this monitor is for
	Description string `json:"description"`

	// How often the monitoring script will run for each of the locations
	Interval int `json:"interval"`

	// The id of the test script that this monitor should run
	TestScript string `json:"testScript"`

	// A CSV list of locations that this monitor should run from
	Locations string `json:"Locations"`

	// The id of the alert policy that this monitor should run
	AlertPolicy string `json:"alertPolicy"`

	// Specifies the browser type that this monitor should use. Note: IE is
	// available for Enterprise customers only
	Browser string `json:"browser"`

	// Enables or disables this monitor from taking samples
	Active string `json:"active"`

	// Set to network monitor type such as 'dns'. See related settings below.
	// Leave this blank for script-based monitors. Note: this interface will
	// not allow you to test network monitor creation. Please use your API client.
	Type string `json:"type"`

	DNSSettings  DNSSettings  `json:"dnsSettings"`
	PingSettings PingSettings `json:"pingSettings"`
	PopSettings  PopSettings  `json:"popSettings"`
	PortSettings PortSettings `json:"portSettings"`
	SMTPSettings SMTPSettings `json:"smtpSettings"`
}

// AggregateParameters holds the allowed options for getting
// aggregate sample data
type AggregateParameters struct {
	StartDate string `json:"startDate"`
	EndDate   string `json:"endDate"`
	Offset    int    `json:"offset"`
	Frequency string `json:"frequency"`
	GroupBy   string `json:"groupBy"`
}

// AggregateSampleDataResponse holds the returned data fro mthe call
type AggregateSampleDataResponse struct {
	Count      int         `json:"count"`
	Uptime     float64     `json:"uptime"`
	Min        int         `json:"min"`
	Max        int         `json:"max"`
	Date       interface{} `json:"date"`
	Avg        float64     `json:"avg"`
	STDDev     float64     `json:"stdDev"`
	Location   string      `json:"location"`
	StepName   string      `json:"stepName"`
	StepNumber int         `json:"stepNumber"`
	TP50       float64     `json:"tp50"`
	TP90       float64     `json:"tp90"`
}

// DNSSettings is a an object containing all DNS-related settings:
// {"timeout": int, "lookups": array}. The "lookups" array contains
// JSON objects with this format: {"lookupType": string ("A" or "AAAA"),
// "authoritative": boolean, "hostname": string, "dnsServer": string, "expectedIps":
// string of comma-separated IP addresses}
type DNSSettings struct {
	LookupType    string `json:"lookupType"`
	Authoritative bool   `json:"authoritative"`
	Hostname      string `json:"hostname"`
	DNSServer     string `json:"dnsServer"`
	ExpectedIPs   string `json:"expectedIps"`
}

// PingSettings is an object containing all PING-related settings: {"timeout": int, "host": string}.
type PingSettings struct {
	Timeout int    `json:"timeout"`
	Host    string `json:"host"`
}

// PortSettings holds port settings
type PortSettings struct {
	Timeout          int    `json:"timeout"`
	Server           string `json:"server"`
	Port             int    `json:"port"`
	Protocol         string `json:"protocol"`
	Command          string `json:"command"`
	ExpectedResponse string `json:"expected_response"`
	DataFormat       string `json:"data_format"`
}

// SMTPSettings holds SMTP setttings
type SMTPSettings struct {
	Timeout int    `json:"timeout"`
	Server  string `json:"server"`
	Email   string `json:"email"`
}

// PopSettings holds port settings
type PopSettings struct {
	Timeout  int    `json:"timeout"`
	Server   string `json:"server"`
	Username string `json:"username"`
	Password string `json:"password"`
}

// Monitoring holds monitoring config
type Monitoring struct {
	config *Configuration
}

// Monitor hold monitoring data
type Monitor struct {
	// The ID of the monitor
	ID string `json:"id"`

	// The ID of the alerting policy associated with this monitor
	AlertPolicy interface{} `json:"alertPolicy"`

	// A list of monitoring locations that this monitor is run from
	Locations []string `json:"locations"`

	// The version, id and name of the script associated with this monitor
	Script Script `json:"script"`

	// The description of the monitor
	Description string `json:"description"`

	// How often this monitor runs
	Interval int `json:"interval"`

	// The name of this monitor
	Name string `json:"name"`

	// The time of the last monitoring sample for this monitor
	LastSampleAt interface{} `json:"lastSampleAt"`

	// Describes whether this monitor is actively monitoring or not
	Active bool `json:"active"`

	// Whether this monitor is in a maintenance window or not
	InMaintenanceWindow bool `json:"inMaintenanceWindow"`

	// Describes the type of browser that monitor is using, 'FF' for
	// Firefox or 'CHROME' for Chrome, or 'IE' for Internet Explorer
	Browser string `json:"browser"`

	// The type of monitor ('RealBrowserUser', 'VirtualUser', 'dns')
	Type string `json:"type"`

	SMTPSettings SMTPSettings `json:"smtpSettings"`
	SLASettings  interface{}  `json:"slaSettings,omitempty"`
	DNSSettings  DNSSettings  `json:"dnsSettings,omitempty"`
	PopSettings  PopSettings  `json:"popSettings,omitempty"`
	PortSettings PortSettings `json:"portSettings,omitempty"`
	PingSettings PingSettings `json:"pingSettings"`
}

// NewMonitor creates a new Monitoring object
func NewMonitor(key, secret string) *Monitoring {
	return &Monitoring{
		config: LoadConfig(key, secret),
	}
}

// Create creates a new monitor and returns the monitor id of the newly
// created monitor. Name, interval, testScript and locations are required.
// Use the Get Monitoring Locations api to retrieve a list of monitoring locations.
func (m *Monitoring) Create() {}

// List retrieves a list of all monitors associated with your account,
//along with information about each. The monitor id that is returned
// is used to make other api calls.
func (m *Monitoring) List() ([]Monitor, int, error) {
	var response *http.Response
	var data map[string]map[string][]Monitor
	response, err := http.Get(fmt.Sprintf("%s%s?apikey=%s&sig=%s", BaseURL, MonitorURI, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return nil, response.StatusCode, err
	}
	defer response.Body.Close()
	if err := json.NewDecoder(response.Body).Decode(&data); err != nil {
		return nil, response.StatusCode, err
	}
	return data["data"]["items"], response.StatusCode, nil
}

// Get retrieves information for a specific monitor associated with your
// account. The monitor id that is returned is used to make other api calls.
func (m *Monitoring) Get(id string) ([]Monitor, int, error) {
	var response *http.Response
	var data map[string]map[string][]Monitor
	response, err := http.Get(fmt.Sprintf("%s%s/%s?apikey=%s&sig=%s", BaseURL, MonitorURI, id, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return nil, response.StatusCode, err
	}
	defer response.Body.Close()
	if err := json.NewDecoder(response.Body).Decode(&data); err != nil {
		return nil, response.StatusCode, err
	}
	return data["data"]["items"], response.StatusCode, nil
}

// Update changes some or all of the parameters of an existing monitor.
// Requires the monitor ID retrieved from the List Monitors api.
func (m *Monitoring) Update() {}

// Delete deletes the given monitor, stopping it from monitoring and removing
// all its monitoring data.
func (m *Monitoring) Delete(id string) (int, error) {
	var response *http.Response
	response, err := http.Get(fmt.Sprintf("%s%s/%s?apikey=%s&sig=%s", BaseURL, MonitorURI, id, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return response.StatusCode, err
	}
	if response.StatusCode != 200 {
		return response.StatusCode, nil
	}
	return response.StatusCode, nil
}

// Samples returns all samples associated to this monitor for a given time period.
// This data is returned at a high level, which timing for the overall sample. To
// get the details for the specific sample, call the get raw sample data api. At a
// maximum, this api will return 2000 samples. If there are more than 2000 results
// returned, the 'more' field will be set to true and you can make another api call
// specifying an offset which would be equal to the number of results returned in the
// first api call plus the offset of that call.
func (m *Monitoring) Samples() ([]string, int, error) {
	var response *http.Response
	var data map[string]map[string][]string
	response, err := http.Get(fmt.Sprintf("%s%s?apikey=%s&sig=%s", BaseURL, MonitorURI, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return nil, response.StatusCode, err
	}
	defer response.Body.Close()
	if err := json.NewDecoder(response.Body).Decode(&data); err != nil {
		return nil, response.StatusCode, err
	}
	return data["data"]["items"], response.StatusCode, nil
}

// RawSampleData retrieves the raw, HTTP Archive (HAR) data for a particular sample
func (m *Monitoring) RawSampleData(monitorID, sampleID string) {}

// AggregateSampleData retrieves the aggregated sample information for a given period
// of time. You can choose to aggregate the data for each hour or each day. This is
// more effecient than getting all the individual samples for a period of time and
// performing the aggregation yourself.
func (m *Monitoring) AggregateSampleData(monitorID string, asd AggregateSampleDataResponse) (int, error) {
	var response *http.Response
	response, err := http.Get(fmt.Sprintf("%s%s/%s?apikey=%s&sig=%s", BaseURL, MonitorURI, monitorID, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return response.StatusCode, err
	}
	if response.StatusCode != 200 {
		return response.StatusCode, nil
	}
	return response.StatusCode, nil
}

// Summary provides the monitor summary api returns all of the data that is found when looking at your
// list of monitors in the web portal. This includes things such as the average load
// time, sample count and uptime for the day, week, month or year, the last time an
// error occurred, and the last error message.
func (m *Monitoring) Summary(monitorID string) {}

// Locations gets a list of all monitoring locations available
func (m *Monitoring) Locations() ([]string, int, error) {
	var response *http.Response
	var data map[string]map[string][]string
	response, err := http.Get(fmt.Sprintf("%s%s%s?apikey=%s&sig=%s", BaseURL, MonitorURI, LocationsURI, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return nil, response.StatusCode, err
	}
	defer response.Body.Close()
	if err := json.NewDecoder(response.Body).Decode(&data); err != nil {
		return nil, response.StatusCode, err
	}
	return data["data"]["items"], response.StatusCode, nil
}

// ValidMonitorType validates the given monitor type is valid
func ValidMonitorType(monitorType string) bool { return false }

// ValidBrowserType validates the given browser type is valid
func ValidBrowserType(browserType string) bool { return false }

// ValidUpdateInterval validates the given interval is valid
func ValidUpdateInterval(interval int) bool { return false }
<commit_msg>validation functions added for monitoring<commit_after>package neustar

import (
	"encoding/json"
	"fmt"
	"net/http"
)

const (
	// MonitorURI is the endpoint for calls to the monitoring API
	MonitorURI = "monitor/1.0"

	// LocationsURI is the endpoint for locations calls
	LocationsURI = "/locations"

	// SummaryURI is the summary endpoint for summary calls
	SummaryURI = "/summary"
)

// MonitorTypes is a slice of valid monitor types
var MonitorTypes = []string{"RealBrowserUser", "VirtualUser", "dns"}

// BrowserTypes is a slice of valid browser types
var BrowserTypes = []string{"FF", "CHROME", "IE"}

// UpdateIntervals is a slice of valid intervals
var UpdateIntervals = []int{1, 2, 3, 4, 5, 10, 15, 20, 30, 60}

// UpdateMonitorParameters holds the allowed options for updating
// a monitor
type UpdateMonitorParameters struct {
	Name        string `json:"name"`
	Description string `json:"name"`
	Interval    int    `json:"interval"`
	TestScript  string `json:"testScript"`
	Locations   string `json:"locations"`
	AlertPolicy string `json:"alertPolicy"`
	Browser     string `json:"browser"`
	Active      string `json:"Active"`
}

// CreateMonitorParameters holds the parameters needed by the create
// monitor endpoint
type CreateMonitorParameters struct {
	// The name of the monitor
	Name string `json:"name"`

	// A description of what this monitor is for
	Description string `json:"description"`

	// How often the monitoring script will run for each of the locations
	Interval int `json:"interval"`

	// The id of the test script that this monitor should run
	TestScript string `json:"testScript"`

	// A CSV list of locations that this monitor should run from
	Locations string `json:"Locations"`

	// The id of the alert policy that this monitor should run
	AlertPolicy string `json:"alertPolicy"`

	// Specifies the browser type that this monitor should use. Note: IE is
	// available for Enterprise customers only
	Browser string `json:"browser"`

	// Enables or disables this monitor from taking samples
	Active string `json:"active"`

	// Set to network monitor type such as 'dns'. See related settings below.
	// Leave this blank for script-based monitors. Note: this interface will
	// not allow you to test network monitor creation. Please use your API client.
	Type string `json:"type"`

	DNSSettings  DNSSettings  `json:"dnsSettings"`
	PingSettings PingSettings `json:"pingSettings"`
	PopSettings  PopSettings  `json:"popSettings"`
	PortSettings PortSettings `json:"portSettings"`
	SMTPSettings SMTPSettings `json:"smtpSettings"`
}

// AggregateParameters holds the allowed options for getting
// aggregate sample data
type AggregateParameters struct {
	StartDate string `json:"startDate"`
	EndDate   string `json:"endDate"`
	Offset    int    `json:"offset"`
	Frequency string `json:"frequency"`
	GroupBy   string `json:"groupBy"`
}

// AggregateSampleDataResponse holds the returned data fro mthe call
type AggregateSampleDataResponse struct {
	Count      int         `json:"count"`
	Uptime     float64     `json:"uptime"`
	Min        int         `json:"min"`
	Max        int         `json:"max"`
	Date       interface{} `json:"date"`
	Avg        float64     `json:"avg"`
	STDDev     float64     `json:"stdDev"`
	Location   string      `json:"location"`
	StepName   string      `json:"stepName"`
	StepNumber int         `json:"stepNumber"`
	TP50       float64     `json:"tp50"`
	TP90       float64     `json:"tp90"`
}

// DNSSettings is a an object containing all DNS-related settings:
// {"timeout": int, "lookups": array}. The "lookups" array contains
// JSON objects with this format: {"lookupType": string ("A" or "AAAA"),
// "authoritative": boolean, "hostname": string, "dnsServer": string, "expectedIps":
// string of comma-separated IP addresses}
type DNSSettings struct {
	LookupType    string `json:"lookupType"`
	Authoritative bool   `json:"authoritative"`
	Hostname      string `json:"hostname"`
	DNSServer     string `json:"dnsServer"`
	ExpectedIPs   string `json:"expectedIps"`
}

// PingSettings is an object containing all PING-related settings: {"timeout": int, "host": string}.
type PingSettings struct {
	Timeout int    `json:"timeout"`
	Host    string `json:"host"`
}

// PortSettings holds port settings
type PortSettings struct {
	Timeout          int    `json:"timeout"`
	Server           string `json:"server"`
	Port             int    `json:"port"`
	Protocol         string `json:"protocol"`
	Command          string `json:"command"`
	ExpectedResponse string `json:"expected_response"`
	DataFormat       string `json:"data_format"`
}

// SMTPSettings holds SMTP setttings
type SMTPSettings struct {
	Timeout int    `json:"timeout"`
	Server  string `json:"server"`
	Email   string `json:"email"`
}

// PopSettings holds port settings
type PopSettings struct {
	Timeout  int    `json:"timeout"`
	Server   string `json:"server"`
	Username string `json:"username"`
	Password string `json:"password"`
}

// Monitoring holds monitoring config
type Monitoring struct {
	config *Configuration
}

// Monitor hold monitoring data
type Monitor struct {
	// The ID of the monitor
	ID string `json:"id"`

	// The ID of the alerting policy associated with this monitor
	AlertPolicy interface{} `json:"alertPolicy"`

	// A list of monitoring locations that this monitor is run from
	Locations []string `json:"locations"`

	// The version, id and name of the script associated with this monitor
	Script Script `json:"script"`

	// The description of the monitor
	Description string `json:"description"`

	// How often this monitor runs
	Interval int `json:"interval"`

	// The name of this monitor
	Name string `json:"name"`

	// The time of the last monitoring sample for this monitor
	LastSampleAt interface{} `json:"lastSampleAt"`

	// Describes whether this monitor is actively monitoring or not
	Active bool `json:"active"`

	// Whether this monitor is in a maintenance window or not
	InMaintenanceWindow bool `json:"inMaintenanceWindow"`

	// Describes the type of browser that monitor is using, 'FF' for
	// Firefox or 'CHROME' for Chrome, or 'IE' for Internet Explorer
	Browser string `json:"browser"`

	// The type of monitor ('RealBrowserUser', 'VirtualUser', 'dns')
	Type string `json:"type"`

	SMTPSettings SMTPSettings `json:"smtpSettings"`
	SLASettings  interface{}  `json:"slaSettings,omitempty"`
	DNSSettings  DNSSettings  `json:"dnsSettings,omitempty"`
	PopSettings  PopSettings  `json:"popSettings,omitempty"`
	PortSettings PortSettings `json:"portSettings,omitempty"`
	PingSettings PingSettings `json:"pingSettings"`
}

// NewMonitor creates a new Monitoring object
func NewMonitor(key, secret string) *Monitoring {
	return &Monitoring{
		config: LoadConfig(key, secret),
	}
}

// Create creates a new monitor and returns the monitor id of the newly
// created monitor. Name, interval, testScript and locations are required.
// Use the Get Monitoring Locations api to retrieve a list of monitoring locations.
func (m *Monitoring) Create() {}

// List retrieves a list of all monitors associated with your account,
//along with information about each. The monitor id that is returned
// is used to make other api calls.
func (m *Monitoring) List() ([]Monitor, int, error) {
	var response *http.Response
	var data map[string]map[string][]Monitor
	response, err := http.Get(fmt.Sprintf("%s%s?apikey=%s&sig=%s", BaseURL, MonitorURI, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return nil, response.StatusCode, err
	}
	defer response.Body.Close()
	if err := json.NewDecoder(response.Body).Decode(&data); err != nil {
		return nil, response.StatusCode, err
	}
	return data["data"]["items"], response.StatusCode, nil
}

// Get retrieves information for a specific monitor associated with your
// account. The monitor id that is returned is used to make other api calls.
func (m *Monitoring) Get(id string) ([]Monitor, int, error) {
	var response *http.Response
	var data map[string]map[string][]Monitor
	response, err := http.Get(fmt.Sprintf("%s%s/%s?apikey=%s&sig=%s", BaseURL, MonitorURI, id, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return nil, response.StatusCode, err
	}
	defer response.Body.Close()
	if err := json.NewDecoder(response.Body).Decode(&data); err != nil {
		return nil, response.StatusCode, err
	}
	return data["data"]["items"], response.StatusCode, nil
}

// Update changes some or all of the parameters of an existing monitor.
// Requires the monitor ID retrieved from the List Monitors api.
func (m *Monitoring) Update() {}

// Delete deletes the given monitor, stopping it from monitoring and removing
// all its monitoring data.
func (m *Monitoring) Delete(id string) (int, error) {
	var response *http.Response
	response, err := http.Get(fmt.Sprintf("%s%s/%s?apikey=%s&sig=%s", BaseURL, MonitorURI, id, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return response.StatusCode, err
	}
	if response.StatusCode != 200 {
		return response.StatusCode, nil
	}
	return response.StatusCode, nil
}

// Samples returns all samples associated to this monitor for a given time period.
// This data is returned at a high level, which timing for the overall sample. To
// get the details for the specific sample, call the get raw sample data api. At a
// maximum, this api will return 2000 samples. If there are more than 2000 results
// returned, the 'more' field will be set to true and you can make another api call
// specifying an offset which would be equal to the number of results returned in the
// first api call plus the offset of that call.
func (m *Monitoring) Samples() ([]string, int, error) {
	var response *http.Response
	var data map[string]map[string][]string
	response, err := http.Get(fmt.Sprintf("%s%s?apikey=%s&sig=%s", BaseURL, MonitorURI, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return nil, response.StatusCode, err
	}
	defer response.Body.Close()
	if err := json.NewDecoder(response.Body).Decode(&data); err != nil {
		return nil, response.StatusCode, err
	}
	return data["data"]["items"], response.StatusCode, nil
}

// RawSampleData retrieves the raw, HTTP Archive (HAR) data for a particular sample
func (m *Monitoring) RawSampleData(monitorID, sampleID string) {}

// AggregateSampleData retrieves the aggregated sample information for a given period
// of time. You can choose to aggregate the data for each hour or each day. This is
// more effecient than getting all the individual samples for a period of time and
// performing the aggregation yourself.
func (m *Monitoring) AggregateSampleData(monitorID string, asd AggregateSampleDataResponse) (int, error) {
	var response *http.Response
	response, err := http.Get(fmt.Sprintf("%s%s/%s?apikey=%s&sig=%s", BaseURL, MonitorURI, monitorID, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return response.StatusCode, err
	}
	if response.StatusCode != 200 {
		return response.StatusCode, nil
	}
	return response.StatusCode, nil
}

// Summary provides the monitor summary api returns all of the data that is found when looking at your
// list of monitors in the web portal. This includes things such as the average load
// time, sample count and uptime for the day, week, month or year, the last time an
// error occurred, and the last error message.
func (m *Monitoring) Summary(monitorID string) {}

// Locations gets a list of all monitoring locations available
func (m *Monitoring) Locations() ([]string, int, error) {
	var response *http.Response
	var data map[string]map[string][]string
	response, err := http.Get(fmt.Sprintf("%s%s%s?apikey=%s&sig=%s", BaseURL, MonitorURI, LocationsURI, m.config.API.Key, m.config.DigitalSignature()))
	if err != nil {
		return nil, response.StatusCode, err
	}
	defer response.Body.Close()
	if err := json.NewDecoder(response.Body).Decode(&data); err != nil {
		return nil, response.StatusCode, err
	}
	return data["data"]["items"], response.StatusCode, nil
}

// ValidMonitorType validates the given monitor type is valid
func ValidMonitorType(monitorType string) bool {
	for _, i := range MonitorTypes {
		if i == monitorType {
			return true
		}
	}
	return false
}

// ValidBrowserType validates the given browser type is valid
func ValidBrowserType(browserType string) bool {
	for _, i := range BrowserTypes {
		if i == browserType {
			return true
		}
	}
	return false
}

// ValidUpdateInterval validates the given interval is valid
func ValidUpdateInterval(interval int) bool {
	for _, i := range UpdateIntervals {
		if i == interval {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package connections

import (
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/pquerna/ffjson/ffjson"
	"github.com/sirupsen/logrus"

	"github.com/ivan1993spb/snake-server/game"
)

const (
	chanOutputMessageBuffer      = 128
	chanReadMessagesBuffer       = 128
	chanDecodeMessageBuffer      = 128
	chanEncodeMessageBuffer      = 128
	chanProxyInputMessageBuffer  = 64
	chanListenInputMessageBuffer = 32

	sendInputMessageTimeout = time.Millisecond * 50
)

type ConnectionWorker struct {
	conn   *websocket.Conn
	logger *logrus.Logger

	chStop      chan struct{}
	chOutput    chan OutputMessage
	chsInput    []chan InputMessage
	chsInputMux *sync.RWMutex

	flagStarted bool
}

func NewConnectionWorker(conn *websocket.Conn, logger *logrus.Logger) *ConnectionWorker {
	return &ConnectionWorker{
		conn:        conn,
		logger:      logger,
		chStop:      make(chan struct{}, 0),
		chOutput:    make(chan OutputMessage, chanOutputMessageBuffer),
		chsInput:    make([]chan InputMessage, 0),
		chsInputMux: &sync.RWMutex{},
	}
}

type ErrStartConnectionWorker string

func (e ErrStartConnectionWorker) Error() string {
	return "error start connection worker: " + string(e)
}

func (cw *ConnectionWorker) Start(game *game.Game) error {
	if cw.flagStarted {
		return ErrStartConnectionWorker("connection worker already started")
	}

	cw.flagStarted = true

	chInputBytes := cw.read()
	chInputMessages := cw.decode(chInputBytes, cw.chStop)
	cw.broadcastInputMessage(chInputMessages, cw.chStop)
	cw.listen(game.Events(cw.chStop), cw.chStop)
	chOutputBytes := cw.encode(cw.chOutput, cw.chStop)
	cw.write(chOutputBytes, cw.chStop)

	return nil
}

func (cw *ConnectionWorker) listenErrors(cherr <-chan error) <-chan struct{} {
	stop := make(chan struct{}, 0)

	go func() {
		for err := range cherr {
			if err != nil {
				close(stop)
				return
			}
		}
	}()

	return stop
}

func (cw *ConnectionWorker) read() <-chan []byte {
	chout := make(chan []byte, chanReadMessagesBuffer)

	go func() {
		defer close(chout)

		for {
			messageType, data, err := cw.conn.ReadMessage()
			if err != nil {
				// TODO: Handle error?
				return
			}

			if websocket.TextMessage != messageType {
				// TODO: Handle case - unexpected message type?
				continue
			}

			chout <- data
		}
	}()

	return chout
}

func (cw *ConnectionWorker) decode(chin <-chan []byte, stop <-chan struct{}) <-chan InputMessage {
	chout := make(chan InputMessage, chanDecodeMessageBuffer)

	go func() {
		defer close(chout)

		var decoder = ffjson.NewDecoder()

		for {
			select {
			case data := <-chin:
				var inputMessage *InputMessage
				if err := decoder.Decode(data, &inputMessage); err != nil {
					// TODO: Handler error.
				} else {
					chout <- *inputMessage
				}
			case <-stop:
				return
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) broadcastInputMessage(chin <-chan InputMessage, stop <-chan struct{}) {
	go func() {
		for {
			select {
			case inputMessage := <-chin:
				cw.chsInputMux.RLock()
				for _, ch := range cw.chsInput {
					select {
					case ch <- inputMessage:
					case <-stop:
						return
					}
				}
				cw.chsInputMux.RUnlock()
			case <-stop:
				return
			}
		}
	}()
}

func (cw *ConnectionWorker) Input(stop <-chan struct{}) <-chan InputMessage {
	chProxy := make(chan InputMessage, chanProxyInputMessageBuffer)

	cw.chsInputMux.Lock()
	cw.chsInput = append(cw.chsInput, chProxy)
	cw.chsInputMux.Unlock()

	chout := make(chan InputMessage, chanListenInputMessageBuffer)

	go func() {
		defer close(chout)
		defer func() {
			cw.chsInputMux.Lock()
			for i := range cw.chsInput {
				if cw.chsInput[i] == chProxy {
					cw.chsInput = append(cw.chsInput[:i], cw.chsInput[i+1:]...)
					close(chProxy)
					break
				}
			}
			cw.chsInputMux.Unlock()
		}()

		for {
			select {
			case <-stop:
				return
			case <-cw.chStop:
				return
			case inputMessage := <-chProxy:
				cw.sendInputMessage(chout, inputMessage, stop, sendInputMessageTimeout)
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) sendInputMessage(ch chan InputMessage, inputMessage InputMessage, stop <-chan struct{}, timeout time.Duration) {
	var timer = time.NewTimer(timeout)
	defer timer.Stop()
	if cap(ch) == 0 {
		select {
		case ch <- inputMessage:
		case <-cw.chStop:
		case <-stop:
		case <-timer.C:
		}
	} else {
		for {
			select {
			case ch <- inputMessage:
				return
			case <-cw.chStop:
				return
			case <-stop:
				return
			case <-timer.C:
				return
			default:
				if len(ch) == cap(ch) {
					<-ch
				}
			}
		}
	}
}

func (cw *ConnectionWorker) write(chin <-chan []byte, stop <-chan struct{}) {
	go func() {
		for {
			select {
			case data := <-chin:
				if err := cw.conn.WriteMessage(websocket.TextMessage, data); err != nil {
					// TODO: Handler error.
				}
			case <-stop:
				return
			}
		}
	}()
}

func (cw *ConnectionWorker) outputMessage(outputMessage OutputMessage) {
	select {
	case cw.chOutput <- outputMessage:
	case <-cw.chStop:
	}
}

func (cw *ConnectionWorker) encode(chin <-chan OutputMessage, stop <-chan struct{}) <-chan []byte {
	chout := make(chan []byte, chanEncodeMessageBuffer)

	go func() {
		defer close(chout)

		for {
			select {
			case message := <-chin:
				if data, err := ffjson.Marshal(message); err != nil {
					// TODO: Handler error.
				} else {
					chout <- data
				}
			case <-stop:
				return
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) listen(chin <-chan game.Event, stop <-chan struct{}) {
	for {
		select {
		case event := <-chin:
			// TODO: Do stuff.
			cw.logger.Info(event)

			// ...

			cw.outputMessage(OutputMessage{
				Type:    OutputMessageTypeGameEvent,
				Payload: event,
			})
		case <-stop:
			return
		}
	}
}
<commit_msg>Refactor connection worker goroutines<commit_after>package connections

import (
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/pquerna/ffjson/ffjson"
	"github.com/sirupsen/logrus"

	"github.com/ivan1993spb/snake-server/game"
)

const (
	chanOutputMessageBuffer      = 128
	chanReadMessagesBuffer       = 128
	chanDecodeMessageBuffer      = 128
	chanEncodeMessageBuffer      = 128
	chanProxyInputMessageBuffer  = 64
	chanListenInputMessageBuffer = 32

	sendInputMessageTimeout = time.Millisecond * 50

	sendCloseConnectionTimeout = time.Second

	readMessageLimit = 1024
)

type ConnectionWorker struct {
	conn   *websocket.Conn
	logger *logrus.Logger

	chStop      <-chan struct{}
	chsInput    []chan InputMessage
	chsInputMux *sync.RWMutex

	flagStarted bool
}

func NewConnectionWorker(conn *websocket.Conn, logger *logrus.Logger) *ConnectionWorker {
	return &ConnectionWorker{
		conn:        conn,
		logger:      logger,
		chsInput:    make([]chan InputMessage, 0),
		chsInputMux: &sync.RWMutex{},
	}
}

type ErrStartConnectionWorker string

func (e ErrStartConnectionWorker) Error() string {
	return "error start connection worker: " + string(e)
}

func (cw *ConnectionWorker) Start(game *game.Game) error {
	if cw.flagStarted {
		return ErrStartConnectionWorker("connection worker already started")
	}

	cw.flagStarted = true

	cw.conn.SetCloseHandler(cw.handleCloseConnection)
	cw.conn.SetReadLimit(readMessageLimit)

	// Input
	chInputBytes, chStop := cw.read()
	chInputMessages := cw.decode(chInputBytes, chStop)
	cw.broadcastInputMessage(chInputMessages, chStop)

	// Output
	chOutputMessages := cw.listenGameEvents(game.Events(chStop), chStop)
	chOutputBytes := cw.encode(chOutputMessages, chStop)
	cw.write(chOutputBytes, chStop)

	cw.chStop = chStop

	<-chStop

	cw.stopInputs()

	return nil
}

func (cw *ConnectionWorker) handleCloseConnection(code int, text string) error {
	message := websocket.FormatCloseMessage(code, "")
	cw.conn.WriteControl(websocket.CloseMessage, message, time.Now().Add(sendCloseConnectionTimeout))
	return nil
}

func (cw *ConnectionWorker) stopInputs() {
	cw.chsInputMux.Lock()
	defer cw.chsInputMux.Unlock()

	for _, ch := range cw.chsInput {
		close(ch)
	}

	cw.chsInput = cw.chsInput[:0]
}

func (cw *ConnectionWorker) read() (<-chan []byte, <-chan struct{}) {
	chout := make(chan []byte, chanReadMessagesBuffer)
	chstop := make(chan struct{}, 0)

	go func() {
		defer close(chout)
		defer close(chstop)

		for {
			messageType, data, err := cw.conn.ReadMessage()
			if err != nil {
				// TODO: Handle error?
				return
			}

			if websocket.TextMessage != messageType {
				// TODO: Handle case - unexpected message type?
				continue
			}

			chout <- data
		}
	}()

	return chout, chstop
}

func (cw *ConnectionWorker) decode(chin <-chan []byte, stop <-chan struct{}) <-chan InputMessage {
	chout := make(chan InputMessage, chanDecodeMessageBuffer)

	go func() {
		defer close(chout)

		var decoder = ffjson.NewDecoder()

		for {
			select {
			case data := <-chin:
				var inputMessage *InputMessage
				if err := decoder.Decode(data, &inputMessage); err != nil {
					// TODO: Handler error.
				} else {
					chout <- *inputMessage
				}
			case <-stop:
				return
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) broadcastInputMessage(chin <-chan InputMessage, stop <-chan struct{}) {
	go func() {
		for {
			select {
			case inputMessage := <-chin:
				cw.chsInputMux.RLock()
				for _, ch := range cw.chsInput {
					select {
					case ch <- inputMessage:
					case <-stop:
						return
					}
				}
				cw.chsInputMux.RUnlock()
			case <-stop:
				return
			}
		}
	}()
}

func (cw *ConnectionWorker) Input(stop <-chan struct{}) <-chan InputMessage {
	chProxy := make(chan InputMessage, chanProxyInputMessageBuffer)

	cw.chsInputMux.Lock()
	cw.chsInput = append(cw.chsInput, chProxy)
	cw.chsInputMux.Unlock()

	chout := make(chan InputMessage, chanListenInputMessageBuffer)

	go func() {
		defer close(chout)
		defer func() {
			cw.chsInputMux.Lock()
			for i := range cw.chsInput {
				if cw.chsInput[i] == chProxy {
					cw.chsInput = append(cw.chsInput[:i], cw.chsInput[i+1:]...)
					close(chProxy)
					break
				}
			}
			cw.chsInputMux.Unlock()
		}()

		for {
			select {
			case <-stop:
				return
			case <-cw.chStop:
				return
			case inputMessage := <-chProxy:
				cw.sendInputMessage(chout, inputMessage, stop, sendInputMessageTimeout)
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) sendInputMessage(ch chan InputMessage, inputMessage InputMessage, stop <-chan struct{}, timeout time.Duration) {
	var timer = time.NewTimer(timeout)
	defer timer.Stop()
	if cap(ch) == 0 {
		select {
		case ch <- inputMessage:
		case <-cw.chStop:
		case <-stop:
		case <-timer.C:
		}
	} else {
		for {
			select {
			case ch <- inputMessage:
				return
			case <-cw.chStop:
				return
			case <-stop:
				return
			case <-timer.C:
				return
			default:
				if len(ch) == cap(ch) {
					<-ch
				}
			}
		}
	}
}

func (cw *ConnectionWorker) write(chin <-chan []byte, stop <-chan struct{}) {
	go func() {
		for {
			select {
			case data := <-chin:
				if err := cw.conn.WriteMessage(websocket.TextMessage, data); err != nil {
					// TODO: Handler error.
				}
			case <-stop:
				return
			}
		}
	}()
}

func (cw *ConnectionWorker) encode(chin <-chan OutputMessage, stop <-chan struct{}) <-chan []byte {
	chout := make(chan []byte, chanEncodeMessageBuffer)

	go func() {
		defer close(chout)

		for {
			select {
			case message := <-chin:
				if data, err := ffjson.Marshal(message); err != nil {
					// TODO: Handler error.
				} else {
					chout <- data
				}
			case <-stop:
				return
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) listenGameEvents(chin <-chan game.Event, stop <-chan struct{}) <-chan OutputMessage {
	chout := make(chan OutputMessage, chanOutputMessageBuffer)

	go func() {
		defer close(chout)

		for {
			select {
			case event := <-chin:
				// TODO: Do stuff.

				outputMessage := OutputMessage{
					Type:    OutputMessageTypeGameEvent,
					Payload: event,
				}

				select {
				case chout <- outputMessage:
				case <-stop:
					return
				}
			case <-stop:
				return
			}
		}
	}()

	return chout
}
<|endoftext|>


<commit_before>package consul

import (
	"bytes"
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/testutil"
	"github.com/hashicorp/net-rpc-msgpackrpc"
)

// verifySnapshot is a helper that does a snapshot and restore.
func verifySnapshot(t *testing.T, s *Server, dc, token string) {
	codec := rpcClient(t, s)
	defer codec.Close()

	// Set a key to a before value.
	{
		args := structs.KVSRequest{
			Datacenter: dc,
			Op:         structs.KVSSet,
			DirEnt: structs.DirEntry{
				Key:   "test",
				Value: []byte("hello"),
			},
			WriteRequest: structs.WriteRequest{
				Token: token,
			},
		}
		var out bool
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Apply", &args, &out); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Take a snapshot.
	args := structs.SnapshotRequest{
		Datacenter: dc,
		Token:      token,
		Op:         structs.SnapshotSave,
	}
	var reply structs.SnapshotResponse
	snap, err := SnapshotRPC(s.connPool, s.config.Datacenter, s.config.RPCAddr,
		&args, bytes.NewReader([]byte("")), &reply)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer snap.Close()

	// Read back the before value.
	{
		getR := structs.KeyRequest{
			Datacenter: dc,
			Key:        "test",
			QueryOptions: structs.QueryOptions{
				Token: token,
			},
		}
		var dirent structs.IndexedDirEntries
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Get", &getR, &dirent); err != nil {
			t.Fatalf("err: %v", err)
		}
		if len(dirent.Entries) != 1 {
			t.Fatalf("Bad: %v", dirent)
		}
		d := dirent.Entries[0]
		if string(d.Value) != "hello" {
			t.Fatalf("bad: %v", d)
		}
	}

	// Set a key to an after value.
	{
		args := structs.KVSRequest{
			Datacenter: dc,
			Op:         structs.KVSSet,
			DirEnt: structs.DirEntry{
				Key:   "test",
				Value: []byte("goodbye"),
			},
			WriteRequest: structs.WriteRequest{
				Token: token,
			},
		}
		var out bool
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Apply", &args, &out); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Read back the before value.
	{
		getR := structs.KeyRequest{
			Datacenter: dc,
			Key:        "test",
			QueryOptions: structs.QueryOptions{
				Token: token,
			},
		}
		var dirent structs.IndexedDirEntries
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Get", &getR, &dirent); err != nil {
			t.Fatalf("err: %v", err)
		}
		if len(dirent.Entries) != 1 {
			t.Fatalf("Bad: %v", dirent)
		}
		d := dirent.Entries[0]
		if string(d.Value) != "goodbye" {
			t.Fatalf("bad: %v", d)
		}
	}

	// Restore the snapshot.
	args.Op = structs.SnapshotRestore
	restore, err := SnapshotRPC(s.connPool, s.config.Datacenter, s.config.RPCAddr,
		&args, snap, &reply)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer restore.Close()

	// Read back the before value post-snapshot.
	{
		getR := structs.KeyRequest{
			Datacenter: dc,
			Key:        "test",
			QueryOptions: structs.QueryOptions{
				Token: token,
			},
		}
		var dirent structs.IndexedDirEntries
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Get", &getR, &dirent); err != nil {
			t.Fatalf("err: %v", err)
		}
		if len(dirent.Entries) != 1 {
			t.Fatalf("Bad: %v", dirent)
		}
		d := dirent.Entries[0]
		if string(d.Value) != "hello" {
			t.Fatalf("bad: %v", d)
		}
	}
}

func TestSnapshot(t *testing.T) {
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testutil.WaitForLeader(t, s1.RPC, "dc1")
	verifySnapshot(t, s1, "dc1", "")
}

func TestSnapshot_LeaderState(t *testing.T) {
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testutil.WaitForLeader(t, s1.RPC, "dc1")

	codec := rpcClient(t, s1)
	defer codec.Close()

	// Make a before session.
	var before string
	{
		args := structs.SessionRequest{
			Datacenter: s1.config.Datacenter,
			Op:         structs.SessionCreate,
			Session: structs.Session{
				Node: s1.config.NodeName,
				TTL:  "60s",
			},
		}
		if err := msgpackrpc.CallWithCodec(codec, "Session.Apply", &args, &before); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Take a snapshot.
	args := structs.SnapshotRequest{
		Datacenter: s1.config.Datacenter,
		Op:         structs.SnapshotSave,
	}
	var reply structs.SnapshotResponse
	snap, err := SnapshotRPC(s1.connPool, s1.config.Datacenter, s1.config.RPCAddr,
		&args, bytes.NewReader([]byte("")), &reply)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer snap.Close()

	// Make an after session.
	var after string
	{
		args := structs.SessionRequest{
			Datacenter: s1.config.Datacenter,
			Op:         structs.SessionCreate,
			Session: structs.Session{
				Node: s1.config.NodeName,
				TTL:  "60s",
			},
		}
		if err := msgpackrpc.CallWithCodec(codec, "Session.Apply", &args, &after); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Make sure the leader has timers setup.
	if _, ok := s1.sessionTimers[before]; !ok {
		t.Fatalf("missing session timer")
	}
	if _, ok := s1.sessionTimers[after]; !ok {
		t.Fatalf("missing session timer")
	}

	// Restore the snapshot.
	args.Op = structs.SnapshotRestore
	restore, err := SnapshotRPC(s1.connPool, s1.config.Datacenter, s1.config.RPCAddr,
		&args, snap, &reply)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer restore.Close()

	// Make sure the before time is still there, and that the after timer
	// got reverted. This proves we fully cycled the leader state.
	if _, ok := s1.sessionTimers[before]; !ok {
		t.Fatalf("missing session timer")
	}
	if _, ok := s1.sessionTimers[after]; ok {
		t.Fatalf("unexpected session timer")
	}
}

func TestSnapshot_ACLDeny(t *testing.T) {
	dir1, s1 := testServerWithConfig(t, func(c *Config) {
		c.ACLDatacenter = "dc1"
		c.ACLMasterToken = "root"
		c.ACLDefaultPolicy = "deny"
	})
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()
	codec := rpcClient(t, s1)
	defer codec.Close()

	testutil.WaitForLeader(t, s1.RPC, "dc1")

	// Take a snapshot.
	func() {
		args := structs.SnapshotRequest{
			Datacenter: "dc1",
			Op:         structs.SnapshotSave,
		}
		var reply structs.SnapshotResponse
		_, err := SnapshotRPC(s1.connPool, s1.config.Datacenter, s1.config.RPCAddr,
			&args, bytes.NewReader([]byte("")), &reply)
		if err == nil || !strings.Contains(err.Error(), permissionDenied) {
			t.Fatalf("err: %v", err)
		}
	}()

	// Restore a snapshot.
	func() {
		args := structs.SnapshotRequest{
			Datacenter: "dc1",
			Op:         structs.SnapshotRestore,
		}
		var reply structs.SnapshotResponse
		_, err := SnapshotRPC(s1.connPool, s1.config.Datacenter, s1.config.RPCAddr,
			&args, bytes.NewReader([]byte("")), &reply)
		if err == nil || !strings.Contains(err.Error(), permissionDenied) {
			t.Fatalf("err: %v", err)
		}
	}()

	// With the token in place everything should go through.
	verifySnapshot(t, s1, "dc1", "root")
}

func TestSnapshot_Forward_Leader(t *testing.T) {
	dir1, s1 := testServerWithConfig(t, func(c *Config) {
		c.Bootstrap = true
	})
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	dir2, s2 := testServerWithConfig(t, func(c *Config) {
		c.Bootstrap = false
	})
	defer os.RemoveAll(dir2)
	defer s2.Shutdown()

	// Try to join.
	addr := fmt.Sprintf("127.0.0.1:%d",
		s1.config.SerfLANConfig.MemberlistConfig.BindPort)
	if _, err := s2.JoinLAN([]string{addr}); err != nil {
		t.Fatalf("err: %v", err)
	}

	testutil.WaitForLeader(t, s1.RPC, "dc1")
	testutil.WaitForLeader(t, s2.RPC, "dc1")

	// Run against the leader and the follower to ensure we forward.
	for _, s := range []*Server{s1, s2} {
		verifySnapshot(t, s, "dc1", "")
		verifySnapshot(t, s, "dc1", "")
	}
}

func TestSnapshot_Forward_Datacenter(t *testing.T) {
	dir1, s1 := testServerDC(t, "dc1")
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	dir2, s2 := testServerDC(t, "dc2")
	defer os.RemoveAll(dir2)
	defer s2.Shutdown()

	testutil.WaitForLeader(t, s1.RPC, "dc1")
	testutil.WaitForLeader(t, s2.RPC, "dc2")

	// Try to WAN join.
	addr := fmt.Sprintf("127.0.0.1:%d",
		s1.config.SerfWANConfig.MemberlistConfig.BindPort)
	if _, err := s2.JoinWAN([]string{addr}); err != nil {
		t.Fatalf("err: %v", err)
	}
	testutil.WaitForResult(
		func() (bool, error) {
			return len(s1.WANMembers()) > 1, nil
		},
		func(err error) {
			t.Fatalf("Failed waiting for WAN join: %v", err)
		})

	// Run a snapshot from each server locally and remotely to ensure we
	// forward.
	for _, s := range []*Server{s1, s2} {
		verifySnapshot(t, s, "dc1", "")
		verifySnapshot(t, s, "dc2", "")
	}
}

func TestSnapshot_AllowStale(t *testing.T) {
	dir1, s1 := testServerWithConfig(t, func(c *Config) {
		c.Bootstrap = false
	})
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	dir2, s2 := testServerWithConfig(t, func(c *Config) {
		c.Bootstrap = false
	})
	defer os.RemoveAll(dir2)
	defer s2.Shutdown()

	// Run against the servers which aren't haven't been set up to establish
	// a leader and make sure we get a no leader error.
	for _, s := range []*Server{s1, s2} {
		// Take a snapshot.
		args := structs.SnapshotRequest{
			Datacenter: s.config.Datacenter,
			Op:         structs.SnapshotSave,
		}
		var reply structs.SnapshotResponse
		_, err := SnapshotRPC(s.connPool, s.config.Datacenter, s.config.RPCAddr,
			&args, bytes.NewReader([]byte("")), &reply)
		if err == nil || !strings.Contains(err.Error(), structs.ErrNoLeader.Error()) {
			t.Fatalf("err: %v", err)
		}
	}

	// Run in stale mode and make sure we get an error from Raft (snapshot
	// was attempted), and not a no leader error.
	for _, s := range []*Server{s1, s2} {
		// Take a snapshot.
		args := structs.SnapshotRequest{
			Datacenter: s.config.Datacenter,
			AllowStale: true,
			Op:         structs.SnapshotSave,
		}
		var reply structs.SnapshotResponse
		_, err := SnapshotRPC(s.connPool, s.config.Datacenter, s.config.RPCAddr,
			&args, bytes.NewReader([]byte("")), &reply)
		if err == nil || !strings.Contains(err.Error(), "Raft error when taking snapshot") {
			t.Fatalf("err: %v", err)
		}
	}
}
<commit_msg>Fixes snapshot endpoint after merge.<commit_after>package consul

import (
	"bytes"
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/testutil"
	"github.com/hashicorp/net-rpc-msgpackrpc"
)

// verifySnapshot is a helper that does a snapshot and restore.
func verifySnapshot(t *testing.T, s *Server, dc, token string) {
	codec := rpcClient(t, s)
	defer codec.Close()

	// Set a key to a before value.
	{
		args := structs.KVSRequest{
			Datacenter: dc,
			Op:         structs.KVSSet,
			DirEnt: structs.DirEntry{
				Key:   "test",
				Value: []byte("hello"),
			},
			WriteRequest: structs.WriteRequest{
				Token: token,
			},
		}
		var out bool
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Apply", &args, &out); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Take a snapshot.
	args := structs.SnapshotRequest{
		Datacenter: dc,
		Token:      token,
		Op:         structs.SnapshotSave,
	}
	var reply structs.SnapshotResponse
	snap, err := SnapshotRPC(s.connPool, s.config.Datacenter, s.config.RPCAddr.String(),
		&args, bytes.NewReader([]byte("")), &reply)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer snap.Close()

	// Read back the before value.
	{
		getR := structs.KeyRequest{
			Datacenter: dc,
			Key:        "test",
			QueryOptions: structs.QueryOptions{
				Token: token,
			},
		}
		var dirent structs.IndexedDirEntries
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Get", &getR, &dirent); err != nil {
			t.Fatalf("err: %v", err)
		}
		if len(dirent.Entries) != 1 {
			t.Fatalf("Bad: %v", dirent)
		}
		d := dirent.Entries[0]
		if string(d.Value) != "hello" {
			t.Fatalf("bad: %v", d)
		}
	}

	// Set a key to an after value.
	{
		args := structs.KVSRequest{
			Datacenter: dc,
			Op:         structs.KVSSet,
			DirEnt: structs.DirEntry{
				Key:   "test",
				Value: []byte("goodbye"),
			},
			WriteRequest: structs.WriteRequest{
				Token: token,
			},
		}
		var out bool
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Apply", &args, &out); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Read back the before value.
	{
		getR := structs.KeyRequest{
			Datacenter: dc,
			Key:        "test",
			QueryOptions: structs.QueryOptions{
				Token: token,
			},
		}
		var dirent structs.IndexedDirEntries
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Get", &getR, &dirent); err != nil {
			t.Fatalf("err: %v", err)
		}
		if len(dirent.Entries) != 1 {
			t.Fatalf("Bad: %v", dirent)
		}
		d := dirent.Entries[0]
		if string(d.Value) != "goodbye" {
			t.Fatalf("bad: %v", d)
		}
	}

	// Restore the snapshot.
	args.Op = structs.SnapshotRestore
	restore, err := SnapshotRPC(s.connPool, s.config.Datacenter, s.config.RPCAddr.String(),
		&args, snap, &reply)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer restore.Close()

	// Read back the before value post-snapshot.
	{
		getR := structs.KeyRequest{
			Datacenter: dc,
			Key:        "test",
			QueryOptions: structs.QueryOptions{
				Token: token,
			},
		}
		var dirent structs.IndexedDirEntries
		if err := msgpackrpc.CallWithCodec(codec, "KVS.Get", &getR, &dirent); err != nil {
			t.Fatalf("err: %v", err)
		}
		if len(dirent.Entries) != 1 {
			t.Fatalf("Bad: %v", dirent)
		}
		d := dirent.Entries[0]
		if string(d.Value) != "hello" {
			t.Fatalf("bad: %v", d)
		}
	}
}

func TestSnapshot(t *testing.T) {
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testutil.WaitForLeader(t, s1.RPC, "dc1")
	verifySnapshot(t, s1, "dc1", "")
}

func TestSnapshot_LeaderState(t *testing.T) {
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testutil.WaitForLeader(t, s1.RPC, "dc1")

	codec := rpcClient(t, s1)
	defer codec.Close()

	// Make a before session.
	var before string
	{
		args := structs.SessionRequest{
			Datacenter: s1.config.Datacenter,
			Op:         structs.SessionCreate,
			Session: structs.Session{
				Node: s1.config.NodeName,
				TTL:  "60s",
			},
		}
		if err := msgpackrpc.CallWithCodec(codec, "Session.Apply", &args, &before); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Take a snapshot.
	args := structs.SnapshotRequest{
		Datacenter: s1.config.Datacenter,
		Op:         structs.SnapshotSave,
	}
	var reply structs.SnapshotResponse
	snap, err := SnapshotRPC(s1.connPool, s1.config.Datacenter, s1.config.RPCAddr.String(),
		&args, bytes.NewReader([]byte("")), &reply)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer snap.Close()

	// Make an after session.
	var after string
	{
		args := structs.SessionRequest{
			Datacenter: s1.config.Datacenter,
			Op:         structs.SessionCreate,
			Session: structs.Session{
				Node: s1.config.NodeName,
				TTL:  "60s",
			},
		}
		if err := msgpackrpc.CallWithCodec(codec, "Session.Apply", &args, &after); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Make sure the leader has timers setup.
	if _, ok := s1.sessionTimers[before]; !ok {
		t.Fatalf("missing session timer")
	}
	if _, ok := s1.sessionTimers[after]; !ok {
		t.Fatalf("missing session timer")
	}

	// Restore the snapshot.
	args.Op = structs.SnapshotRestore
	restore, err := SnapshotRPC(s1.connPool, s1.config.Datacenter, s1.config.RPCAddr.String(),
		&args, snap, &reply)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer restore.Close()

	// Make sure the before time is still there, and that the after timer
	// got reverted. This proves we fully cycled the leader state.
	if _, ok := s1.sessionTimers[before]; !ok {
		t.Fatalf("missing session timer")
	}
	if _, ok := s1.sessionTimers[after]; ok {
		t.Fatalf("unexpected session timer")
	}
}

func TestSnapshot_ACLDeny(t *testing.T) {
	dir1, s1 := testServerWithConfig(t, func(c *Config) {
		c.ACLDatacenter = "dc1"
		c.ACLMasterToken = "root"
		c.ACLDefaultPolicy = "deny"
	})
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()
	codec := rpcClient(t, s1)
	defer codec.Close()

	testutil.WaitForLeader(t, s1.RPC, "dc1")

	// Take a snapshot.
	func() {
		args := structs.SnapshotRequest{
			Datacenter: "dc1",
			Op:         structs.SnapshotSave,
		}
		var reply structs.SnapshotResponse
		_, err := SnapshotRPC(s1.connPool, s1.config.Datacenter, s1.config.RPCAddr.String(),
			&args, bytes.NewReader([]byte("")), &reply)
		if err == nil || !strings.Contains(err.Error(), permissionDenied) {
			t.Fatalf("err: %v", err)
		}
	}()

	// Restore a snapshot.
	func() {
		args := structs.SnapshotRequest{
			Datacenter: "dc1",
			Op:         structs.SnapshotRestore,
		}
		var reply structs.SnapshotResponse
		_, err := SnapshotRPC(s1.connPool, s1.config.Datacenter, s1.config.RPCAddr.String(),
			&args, bytes.NewReader([]byte("")), &reply)
		if err == nil || !strings.Contains(err.Error(), permissionDenied) {
			t.Fatalf("err: %v", err)
		}
	}()

	// With the token in place everything should go through.
	verifySnapshot(t, s1, "dc1", "root")
}

func TestSnapshot_Forward_Leader(t *testing.T) {
	dir1, s1 := testServerWithConfig(t, func(c *Config) {
		c.Bootstrap = true
	})
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	dir2, s2 := testServerWithConfig(t, func(c *Config) {
		c.Bootstrap = false
	})
	defer os.RemoveAll(dir2)
	defer s2.Shutdown()

	// Try to join.
	addr := fmt.Sprintf("127.0.0.1:%d",
		s1.config.SerfLANConfig.MemberlistConfig.BindPort)
	if _, err := s2.JoinLAN([]string{addr}); err != nil {
		t.Fatalf("err: %v", err)
	}

	testutil.WaitForLeader(t, s1.RPC, "dc1")
	testutil.WaitForLeader(t, s2.RPC, "dc1")

	// Run against the leader and the follower to ensure we forward.
	for _, s := range []*Server{s1, s2} {
		verifySnapshot(t, s, "dc1", "")
		verifySnapshot(t, s, "dc1", "")
	}
}

func TestSnapshot_Forward_Datacenter(t *testing.T) {
	dir1, s1 := testServerDC(t, "dc1")
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	dir2, s2 := testServerDC(t, "dc2")
	defer os.RemoveAll(dir2)
	defer s2.Shutdown()

	testutil.WaitForLeader(t, s1.RPC, "dc1")
	testutil.WaitForLeader(t, s2.RPC, "dc2")

	// Try to WAN join.
	addr := fmt.Sprintf("127.0.0.1:%d",
		s1.config.SerfWANConfig.MemberlistConfig.BindPort)
	if _, err := s2.JoinWAN([]string{addr}); err != nil {
		t.Fatalf("err: %v", err)
	}
	testutil.WaitForResult(
		func() (bool, error) {
			return len(s1.WANMembers()) > 1, nil
		},
		func(err error) {
			t.Fatalf("Failed waiting for WAN join: %v", err)
		})

	// Run a snapshot from each server locally and remotely to ensure we
	// forward.
	for _, s := range []*Server{s1, s2} {
		verifySnapshot(t, s, "dc1", "")
		verifySnapshot(t, s, "dc2", "")
	}
}

func TestSnapshot_AllowStale(t *testing.T) {
	dir1, s1 := testServerWithConfig(t, func(c *Config) {
		c.Bootstrap = false
	})
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	dir2, s2 := testServerWithConfig(t, func(c *Config) {
		c.Bootstrap = false
	})
	defer os.RemoveAll(dir2)
	defer s2.Shutdown()

	// Run against the servers which aren't haven't been set up to establish
	// a leader and make sure we get a no leader error.
	for _, s := range []*Server{s1, s2} {
		// Take a snapshot.
		args := structs.SnapshotRequest{
			Datacenter: s.config.Datacenter,
			Op:         structs.SnapshotSave,
		}
		var reply structs.SnapshotResponse
		_, err := SnapshotRPC(s.connPool, s.config.Datacenter, s.config.RPCAddr.String(),
			&args, bytes.NewReader([]byte("")), &reply)
		if err == nil || !strings.Contains(err.Error(), structs.ErrNoLeader.Error()) {
			t.Fatalf("err: %v", err)
		}
	}

	// Run in stale mode and make sure we get an error from Raft (snapshot
	// was attempted), and not a no leader error.
	for _, s := range []*Server{s1, s2} {
		// Take a snapshot.
		args := structs.SnapshotRequest{
			Datacenter: s.config.Datacenter,
			AllowStale: true,
			Op:         structs.SnapshotSave,
		}
		var reply structs.SnapshotResponse
		_, err := SnapshotRPC(s.connPool, s.config.Datacenter, s.config.RPCAddr.String(),
			&args, bytes.NewReader([]byte("")), &reply)
		if err == nil || !strings.Contains(err.Error(), "Raft error when taking snapshot") {
			t.Fatalf("err: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package zhihu

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strconv"
	"strings"

	"github.com/PuerkitoBio/goquery"
)

// Collection 是一个知乎的收藏夹页面
type Collection struct {
	*Page

	// creator 是该收藏夹的创建者
	creator *User

	// name 是该收藏夹的名称
	name string
}

// NewCollection 创建一个收藏夹对象，返回 *Collection
func NewCollection(link string, name string, creator *User) *Collection {
	if !validCollectionURL(link) {
		panic("收藏夹链接不正确：" + link)
	}

	return &Collection{
		Page:    newZhihuPage(link),
		creator: creator,
		name:    name,
	}
}

// GetName 返回收藏夹的名字
func (c *Collection) GetName() string {
	if c.name != "" {
		return c.name
	}

	doc := c.Doc()

	//

	//   恩恩恩 大力一点，不要停～
	//

	c.name = strip(doc.Find("h2#zh-fav-head-title").Text())
	return c.name
}

// GetCreator 返回收藏夹的创建者
func (c *Collection) GetCreator() *User {
	if c.creator != nil {
		return c.creator
	}

	doc := c.Doc()

	//

	//   李阳良
	//

	sel := doc.Find("h2.zm-list-content-title a")
	userId := strip(sel.Text())
	linkPath, _ := sel.Attr("href")
	c.creator = NewUser(makeZhihuLink(linkPath), userId)
	return c.creator
}

// GetFollowersNum 返回收藏夹的关注者数量
func (c *Collection) GetFollowersNum() int {
	if got, ok := c.getIntField("followers-num"); ok {
		return got
	}

	doc := c.Doc()

	//
	//   7516
	//
	text := strip(doc.Find(`a[data-za-a="visit_collection_followers"]`).Text())
	num, _ := strconv.Atoi(text)
	c.setField("followers-num", num)
	return num
}

// GetFollowersN 返回 n 个关注该收藏夹的用户，如果 n < 0，返回所有关注者
func (c *Collection) GetFollowersN(n int) []*User {
	var (
		link = urlJoin(c.Link, "/followers")
		xsrf = c.GetXSRF()
	)
	users, err := ajaxGetFollowers(link, xsrf, n)
	if err != nil {
		return nil
	}
	return users
}

// GetFollowers 返回关注该收藏夹的用户
func (c *Collection) GetFollowers() []*User {
	return c.GetFollowersN(c.GetFollowersNum())
}

// GetQuestionsN 返回前 n 个问题，如果 n < 0，返回所有问题
func (c *Collection) GetQuestionsN(n int) []*Question {
	if n == 0 {
		return nil
	}

	// 先获取第一页的问题
	questions := getQuestionsFromDoc(c.Doc())

	totalPages := c.totalPages()
	if totalPages == 1 {
		if n < 0 || n > len(questions) {
			return questions
		}
		return questions[0:n]
	}

	// 再分页查询其他问题
	currentPage := 2
	for currentPage <= totalPages {
		link := fmt.Sprintf("%s?page=%d", c.Link, currentPage)
		doc, err := newDocumentFromURL(link)
		if err != nil {
			logger.Error("解析页面失败：%s, %s", link, err.Error())
			return nil
		}

		newQuestions := getQuestionsFromDoc(doc)
		questions = append(questions, newQuestions...)
		if n > 0 && len(questions) >= n {
			return questions[0:n]
		}
		currentPage++
	}

	return questions
}

// GetQuestions 返回收藏夹里所有的问题
func (c *Collection) GetQuestions() []*Question {
	return c.GetQuestionsN(-1)
}

// GetAnswersN 返回 n 个回答，如果 n < 0，返回所有回答
func (c *Collection) GetAnswersN(n int) []*Answer {
	if n == 0 {
		return nil
	}

	// 先获取第一页的回答
	answers := getAnswersFromDoc(c.Doc())

	totalPages := c.totalPages()
	if totalPages == 1 {
		if n < 0 || n > len(answers) {
			return answers
		}
		return answers[0:n]
	}

	// 在分页查询
	currentPage := 2
	for currentPage <= totalPages {
		link := fmt.Sprintf("%s?page=%d", c.Link, currentPage)
		doc, err := newDocumentFromURL(link)
		if err != nil {
			logger.Error("解析页面失败：%s, %s", link, err.Error())
			return nil
		}

		newAnswers := getAnswersFromDoc(doc)
		answers = append(answers, newAnswers...)
		if n > 0 && len(answers) >= n {
			return answers[0:n]
		}
		currentPage++
	}
	return answers
}

// GetAnswers 返回收藏夹里所有的回答
func (c *Collection) GetAnswers() []*Answer {
	return c.GetAnswersN(-1)
}

func (c *Collection) String() string {
	return fmt.Sprintf("<Collection: %s - %s>", c.GetName(), c.Link)
}

func (c *Collection) totalPages() int {
	pager := c.Doc().Find("div.zm-invite-pager")
	if pager.Size() == 0 {
		// 只有一页
		return 1
	}
	text := pager.Find("span").Eq(-2).Text()
	pages, _ := strconv.Atoi(text)
	return pages
}

func ajaxGetFollowers(link string, xsrf string, total int) ([]*User, error) {
	if total == 0 {
		return nil, nil
	}

	var (
		offset     = 0
		gotDataNum = pageSize
		initCap    = total
	)

	if initCap < 0 {
		initCap = pageSize
	}
	users := make([]*User, 0, initCap)

	form := url.Values{}
	form.Set("_xsrf", xsrf)

	for gotDataNum == pageSize {
		form.Set("offset", strconv.Itoa(offset))
		doc, dataNum, err := newDocByNormalAjax(link, form)
		if err != nil {
			return nil, err
		}

		doc.Find("div.zm-profile-card").Each(func(index int, sel *goquery.Selection) {
			thisUser := newUserFromSelector(sel)
			users = append(users, thisUser)
		})

		if total > 0 && len(users) >= total {
			return users[:total], nil
		}

		gotDataNum = dataNum
		offset += gotDataNum
	}
	return users, nil
}

func newDocByNormalAjax(link string, form url.Values) (*goquery.Document, int, error) {
	gotDataNum := 0
	body := strings.NewReader(form.Encode())
	resp, err := gSession.Ajax(link, body, link)
	if err != nil {
		logger.Error("查询关注的话题失败, 链接：%s, 参数：%s，错误：%s", link, form.Encode(), err.Error())
		return nil, gotDataNum, err
	}

	defer resp.Body.Close()
	result := normalAjaxResult{}
	err = json.NewDecoder(resp.Body).Decode(&result)
	if err != nil {
		logger.Error("解析返回值 json 失败：%s", err.Error())
		return nil, gotDataNum, err
	}

	topicsHtml := result.Msg[1].(string)
	doc, err := goquery.NewDocumentFromReader(strings.NewReader(topicsHtml))
	if err != nil {
		logger.Error("解析返回的 HTML 失败：%s", err.Error())
		return nil, gotDataNum, err
	}
	gotDataNum = int(result.Msg[0].(float64))
	return doc, gotDataNum, err
}

func getQuestionsFromDoc(doc *goquery.Document) []*Question {
	questions := make([]*Question, 0, pageSize)
	items := doc.Find("div#zh-list-answer-wrap").Find("h2.zm-item-title")
	items.Each(func(index int, sel *goquery.Selection) {
		a := sel.Find("a")
		qTitle := strip(a.Text())
		qHref, _ := a.Attr("href")
		thisQuestion := NewQuestion(makeZhihuLink(qHref), qTitle)
		questions = append(questions, thisQuestion)
	})
	return questions
}

func getAnswersFromDoc(doc *goquery.Document) []*Answer {
	var answers []*Answer
	var lastQuestion *Question

	doc.Find("div.zm-item").Each(func(index int, sel *goquery.Selection) {
		// 回答
		contentTag := sel.Find("div.zm-item-rich-text")
		if contentTag.Size() == 0 {
			// 回答被建议修改
			reason := strip(sel.Find("div.answer-status").Text())
			logger.Warn("忽略一个问题，原因：%s", reason)
			return
		}

		// 获取问题，如果同一个问题下收藏了多个回答，则除了第一个外，后面的回答的 HTML 部分，
		// 也就是 div.zm-item 里面不会有该问题的链接（a 标签），所以用 lastQuestion 标记
		// 最近的一个问题
		var thisQuestion *Question
		if qTag := sel.Find("h2.zm-item-title").Find("a"); qTag.Size() > 0 {
			qTitle := strip(qTag.Text())
			qHref, _ := qTag.Attr("href")
			thisQuestion = NewQuestion(makeZhihuLink(qHref), qTitle)
			lastQuestion = thisQuestion
		} else {
			thisQuestion = lastQuestion
		}

		// 答主
		author := newUserFromAnswerAuthorTag(sel.Find("div.zm-item-answer-author-info"))

		answerHref, _ := contentTag.Attr("data-entry-url")
		voteText, _ := sel.Find("a.zm-item-vote-count").Attr("data-votecount")
		vote, _ := strconv.Atoi(voteText)
		thisAnswer := NewAnswer(makeZhihuLink(answerHref), thisQuestion, author)
		thisAnswer.setUpvote(vote)

		answers = append(answers, thisAnswer)
	})
	return answers
}
<commit_msg>collection.go 里增加 GetQuestionsNum 和 GetAnswersNum 两个方法的签名<commit_after>package zhihu

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strconv"
	"strings"

	"github.com/PuerkitoBio/goquery"
)

// Collection 是一个知乎的收藏夹页面
type Collection struct {
	*Page

	// creator 是该收藏夹的创建者
	creator *User

	// name 是该收藏夹的名称
	name string
}

// NewCollection 创建一个收藏夹对象，返回 *Collection
func NewCollection(link string, name string, creator *User) *Collection {
	if !validCollectionURL(link) {
		panic("收藏夹链接不正确：" + link)
	}

	return &Collection{
		Page:    newZhihuPage(link),
		creator: creator,
		name:    name,
	}
}

// GetName 返回收藏夹的名字
func (c *Collection) GetName() string {
	if c.name != "" {
		return c.name
	}

	doc := c.Doc()

	//

	//   恩恩恩 大力一点，不要停～
	//

	c.name = strip(doc.Find("h2#zh-fav-head-title").Text())
	return c.name
}

// GetCreator 返回收藏夹的创建者
func (c *Collection) GetCreator() *User {
	if c.creator != nil {
		return c.creator
	}

	doc := c.Doc()

	//

	//   李阳良
	//

	sel := doc.Find("h2.zm-list-content-title a")
	userId := strip(sel.Text())
	linkPath, _ := sel.Attr("href")
	c.creator = NewUser(makeZhihuLink(linkPath), userId)
	return c.creator
}

// GetFollowersNum 返回收藏夹的关注者数量
func (c *Collection) GetFollowersNum() int {
	if got, ok := c.getIntField("followers-num"); ok {
		return got
	}

	doc := c.Doc()

	//
	//   7516
	//
	text := strip(doc.Find(`a[data-za-a="visit_collection_followers"]`).Text())
	num, _ := strconv.Atoi(text)
	c.setField("followers-num", num)
	return num
}

// GetFollowersN 返回 n 个关注该收藏夹的用户，如果 n < 0，返回所有关注者
func (c *Collection) GetFollowersN(n int) []*User {
	var (
		link = urlJoin(c.Link, "/followers")
		xsrf = c.GetXSRF()
	)
	users, err := ajaxGetFollowers(link, xsrf, n)
	if err != nil {
		return nil
	}
	return users
}

// GetFollowers 返回关注该收藏夹的用户
func (c *Collection) GetFollowers() []*User {
	return c.GetFollowersN(c.GetFollowersNum())
}

// GetQuestionsN 返回前 n 个问题，如果 n < 0，返回所有问题
func (c *Collection) GetQuestionsN(n int) []*Question {
	if n == 0 {
		return nil
	}

	// 先获取第一页的问题
	questions := getQuestionsFromDoc(c.Doc())

	totalPages := c.totalPages()
	if totalPages == 1 {
		if n < 0 || n > len(questions) {
			return questions
		}
		return questions[0:n]
	}

	// 再分页查询其他问题
	currentPage := 2
	for currentPage <= totalPages {
		link := fmt.Sprintf("%s?page=%d", c.Link, currentPage)
		doc, err := newDocumentFromURL(link)
		if err != nil {
			logger.Error("解析页面失败：%s, %s", link, err.Error())
			return nil
		}

		newQuestions := getQuestionsFromDoc(doc)
		questions = append(questions, newQuestions...)
		if n > 0 && len(questions) >= n {
			return questions[0:n]
		}
		currentPage++
	}

	return questions
}

// GetQuestions 返回收藏夹里所有的问题
func (c *Collection) GetQuestions() []*Question {
	return c.GetQuestionsN(-1)
}

// GetAnswersN 返回 n 个回答，如果 n < 0，返回所有回答
func (c *Collection) GetAnswersN(n int) []*Answer {
	if n == 0 {
		return nil
	}

	// 先获取第一页的回答
	answers := getAnswersFromDoc(c.Doc())

	totalPages := c.totalPages()
	if totalPages == 1 {
		if n < 0 || n > len(answers) {
			return answers
		}
		return answers[0:n]
	}

	// 在分页查询
	currentPage := 2
	for currentPage <= totalPages {
		link := fmt.Sprintf("%s?page=%d", c.Link, currentPage)
		doc, err := newDocumentFromURL(link)
		if err != nil {
			logger.Error("解析页面失败：%s, %s", link, err.Error())
			return nil
		}

		newAnswers := getAnswersFromDoc(doc)
		answers = append(answers, newAnswers...)
		if n > 0 && len(answers) >= n {
			return answers[0:n]
		}
		currentPage++
	}
	return answers
}

// GetAnswers 返回收藏夹里所有的回答
func (c *Collection) GetAnswers() []*Answer {
	return c.GetAnswersN(-1)
}

// TODO GetQuestionsNum 返回收藏夹的问题数量
func (c *Collection) GetQuestionsNum() int {
	return 0
}

// TODO GetAnswersNum 返回收藏夹的答案数量
func (c *Collection) GetAnswersNum() int {
	return 0
}

func (c *Collection) String() string {
	return fmt.Sprintf("<Collection: %s - %s>", c.GetName(), c.Link)
}

func (c *Collection) totalPages() int {
	pager := c.Doc().Find("div.zm-invite-pager")
	if pager.Size() == 0 {
		// 只有一页
		return 1
	}
	text := pager.Find("span").Eq(-2).Text()
	pages, _ := strconv.Atoi(text)
	return pages
}

func ajaxGetFollowers(link string, xsrf string, total int) ([]*User, error) {
	if total == 0 {
		return nil, nil
	}

	var (
		offset     = 0
		gotDataNum = pageSize
		initCap    = total
	)

	if initCap < 0 {
		initCap = pageSize
	}
	users := make([]*User, 0, initCap)

	form := url.Values{}
	form.Set("_xsrf", xsrf)

	for gotDataNum == pageSize {
		form.Set("offset", strconv.Itoa(offset))
		doc, dataNum, err := newDocByNormalAjax(link, form)
		if err != nil {
			return nil, err
		}

		doc.Find("div.zm-profile-card").Each(func(index int, sel *goquery.Selection) {
			thisUser := newUserFromSelector(sel)
			users = append(users, thisUser)
		})

		if total > 0 && len(users) >= total {
			return users[:total], nil
		}

		gotDataNum = dataNum
		offset += gotDataNum
	}
	return users, nil
}

func newDocByNormalAjax(link string, form url.Values) (*goquery.Document, int, error) {
	gotDataNum := 0
	body := strings.NewReader(form.Encode())
	resp, err := gSession.Ajax(link, body, link)
	if err != nil {
		logger.Error("查询关注的话题失败, 链接：%s, 参数：%s，错误：%s", link, form.Encode(), err.Error())
		return nil, gotDataNum, err
	}

	defer resp.Body.Close()
	result := normalAjaxResult{}
	err = json.NewDecoder(resp.Body).Decode(&result)
	if err != nil {
		logger.Error("解析返回值 json 失败：%s", err.Error())
		return nil, gotDataNum, err
	}

	topicsHtml := result.Msg[1].(string)
	doc, err := goquery.NewDocumentFromReader(strings.NewReader(topicsHtml))
	if err != nil {
		logger.Error("解析返回的 HTML 失败：%s", err.Error())
		return nil, gotDataNum, err
	}
	gotDataNum = int(result.Msg[0].(float64))
	return doc, gotDataNum, err
}

func getQuestionsFromDoc(doc *goquery.Document) []*Question {
	questions := make([]*Question, 0, pageSize)
	items := doc.Find("div#zh-list-answer-wrap").Find("h2.zm-item-title")
	items.Each(func(index int, sel *goquery.Selection) {
		a := sel.Find("a")
		qTitle := strip(a.Text())
		qHref, _ := a.Attr("href")
		thisQuestion := NewQuestion(makeZhihuLink(qHref), qTitle)
		questions = append(questions, thisQuestion)
	})
	return questions
}

func getAnswersFromDoc(doc *goquery.Document) []*Answer {
	var answers []*Answer
	var lastQuestion *Question

	doc.Find("div.zm-item").Each(func(index int, sel *goquery.Selection) {
		// 回答
		contentTag := sel.Find("div.zm-item-rich-text")
		if contentTag.Size() == 0 {
			// 回答被建议修改
			reason := strip(sel.Find("div.answer-status").Text())
			logger.Warn("忽略一个问题，原因：%s", reason)
			return
		}

		// 获取问题，如果同一个问题下收藏了多个回答，则除了第一个外，后面的回答的 HTML 部分，
		// 也就是 div.zm-item 里面不会有该问题的链接（a 标签），所以用 lastQuestion 标记
		// 最近的一个问题
		var thisQuestion *Question
		if qTag := sel.Find("h2.zm-item-title").Find("a"); qTag.Size() > 0 {
			qTitle := strip(qTag.Text())
			qHref, _ := qTag.Attr("href")
			thisQuestion = NewQuestion(makeZhihuLink(qHref), qTitle)
			lastQuestion = thisQuestion
		} else {
			thisQuestion = lastQuestion
		}

		// 答主
		author := newUserFromAnswerAuthorTag(sel.Find("div.zm-item-answer-author-info"))

		answerHref, _ := contentTag.Attr("data-entry-url")
		voteText, _ := sel.Find("a.zm-item-vote-count").Attr("data-votecount")
		vote, _ := strconv.Atoi(voteText)
		thisAnswer := NewAnswer(makeZhihuLink(answerHref), thisQuestion, author)
		thisAnswer.setUpvote(vote)

		answers = append(answers, thisAnswer)
	})
	return answers
}
<|endoftext|>