<commit_before>//Store is a package for managing persistent local server data,
//including data blobs, metadata, user data, and server configerations.
//
//
//Large files are divided into blocks for internal server deduplication, this is
//invisble to the outside. Lower block size gives better deduplication, and
//higher block size gives faster sequential access and less metadata. (use 4MB for now)
//
//Store is designed to store files with the following possibilities:
//
//Importing a local file: data is read, hashed, and saved a block at a time. The
//hash and length is reported at the end.
//
//Requested to be saved locally: hash and length are given (along with locations
//to the server). Inner hash are downloaded to see what blocks are needed, then
//the blocks are downloaded.
//
//As such hashes maybe stored bottom up or top down, and when a tree is traversed,
//it may find missing nodes. An error condition that the server must handle or report
//to client.
package store

import (
	"io"

	ht "github.com/xiegeo/fensan/hashtree"
)

//KV is an interface for a []byte based key value store. For storing small pieces
//of metadata that are often updated.
type KV interface {
	//Get gets the value for the given key. It returns nil if KV does not contain the key.
	//
	//Warning: Get might not see what was Set or Deleted until after sync().
	//This is a temperary work around to easly warp some apis.
	Get(key []byte) []byte
	//Set sets the value for the given key. It overwrites any previous value for that key.
	//Do not set with len(v) == 0, use Delete instead.
	Set(key []byte, v []byte)
	//Delete deletes the value for the given key, it is a no-op if KV does not contain the key.
	Delete(key []byte)
	//Sync commits all the changes to stable storage. It blocks until done.
	Sync()
	//Close closes KV.
	Close() error
	GC(startAfterKey []byte, f func(key []byte, v []byte) (delete bool, stop bool))
}

//Blob can be used to read/write parts of a long '[]byte' values
//without loading the full data.
//It can be backed by os.File, as it is a subset of os.File, except Size.
//All error reporting are turned off, as none are expected (panic otherwise).
type Blob interface {
	//length in bytes for this block
	Size() int64
	//ReadAt reads len(b) bytes from the File starting at byte offset off.
	ReadAt(b []byte, off int64)
	//WriteAt writes len(b) bytes to the File starting at byte offset off.
	//WriteAt, may or may not sync to disk, at any order.
	WriteAt(b []byte, off int64)
	//Sync commits the all WriteAt to stable storage.
	Sync()
	//Close closes the Block, allowing the freeing of resources. This Block can
	//not be used after Close.
	Close()
}

//LV (Large Values) is an interface for a Block based data store. For
//storing large values with length known in advance and as part of key.
//The same key with two sizes represent two different Blobs.
type LV interface {
	//New creates a new block in LV, or nil if the key already exist.
	//size > 0
	New(key []byte, size int64) Blob
	//Get returns the Block referenced by key, or nil if key is not found.
	Get(key []byte, size int64) Blob
	//Move changes the key from from to to, and resizes to newSize.
	//If from does not exist, or to already exist, then an error is reported.
	//The Blob keyed by from should not be open.
	Move(oldKey []byte, oldSize int64, newKey []byte, newSize int64) error

	//Delete removes the Block repersented by key from LV.
	//Returns true if blob is not in LV after Delete.
	Delete(key []byte, size int64) (bool, error)

	Close() error
}

type FileState int

const (
	//FILE_NONE means we don't have this file
	FILE_NONE FileState = iota
	//FILE_PART means we have parts of this file
	FILE_PART
	//FILE_COMPLETE means we have all of this file
	FILE_COMPLETE
)

//MetaStore is a store for hash meta data on top of KV.
//Such as: Child hashes, data length/hash level, and TTL.
type MetaStore interface {
	//InnerHashMinLevel reports the amount of inner hash saved, everything at or
	//above the Level can be retrieved
	InnerHashMinLevel() ht.Level
	//GetInnerHashes reads inner hashes at level and offset, len(hs) should be a
	//multiple of hash length (32).
	//
	//An error is reported if parts of hs is unknown, parts known is still filled
	//with unknown parts as all zeros.
	//
	//If parts of hs are impossible by index range, it panics
	GetInnerHashes(key HLKey, hs []byte, level ht.Level, off ht.Nodes) error

	PutInnerHashes(key HLKey, hs []byte, level ht.Level, off ht.Nodes) (has ht.Nodes, complete bool, err error)

	//TTLGet gets the TTL of a file
	TTLGet(key HLKey) TTL
	//TTLSetAtleast updates the TTL to be at least coved to util, the total
	//increase is multiplied by key.Length() to return costs in storage time by
	//byteMonth.
	//Savings in deduplication can be refected in byteMonth.
	TTLSetAtleast(key HLKey, util TTL) (byteMonth int64)

	//Close closes MetaStore
	Close()
}

//Database is the full interface for all the data of a server,
//WIP
type Database interface {
	MetaStore
	//GetState checks if we have a file or not, or in progress
	GetState(key HLKey) FileState
	//GetAt reads len(b) bytes of file key from offset off.
	//returns error is b can't be read completely
	GetAt(key HLKey, b []byte, off int64) error
	//PutAt writes b to file from offset off.
	//
	//has returns the number of leaf nodes completed. This can report the progress.
	//Iff has is full, complete is true
	//
	//If hash checking fails, an err is reported. The server can us this information
	//to demote the source.
	PutAt(key HLKey, b []byte, off int64) (has ht.Nodes, complete bool, err error)

	//Import a file from reader
	ImportFromReader(r io.Reader) HLKey
}

type metaValue struct {
	ttl TTL
}
<commit_msg>small changes<commit_after>//Store is a package for managing persistent local server data,
//including data blobs, metadata, user data, and server configerations.
//
//
//Large files are divided into blocks for internal server deduplication, this is
//invisble to the outside. Lower block size gives better deduplication, and
//higher block size gives faster sequential access and less metadata. (use 4MB for now)
//
//Store is designed to store files with the following possibilities:
//
//Importing a local file: data is read, hashed, and saved a block at a time. The
//hash and length is reported at the end.
//
//Requested to be saved locally: hash and length are given (along with locations
//to the server). Inner hash are downloaded to see what blocks are needed, then
//the blocks are downloaded.
//
//As such hashes maybe stored bottom up or top down, and when a tree is traversed,
//it may find missing nodes. An error condition that the server must handle or report
//to client.
package store

import (
	"io"

	ht "github.com/xiegeo/fensan/hashtree"
)

//KV is an interface for a []byte based key value store. For storing small pieces
//of metadata that are often updated.
type KV interface {
	//Get gets the value for the given key. It returns nil if KV does not contain the key.
	//
	//Warning: Get might not see what was Set or Deleted until after sync().
	//This is a temperary work around to easly warp some apis.
	Get(key []byte) []byte
	//Set sets the value for the given key. It overwrites any previous value for that key.
	//Do not set with len(v) == 0, use Delete instead.
	Set(key []byte, v []byte)
	//Delete deletes the value for the given key, it is a no-op if KV does not contain the key.
	Delete(key []byte)
	//Sync commits all the changes to stable storage. It blocks until done.
	Sync()
	//Close closes KV.
	Close() error
	GC(startAfterKey []byte, f func(key []byte, v []byte) (delete bool, stop bool))
}

//Blob can be used to read/write parts of a long '[]byte' values
//without loading the full data.
//It can be backed by os.File, as it is a subset of os.File, except Size.
//All error reporting are turned off, as none are expected (panic otherwise).
type Blob interface {
	//length in bytes for this block
	Size() int64
	//ReadAt reads len(b) bytes from the File starting at byte offset off.
	ReadAt(b []byte, off int64)
	//WriteAt writes len(b) bytes to the File starting at byte offset off.
	//WriteAt, may or may not sync to disk, at any order.
	WriteAt(b []byte, off int64)
	//Sync commits the all WriteAt to stable storage.
	Sync()
	//Close closes the Block, allowing the freeing of resources. This Block can
	//not be used after Close.
	Close()
}

//LV (Large Values) is an interface for a Block based data store. For
//storing large values with length known in advance and as part of key.
//The same key with two sizes represent two different Blobs.
type LV interface {
	//New creates a new block in LV, or nil if the key already exist.
	//size > 0
	New(key []byte, size int64) Blob
	//Get returns the Block referenced by key, or nil if key is not found.
	Get(key []byte, size int64) Blob
	//Move changes the key from from to to, and resizes to newSize.
	//If from does not exist, or to already exist, then an error is reported.
	//The Blob keyed by from should not be open.
	Move(oldKey []byte, oldSize int64, newKey []byte, newSize int64) error

	//Delete removes the Block repersented by key from LV.
	//Returns true if blob is not in LV after Delete.
	Delete(key []byte, size int64) (bool, error)

	Close() error
}

type FileState int

const (
	//FileNone means we don't have this file
	FileNone FileState = iota
	//FilePart means we have parts of this file
	FilePart
	//FileComplete means we have all of this file
	FileComplete
	//FileExpired means the ttl is outdated, gc maybe removing this file
	FileExpired
)

//MetaStore is a store for hash meta data on top of KV.
//Such as: Child hashes, data length/hash level, and TTL.
type MetaStore interface {
	//InnerHashMinLevel reports the amount of inner hash saved, everything at or
	//above the Level can be retrieved
	InnerHashMinLevel() ht.Level
	//GetInnerHashes reads inner hashes at level and offset, len(hs) should be a
	//multiple of hash length (32).
	//
	//An error is reported if parts of hs is unknown, parts known is still filled
	//with unknown parts as all zeros.
	//
	//If parts of hs are impossible by index range, it panics
	GetInnerHashes(key HLKey, hs []byte, level ht.Level, off ht.Nodes) error

	PutInnerHashes(key HLKey, hs []byte, level ht.Level, off ht.Nodes) (has ht.Nodes, complete bool, err error)

	//TTLGet gets the TTL of a file
	//TTLGet does not imply having a file. Only that the file is desired to be keeped.
	TTLGet(key HLKey) TTL
	//TTLSetAtleast updates the TTL to be at least coved to util, the total
	//increase is multiplied by key.Length() to return costs in storage time by
	//byteMonth.
	//Savings in deduplication can be reflected in byteMonth.
	TTLSetAtleast(key HLKey, util TTL) (byteMonth int64)

	//Close closes MetaStore
	Close()
}

//Database is the full interface for all the data of a server,
//WIP
type Database interface {
	MetaStore
	//GetState checks if we have a file or not, or in progress
	GetState(key HLKey) FileState
	//GetAt reads len(b) bytes of file key from offset off.
	//returns error if b can't be read completely
	GetAt(key HLKey, b []byte, off int64) error
	//PutAt writes b to file from offset off.
	//
	//has returns the number of leaf nodes completed. This can report the progress.
	//Iff has is full, complete is true
	//
	//If hash checking fails, an err is reported. The server can us this information
	//to demote the source.
	PutAt(key HLKey, b []byte, off int64) (has ht.Nodes, complete bool, err error)

	//Import a file from reader
	ImportFromReader(r io.Reader) HLKey
}

type metaValue struct {
	ttl TTL
}
<|endoftext|>


<commit_before>package sequence

import (
	"errors"
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

var (
	ErrNotSeqfile    = errors.New("not a sequence file")
	ErrFrameExists   = errors.New("frame exists")
	ErrNegativeFrame = errors.New("nagative frame")
)

// Splitter is a file name splitter.
type Splitter struct {
	re *regexp.Regexp
}

// reSplit is default regular expression for Splitter.
var reSplit = regexp.MustCompile(`(.*\D)*(\d+)(.*?)$`)

// NewSplitter creates a new splitter.
func NewSplitter() *Splitter {
	return &Splitter{
		re: reSplit,
	}
}

// SetRegexp let users of this package to make their own splitter.
// Splitter always assume that the regular expression is right.
// So who makes their own splitter should ensure that it is right.
func (s *Splitter) SetRegexp(re *regexp.Regexp) {
	s.re = re
}

// Split takes file name and splits it into 3 parts,
// which is pre, digits, and post.
func (s *Splitter) Split(fname string) (pre, digits, post string, err error) {
	m := s.re.FindStringSubmatch(fname)
	if m == nil {
		return "", "", "", ErrNotSeqfile
	}
	return m[1], m[2], m[3], nil
}

// Fmt{Sharp, DollarF, PrecentD} are pre-defined formatter, that covers most user's need.
var (
	FmtSharp = func(pre, digits, post string) string {
		return pre + strings.Repeat("#", len(digits)) + post
	}
	FmtDollarF = func(pre, digits, post string) string {
		return pre + "$F" + strconv.Itoa(len(digits)) + post
	}
	FmtPercentD = func(pre, digits, post string) string {
		return pre + "%0" + strconv.Itoa(len(digits)) + "d" + post
	}
)

// A Manager is a sequence manager.
type Manager struct {
	Seqs map[string]*Seq

	splitter   *Splitter
	formatting func(pre, digits, post string) string
}

// NewManager creates a new sequence manager.
func NewManager(splitter *Splitter, formatting func(pre, digits, post string) string) *Manager {
	return &Manager{
		Seqs:       make(map[string]*Seq),
		splitter:   splitter,
		formatting: formatting,
	}
}

// Add adds a file to the manager.
// If the file's sequence is not exist yet, it will create a new sequence automatically.
func (m *Manager) Add(fname string) error {
	pre, digits, post, err := m.splitter.Split(fname)
	if err != nil {
		return err
	}

	name := m.formatting(pre, digits, post)
	frame, _ := strconv.Atoi(digits)

	s, ok := m.Seqs[name]
	if !ok {
		s = NewSeq()
		m.Seqs[name] = s
	}
	return s.AddFrame(frame)
}

// SeqNames returns it's sequence names in ascending order.
func (m *Manager) SeqNames() []string {
	names := []string{}
	for n := range m.Seqs {
		names = append(names, n)
	}
	sort.Strings(names)
	return names
}

// String returns a string that shows it's sequences.
// It will be multiple lines if it has more than one sequence.
func (m *Manager) String() string {
	str := ""
	for _, n := range m.SeqNames() {
		if str != "" {
			str += "\n"
		}
		str += fmt.Sprintf("%s %s", n, m.Seqs[n])
	}
	return str
}

type Seq struct {
	frames map[int]struct{}
}

func NewSeq() *Seq {
	return &Seq{
		frames: make(map[int]struct{}),
	}
}

func (s *Seq) AddFrame(f int) error {
	if f < 0 {
		return ErrNegativeFrame
	}
	if _, ok := s.frames[f]; ok {
		return ErrFrameExists
	}
	s.frames[f] = struct{}{}
	return nil
}

func (s *Seq) Ranges() []*Range {
	if len(s.frames) == 0 {
		return []*Range{}
	}

	frames := []int{}
	for f := range s.frames {
		frames = append(frames, f)
	}
	sort.Ints(frames)

	rngs := []*Range{}
	r := NewRange(frames[0])
	rngs = append(rngs, r)
	for _, f := range frames[1:] {
		ok := r.Extend(f)
		if !ok {
			r = NewRange(f)
			rngs = append(rngs, r)
		}
	}
	return rngs
}

func (s *Seq) String() string {
	str := ""
	rngs := s.Ranges()
	for _, r := range rngs {
		if str != "" {
			str += " "
		}
		str += r.String()
	}
	return str
}

type Range struct {
	minf int
	maxf int
}

func NewRange(f int) *Range {
	return &Range{
		minf: f,
		maxf: f,
	}
}

func (r *Range) Extend(f int) bool {
	if f != r.maxf+1 {
		return false
	}
	r.maxf = f
	return true
}

func (r *Range) String() string {
	if r.minf == r.maxf {
		return fmt.Sprintf("%d", r.minf)
	}
	return fmt.Sprintf("%d-%d", r.minf, r.maxf)
}
<commit_msg>add additional comments.<commit_after>package sequence

import (
	"errors"
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

var (
	ErrNotSeqfile    = errors.New("not a sequence file")
	ErrFrameExists   = errors.New("frame exists")
	ErrNegativeFrame = errors.New("nagative frame")
)

// Splitter is a file name splitter.
type Splitter struct {
	re *regexp.Regexp
}

// reSplit is default regular expression for Splitter.
var reSplit = regexp.MustCompile(`(.*\D)*(\d+)(.*?)$`)

// NewSplitter creates a new splitter.
func NewSplitter() *Splitter {
	return &Splitter{
		re: reSplit,
	}
}

// SetRegexp let users of this package to make their own splitter.
// Splitter always assume that the regular expression is right.
// So who makes their own splitter should ensure that it is right.
func (s *Splitter) SetRegexp(re *regexp.Regexp) {
	s.re = re
}

// Split takes file name and splits it into 3 parts,
// which is pre, digits, and post.
func (s *Splitter) Split(fname string) (pre, digits, post string, err error) {
	m := s.re.FindStringSubmatch(fname)
	if m == nil {
		return "", "", "", ErrNotSeqfile
	}
	return m[1], m[2], m[3], nil
}

// Fmt{Sharp, DollarF, PrecentD} are pre-defined formatter, that covers most user's need.
var (
	FmtSharp = func(pre, digits, post string) string {
		return pre + strings.Repeat("#", len(digits)) + post
	}
	FmtDollarF = func(pre, digits, post string) string {
		return pre + "$F" + strconv.Itoa(len(digits)) + post
	}
	FmtPercentD = func(pre, digits, post string) string {
		return pre + "%0" + strconv.Itoa(len(digits)) + "d" + post
	}
)

// A Manager is a sequence manager.
type Manager struct {
	Seqs map[string]*Seq

	splitter   *Splitter
	formatting func(pre, digits, post string) string
}

// NewManager creates a new sequence manager.
func NewManager(splitter *Splitter, formatting func(pre, digits, post string) string) *Manager {
	return &Manager{
		Seqs:       make(map[string]*Seq),
		splitter:   splitter,
		formatting: formatting,
	}
}

// Add adds a file to the manager.
// If the file's sequence is not exist yet, it will create a new sequence automatically.
func (m *Manager) Add(fname string) error {
	pre, digits, post, err := m.splitter.Split(fname)
	if err != nil {
		return err
	}

	name := m.formatting(pre, digits, post)
	frame, _ := strconv.Atoi(digits)

	s, ok := m.Seqs[name]
	if !ok {
		s = NewSeq()
		m.Seqs[name] = s
	}
	return s.AddFrame(frame)
}

// SeqNames returns it's sequence names in ascending order.
func (m *Manager) SeqNames() []string {
	names := []string{}
	for n := range m.Seqs {
		names = append(names, n)
	}
	sort.Strings(names)
	return names
}

// String returns a string that shows it's sequences.
// It will be multiple lines if it has more than one sequence.
func (m *Manager) String() string {
	str := ""
	for _, n := range m.SeqNames() {
		if str != "" {
			str += "\n"
		}
		str += fmt.Sprintf("%s %s", n, m.Seqs[n])
	}
	return str
}

// A Seq is a frame sequence. It does not hold sequence name.
type Seq struct {
	frames map[int]struct{}
}

// NewSeq creates a new sequence.
func NewSeq() *Seq {
	return &Seq{
		frames: make(map[int]struct{}),
	}
}

// AddFrame adds a frame into sequence.
// It treats negative frames are invalid.
// So ErrNegativeFrame error will return when it takes a negative frame.
func (s *Seq) AddFrame(f int) error {
	if f < 0 {
		return ErrNegativeFrame
	}
	if _, ok := s.frames[f]; ok {
		return ErrFrameExists
	}
	s.frames[f] = struct{}{}
	return nil
}

// Ranges convert a sequence to several contiguous ranges.
func (s *Seq) Ranges() []*Range {
	if len(s.frames) == 0 {
		return []*Range{}
	}

	frames := []int{}
	for f := range s.frames {
		frames = append(frames, f)
	}
	sort.Ints(frames)

	rngs := []*Range{}
	r := NewRange(frames[0])
	rngs = append(rngs, r)
	for _, f := range frames[1:] {
		ok := r.Extend(f)
		if !ok {
			r = NewRange(f)
			rngs = append(rngs, r)
		}
	}
	return rngs
}

// String expresses a sequence using ranges.
func (s *Seq) String() string {
	str := ""
	rngs := s.Ranges()
	for _, r := range rngs {
		if str != "" {
			str += " "
		}
		str += r.String()
	}
	return str
}

// Range is a contiguous frame range.
type Range struct {
	minf int
	maxf int
}

// NewRange creates a new range.
func NewRange(f int) *Range {
	return &Range{
		minf: f,
		maxf: f,
	}
}

// Extend extends a range by one if the frame is bigger than current max frame by 1.
// If it extends, it returns true, or it returns false.
func (r *Range) Extend(f int) bool {
	if f != r.maxf+1 {
		return false
	}
	r.maxf = f
	return true
}

// String express the range with dash. Like "1-10".
// But if the min and max is same, it will just show one. Like "5".
func (r *Range) String() string {
	if r.minf == r.maxf {
		return fmt.Sprintf("%d", r.minf)
	}
	return fmt.Sprintf("%d-%d", r.minf, r.maxf)
}
<|endoftext|>


<commit_before>package pgx_test

import (
	"errors"
	"math/rand"
	"testing"
	"time"

	"github.com/jackc/fake"
	"github.com/jackc/pgx"
)

type execer interface {
	Exec(sql string, arguments ...interface{}) (commandTag pgx.CommandTag, err error)
}
type queryer interface {
	Query(sql string, args ...interface{}) (*pgx.Rows, error)
}
type queryRower interface {
	QueryRow(sql string, args ...interface{}) *pgx.Row
}

func TestStressConnPool(t *testing.T) {
	maxConnections := 8
	pool := createConnPool(t, maxConnections)
	defer pool.Close()

	setupStressDB(t, pool)

	actions := []struct {
		name string
		fn   func(*pgx.ConnPool, int) error
	}{
		{"insertUnprepared", func(p *pgx.ConnPool, n int) error { return insertUnprepared(p, n) }},
		{"queryRowWithoutParams", func(p *pgx.ConnPool, n int) error { return queryRowWithoutParams(p, n) }},
		{"query", func(p *pgx.ConnPool, n int) error { return queryCloseEarly(p, n) }},
		{"queryCloseEarly", func(p *pgx.ConnPool, n int) error { return query(p, n) }},
		{"queryErrorWhileReturningRows", func(p *pgx.ConnPool, n int) error { return queryErrorWhileReturningRows(p, n) }},
		{"txInsertRollback", txInsertRollback},
		{"txInsertCommit", txInsertCommit},
		{"txMultipleQueries", txMultipleQueries},
		{"notify", notify},
		{"listenAndPoolUnlistens", listenAndPoolUnlistens},
		{"reset", func(p *pgx.ConnPool, n int) error { p.Reset(); return nil }},
	}

	actionCount := 5000
	if testing.Short() {
		actionCount /= 10
	}
	workerCount := 16

	workChan := make(chan int)
	doneChan := make(chan struct{})
	errChan := make(chan error)

	work := func() {
		for n := range workChan {
			action := actions[rand.Intn(len(actions))]
			err := action.fn(pool, n)
			if err != nil {
				errChan <- err
				break
			}
		}
		doneChan <- struct{}{}
	}

	for i := 0; i < workerCount; i++ {
		go work()
	}

	for i := 0; i < actionCount; i++ {
		select {
		case workChan <- i:
		case err := <-errChan:
			close(workChan)
			t.Fatal(err)
		}
	}
	close(workChan)

	for i := 0; i < workerCount; i++ {
		<-doneChan
	}
}

func TestStressTLSConnection(t *testing.T) {
	t.Parallel()

	if tlsConnConfig == nil {
		t.Skip("Skipping due to undefined tlsConnConfig")
	}

	if testing.Short() {
		t.Skip("Skipping due to testing -short")
	}

	conn, err := pgx.Connect(*tlsConnConfig)
	if err != nil {
		t.Fatalf("Unable to establish connection: %v", err)
	}
	defer conn.Close()

	queryCount := 50
	if testing.Short() {
		queryCount /= 10
	}

	for i := 0; i < queryCount; i++ {
		sql := `select * from generate_series(1, $1)`

		rows, err := conn.Query(sql, 2000000)
		if err != nil {
			t.Fatal(err)
		}

		var n int32
		for rows.Next() {
			rows.Scan(&n)
		}

		if rows.Err() != nil {
			t.Fatalf("queryCount: %d, Row number: %d. %v", i, n, rows.Err())
		}
	}
}

func setupStressDB(t *testing.T, pool *pgx.ConnPool) {
	_, err := pool.Exec(`
		drop table if exists widgets;
		create table widgets(
			id serial primary key,
			name varchar not null,
			description text,
			creation_time timestamptz
		);
`)
	if err != nil {
		t.Fatal(err)
	}
}

func insertUnprepared(e execer, actionNum int) error {
	sql := `
		insert into widgets(name, description, creation_time)
		values($1, $2, $3)`

	_, err := e.Exec(sql, fake.ProductName(), fake.Sentences(), time.Now())
	return err
}

func queryRowWithoutParams(qr queryRower, actionNum int) error {
	var id int32
	var name, description string
	var creationTime time.Time

	sql := `select * from widgets order by random() limit 1`

	err := qr.QueryRow(sql).Scan(&id, &name, &description, &creationTime)
	if err == pgx.ErrNoRows {
		return nil
	}
	return err
}

func query(q queryer, actionNum int) error {
	sql := `select * from widgets order by random() limit $1`

	rows, err := q.Query(sql, 10)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var id int32
		var name, description string
		var creationTime time.Time
		rows.Scan(&id, &name, &description, &creationTime)
	}

	return rows.Err()
}

func queryCloseEarly(q queryer, actionNum int) error {
	sql := `select * from generate_series(1,$1)`

	rows, err := q.Query(sql, 100)
	if err != nil {
		return err
	}
	defer rows.Close()

	for i := 0; i < 10 && rows.Next(); i++ {
		var n int32
		rows.Scan(&n)
	}
	rows.Close()

	return rows.Err()
}

func queryErrorWhileReturningRows(q queryer, actionNum int) error {
	// This query should divide by 0 within the first number of rows
	sql := `select 42 / (random() * 20)::integer from generate_series(1,100000)`

	rows, err := q.Query(sql)
	if err != nil {
		return nil
	}
	defer rows.Close()

	for rows.Next() {
		var n int32
		rows.Scan(&n)
	}

	if _, ok := rows.Err().(pgx.PgError); ok {
		return nil
	}
	return rows.Err()
}

func notify(pool *pgx.ConnPool, actionNum int) error {
	_, err := pool.Exec("notify stress")
	return err
}

func listenAndPoolUnlistens(pool *pgx.ConnPool, actionNum int) error {
	conn, err := pool.Acquire()
	if err != nil {
		return err
	}
	defer pool.Release(conn)

	err = conn.Listen("stress")
	if err != nil {
		return err
	}

	_, err = conn.WaitForNotification(100 * time.Millisecond)
	if err == pgx.ErrNotificationTimeout {
		return nil
	}
	return err
}

func txInsertRollback(pool *pgx.ConnPool, actionNum int) error {
	tx, err := pool.Begin()
	if err != nil {
		return err
	}

	sql := `
		insert into widgets(name, description, creation_time)
		values($1, $2, $3)`

	_, err = tx.Exec(sql, fake.ProductName(), fake.Sentences(), time.Now())
	if err != nil {
		return err
	}

	return tx.Rollback()
}

func txInsertCommit(pool *pgx.ConnPool, actionNum int) error {
	tx, err := pool.Begin()
	if err != nil {
		return err
	}

	sql := `
		insert into widgets(name, description, creation_time)
		values($1, $2, $3)`

	_, err = tx.Exec(sql, fake.ProductName(), fake.Sentences(), time.Now())
	if err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}

func txMultipleQueries(pool *pgx.ConnPool, actionNum int) error {
	tx, err := pool.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	errExpectedTxDeath := errors.New("Expected tx death")

	actions := []struct {
		name string
		fn   func() error
	}{
		{"insertUnprepared", func() error { return insertUnprepared(tx, actionNum) }},
		{"queryRowWithoutParams", func() error { return queryRowWithoutParams(tx, actionNum) }},
		{"query", func() error { return query(tx, actionNum) }},
		{"queryCloseEarly", func() error { return queryCloseEarly(tx, actionNum) }},
		{"queryErrorWhileReturningRows", func() error {
			err := queryErrorWhileReturningRows(tx, actionNum)
			if err != nil {
				return err
			}
			return errExpectedTxDeath
		}},
	}

	for i := 0; i < 20; i++ {
		action := actions[rand.Intn(len(actions))]
		err := action.fn()
		if err == errExpectedTxDeath {
			return nil
		} else if err != nil {
			return err
		}
	}

	return tx.Commit()
}
<commit_msg>Remove dead code<commit_after>package pgx_test

import (
	"errors"
	"math/rand"
	"testing"
	"time"

	"github.com/jackc/fake"
	"github.com/jackc/pgx"
)

type execer interface {
	Exec(sql string, arguments ...interface{}) (commandTag pgx.CommandTag, err error)
}
type queryer interface {
	Query(sql string, args ...interface{}) (*pgx.Rows, error)
}
type queryRower interface {
	QueryRow(sql string, args ...interface{}) *pgx.Row
}

func TestStressConnPool(t *testing.T) {
	maxConnections := 8
	pool := createConnPool(t, maxConnections)
	defer pool.Close()

	setupStressDB(t, pool)

	actions := []struct {
		name string
		fn   func(*pgx.ConnPool, int) error
	}{
		{"insertUnprepared", func(p *pgx.ConnPool, n int) error { return insertUnprepared(p, n) }},
		{"queryRowWithoutParams", func(p *pgx.ConnPool, n int) error { return queryRowWithoutParams(p, n) }},
		{"query", func(p *pgx.ConnPool, n int) error { return queryCloseEarly(p, n) }},
		{"queryCloseEarly", func(p *pgx.ConnPool, n int) error { return query(p, n) }},
		{"queryErrorWhileReturningRows", func(p *pgx.ConnPool, n int) error { return queryErrorWhileReturningRows(p, n) }},
		{"txInsertRollback", txInsertRollback},
		{"txInsertCommit", txInsertCommit},
		{"txMultipleQueries", txMultipleQueries},
		{"notify", notify},
		{"listenAndPoolUnlistens", listenAndPoolUnlistens},
		{"reset", func(p *pgx.ConnPool, n int) error { p.Reset(); return nil }},
	}

	actionCount := 5000
	if testing.Short() {
		actionCount /= 10
	}
	workerCount := 16

	workChan := make(chan int)
	doneChan := make(chan struct{})
	errChan := make(chan error)

	work := func() {
		for n := range workChan {
			action := actions[rand.Intn(len(actions))]
			err := action.fn(pool, n)
			if err != nil {
				errChan <- err
				break
			}
		}
		doneChan <- struct{}{}
	}

	for i := 0; i < workerCount; i++ {
		go work()
	}

	for i := 0; i < actionCount; i++ {
		select {
		case workChan <- i:
		case err := <-errChan:
			close(workChan)
			t.Fatal(err)
		}
	}
	close(workChan)

	for i := 0; i < workerCount; i++ {
		<-doneChan
	}
}

func TestStressTLSConnection(t *testing.T) {
	t.Parallel()

	if tlsConnConfig == nil {
		t.Skip("Skipping due to undefined tlsConnConfig")
	}

	if testing.Short() {
		t.Skip("Skipping due to testing -short")
	}

	conn, err := pgx.Connect(*tlsConnConfig)
	if err != nil {
		t.Fatalf("Unable to establish connection: %v", err)
	}
	defer conn.Close()

	for i := 0; i < 50; i++ {
		sql := `select * from generate_series(1, $1)`

		rows, err := conn.Query(sql, 2000000)
		if err != nil {
			t.Fatal(err)
		}

		var n int32
		for rows.Next() {
			rows.Scan(&n)
		}

		if rows.Err() != nil {
			t.Fatalf("queryCount: %d, Row number: %d. %v", i, n, rows.Err())
		}
	}
}

func setupStressDB(t *testing.T, pool *pgx.ConnPool) {
	_, err := pool.Exec(`
		drop table if exists widgets;
		create table widgets(
			id serial primary key,
			name varchar not null,
			description text,
			creation_time timestamptz
		);
`)
	if err != nil {
		t.Fatal(err)
	}
}

func insertUnprepared(e execer, actionNum int) error {
	sql := `
		insert into widgets(name, description, creation_time)
		values($1, $2, $3)`

	_, err := e.Exec(sql, fake.ProductName(), fake.Sentences(), time.Now())
	return err
}

func queryRowWithoutParams(qr queryRower, actionNum int) error {
	var id int32
	var name, description string
	var creationTime time.Time

	sql := `select * from widgets order by random() limit 1`

	err := qr.QueryRow(sql).Scan(&id, &name, &description, &creationTime)
	if err == pgx.ErrNoRows {
		return nil
	}
	return err
}

func query(q queryer, actionNum int) error {
	sql := `select * from widgets order by random() limit $1`

	rows, err := q.Query(sql, 10)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var id int32
		var name, description string
		var creationTime time.Time
		rows.Scan(&id, &name, &description, &creationTime)
	}

	return rows.Err()
}

func queryCloseEarly(q queryer, actionNum int) error {
	sql := `select * from generate_series(1,$1)`

	rows, err := q.Query(sql, 100)
	if err != nil {
		return err
	}
	defer rows.Close()

	for i := 0; i < 10 && rows.Next(); i++ {
		var n int32
		rows.Scan(&n)
	}
	rows.Close()

	return rows.Err()
}

func queryErrorWhileReturningRows(q queryer, actionNum int) error {
	// This query should divide by 0 within the first number of rows
	sql := `select 42 / (random() * 20)::integer from generate_series(1,100000)`

	rows, err := q.Query(sql)
	if err != nil {
		return nil
	}
	defer rows.Close()

	for rows.Next() {
		var n int32
		rows.Scan(&n)
	}

	if _, ok := rows.Err().(pgx.PgError); ok {
		return nil
	}
	return rows.Err()
}

func notify(pool *pgx.ConnPool, actionNum int) error {
	_, err := pool.Exec("notify stress")
	return err
}

func listenAndPoolUnlistens(pool *pgx.ConnPool, actionNum int) error {
	conn, err := pool.Acquire()
	if err != nil {
		return err
	}
	defer pool.Release(conn)

	err = conn.Listen("stress")
	if err != nil {
		return err
	}

	_, err = conn.WaitForNotification(100 * time.Millisecond)
	if err == pgx.ErrNotificationTimeout {
		return nil
	}
	return err
}

func txInsertRollback(pool *pgx.ConnPool, actionNum int) error {
	tx, err := pool.Begin()
	if err != nil {
		return err
	}

	sql := `
		insert into widgets(name, description, creation_time)
		values($1, $2, $3)`

	_, err = tx.Exec(sql, fake.ProductName(), fake.Sentences(), time.Now())
	if err != nil {
		return err
	}

	return tx.Rollback()
}

func txInsertCommit(pool *pgx.ConnPool, actionNum int) error {
	tx, err := pool.Begin()
	if err != nil {
		return err
	}

	sql := `
		insert into widgets(name, description, creation_time)
		values($1, $2, $3)`

	_, err = tx.Exec(sql, fake.ProductName(), fake.Sentences(), time.Now())
	if err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}

func txMultipleQueries(pool *pgx.ConnPool, actionNum int) error {
	tx, err := pool.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	errExpectedTxDeath := errors.New("Expected tx death")

	actions := []struct {
		name string
		fn   func() error
	}{
		{"insertUnprepared", func() error { return insertUnprepared(tx, actionNum) }},
		{"queryRowWithoutParams", func() error { return queryRowWithoutParams(tx, actionNum) }},
		{"query", func() error { return query(tx, actionNum) }},
		{"queryCloseEarly", func() error { return queryCloseEarly(tx, actionNum) }},
		{"queryErrorWhileReturningRows", func() error {
			err := queryErrorWhileReturningRows(tx, actionNum)
			if err != nil {
				return err
			}
			return errExpectedTxDeath
		}},
	}

	for i := 0; i < 20; i++ {
		action := actions[rand.Intn(len(actions))]
		err := action.fn()
		if err == errExpectedTxDeath {
			return nil
		} else if err != nil {
			return err
		}
	}

	return tx.Commit()
}
<|endoftext|>


<commit_before>package clusterf
/*
 * Internal services states, maintained from config changes.
 */

import (
    "github.com/qmsk/clusterf/config"
    "fmt"
    "log"
)

type Services struct {
    services    map[string]*Service
    routes      Routes

    driver      *IPVSDriver
}

func NewServices() *Services {
    return &Services{
        services:   make(map[string]*Service),
        routes:     makeRoutes(),
    }
}

// Return Service for named service, possibly creating a new (empty) Service.
func (self *Services) get(name string) *Service {
    service, serviceExists := self.services[name]

    if !serviceExists {
        service = newService(name)
        self.services[name] = service
    }

    return service
}

// Return all currently valid Services
func (self *Services) Services() []*Service {
    services := make([]*Service, 0, len(self.services))

    for _, service := range self.services {
        if service.Frontend == nil {
            continue
        }

        services = append(services, service)
    }

    return services
}

/* Configuration actions */

// Configuration action on a service itself
// handle service-delete actions
// new service creation is implicitly handled when calling this
func (self *Services) configService(service *Service, action config.Action, serviceConfig *config.ConfigService) {
    log.Printf("clusterf:Service %s: %s %+v\n", service.Name, action, serviceConfig)

    switch action {
    case config.DelConfig:
        delete(self.services, service.Name)

        service.delFrontend()
    }
}

func (self *Services) configRoute(route *Route, action config.Action, routeConfig *config.ConfigRoute) {
    log.Printf("clusterf:Route %s: %s %+v\n", route.Name, action, routeConfig)

    switch action {
    case config.NewConfig, config.SetConfig:
        if err := route.config(action, routeConfig.Route); err != nil {
            log.Printf("clusterf:Route %s: %s\n", route.Name, err)
        } else {
            log.Printf("clusterf:Route %s: %+v\n", route.Name, route)
        }

    case config.DelConfig:
        self.routes.del(route.Name)
    }

    // TODO: update services?
}

func (self *Services) config(action config.Action, baseConfig config.Config) {
    log.Printf("clusterf: config %s %#v\n", action, baseConfig)

    switch applyConfig:= baseConfig.(type) {
    case *config.ConfigService:
        serviceConfig := baseConfig.(*config.ConfigService)

        if serviceConfig.ServiceName == "" {
            // all services
            for _, service := range self.services {
                self.configService(service, action, serviceConfig)
            }
        } else {
            service := self.get(serviceConfig.ServiceName)

            self.configService(service, action, serviceConfig)
        }

    case *config.ConfigServiceFrontend:
        frontendConfig := baseConfig.(*config.ConfigServiceFrontend)

        service := self.get(frontendConfig.ServiceName)

        service.configFrontend(action, frontendConfig)

    case *config.ConfigServiceBackend:
        backendConfig := baseConfig.(*config.ConfigServiceBackend)

        service := self.get(backendConfig.ServiceName)

        if backendConfig.BackendName == "" {
            // all service backends
            for backendName, _ := range service.Backends {
                service.configBackend(backendName, action, backendConfig)
            }
        } else {
            service.configBackend(backendConfig.BackendName, action, backendConfig)
        }

    case *config.ConfigRoute:
        if applyConfig.RouteName == "" {
            // all routes
            for _, route := range self.routes {
                self.configRoute(route, action, applyConfig)
            }
        } else {
            route := self.routes.get(applyConfig.RouteName)

            self.configRoute(route, action, applyConfig)
        }

    default:
        panic(fmt.Errorf("Unknown config type: %#v", baseConfig))
    }
}

// Initialize config before driver sync
func (self *Services) NewConfig(baseConfig config.Config) {
    if self.driver != nil {
        panic("NewConfig after driver sync")
    }

    self.config(config.NewConfig, baseConfig)
}

// Sync initial configuration loaded via NewConfig() to IPVS
//
// Begins by flushing the IPVS state
func (self *Services) SyncIPVS(ipvsConfig IpvsConfig) (*IPVSDriver, error) {
    if ipvsDriver, err := ipvsConfig.setup(self.routes); err != nil {
        return nil, err
    } else {
        self.driver = ipvsDriver
    }

    // begin sync
    if err := self.driver.sync(); err != nil {
        return nil, err
    }

    for _, service := range self.services {
        service.sync(self.driver)
    }

    return self.driver, nil
}

// Apply changes to the current configuration, updating the running driver
func (self *Services) ConfigEvent(event config.Event) {
    if self.driver == nil {
        panic("ConfigEvent before driver sync")
    }

    self.config(event.Action, event.Config)
}
<commit_msg>fix newly added Services to be sync()'d after initial startup<commit_after>package clusterf
/*
 * Internal services states, maintained from config changes.
 */

import (
    "github.com/qmsk/clusterf/config"
    "fmt"
    "log"
)

type Services struct {
    services    map[string]*Service
    routes      Routes

    driver      *IPVSDriver
}

func NewServices() *Services {
    return &Services{
        services:   make(map[string]*Service),
        routes:     makeRoutes(),
    }
}

// Return Service for named service, possibly creating a new (empty) Service.
func (self *Services) get(name string) *Service {
    service, serviceExists := self.services[name]

    if !serviceExists {
        service = newService(name)
        self.services[name] = service

        // initial sync
        if self.driver != nil {
            service.sync(self.driver)
        }
    }

    return service
}

// Return all currently valid Services
func (self *Services) Services() []*Service {
    services := make([]*Service, 0, len(self.services))

    for _, service := range self.services {
        if service.Frontend == nil {
            continue
        }

        services = append(services, service)
    }

    return services
}

/* Configuration actions */

// Configuration action on a service itself
// handle service-delete actions
// new service creation is implicitly handled when calling this
func (self *Services) configService(service *Service, action config.Action, serviceConfig *config.ConfigService) {
    log.Printf("clusterf:Service %s: %s %+v\n", service.Name, action, serviceConfig)

    switch action {
    case config.DelConfig:
        delete(self.services, service.Name)

        service.delFrontend()
    }
}

func (self *Services) configRoute(route *Route, action config.Action, routeConfig *config.ConfigRoute) {
    log.Printf("clusterf:Route %s: %s %+v\n", route.Name, action, routeConfig)

    switch action {
    case config.NewConfig, config.SetConfig:
        if err := route.config(action, routeConfig.Route); err != nil {
            log.Printf("clusterf:Route %s: %s\n", route.Name, err)
        } else {
            log.Printf("clusterf:Route %s: %+v\n", route.Name, route)
        }

    case config.DelConfig:
        self.routes.del(route.Name)
    }

    // TODO: update services?
}

func (self *Services) config(action config.Action, baseConfig config.Config) {
    log.Printf("clusterf: config %s %#v\n", action, baseConfig)

    switch applyConfig:= baseConfig.(type) {
    case *config.ConfigService:
        serviceConfig := baseConfig.(*config.ConfigService)

        if serviceConfig.ServiceName == "" {
            // all services
            for _, service := range self.services {
                self.configService(service, action, serviceConfig)
            }
        } else {
            service := self.get(serviceConfig.ServiceName)

            self.configService(service, action, serviceConfig)
        }

    case *config.ConfigServiceFrontend:
        frontendConfig := baseConfig.(*config.ConfigServiceFrontend)

        service := self.get(frontendConfig.ServiceName)

        service.configFrontend(action, frontendConfig)

    case *config.ConfigServiceBackend:
        backendConfig := baseConfig.(*config.ConfigServiceBackend)

        service := self.get(backendConfig.ServiceName)

        if backendConfig.BackendName == "" {
            // all service backends
            for backendName, _ := range service.Backends {
                service.configBackend(backendName, action, backendConfig)
            }
        } else {
            service.configBackend(backendConfig.BackendName, action, backendConfig)
        }

    case *config.ConfigRoute:
        if applyConfig.RouteName == "" {
            // all routes
            for _, route := range self.routes {
                self.configRoute(route, action, applyConfig)
            }
        } else {
            route := self.routes.get(applyConfig.RouteName)

            self.configRoute(route, action, applyConfig)
        }

    default:
        panic(fmt.Errorf("Unknown config type: %#v", baseConfig))
    }
}

// Initialize config before driver sync
func (self *Services) NewConfig(baseConfig config.Config) {
    if self.driver != nil {
        panic("NewConfig after driver sync")
    }

    self.config(config.NewConfig, baseConfig)
}

// Sync initial configuration loaded via NewConfig() to IPVS
//
// Begins by flushing the IPVS state
func (self *Services) SyncIPVS(ipvsConfig IpvsConfig) (*IPVSDriver, error) {
    if ipvsDriver, err := ipvsConfig.setup(self.routes); err != nil {
        return nil, err
    } else {
        self.driver = ipvsDriver
    }

    // begin sync
    if err := self.driver.sync(); err != nil {
        return nil, err
    }

    for _, service := range self.services {
        service.sync(self.driver)
    }

    return self.driver, nil
}

// Apply changes to the current configuration, updating the running driver
func (self *Services) ConfigEvent(event config.Event) {
    if self.driver == nil {
        panic("ConfigEvent before driver sync")
    }

    self.config(event.Action, event.Config)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kyma Authors.

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

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

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

package function

import (
	"encoding/json"
	"os"
	"testing"
	"time"

	buildv1alpha1 "github.com/knative/build/pkg/apis/build/v1alpha1"
	servingv1alpha1 "github.com/knative/serving/pkg/apis/serving/v1alpha1"
	runtimev1alpha1 "github.com/kyma-incubator/runtime/pkg/apis/runtime/v1alpha1"
	"github.com/onsi/gomega"
	"golang.org/x/net/context"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/kubernetes/scheme"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"
)

var c client.Client

var expectedRequest = reconcile.Request{NamespacedName: types.NamespacedName{Name: "foo", Namespace: "default"}}
var depKey = types.NamespacedName{Name: "foo", Namespace: "default"}

const timeout = time.Second * 10

func TestReconcile(t *testing.T) {
	g := gomega.NewGomegaWithT(t)
	fnCreated := &runtimev1alpha1.Function{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "foo",
			Namespace: "default",
		},
		Spec: runtimev1alpha1.FunctionSpec{
			Function:            "main() {asdfasdf}",
			FunctionContentType: "plaintext",
			Size:                "L",
			Runtime:             "nodejs6",
		},
	}

	expectedEnv := []corev1.EnvVar{
		{
			Name:  "FUNC_HANDLER",
			Value: "main",
		},
		{
			Name:  "MOD_NAME",
			Value: "handler",
		},
		{
			Name:  "FUNC_TIMEOUT",
			Value: "180",
		},
		{
			Name:  "FUNC_RUNTIME",
			Value: "nodejs8",
		},
		{
			Name:  "FUNC_MEMORY_LIMIT",
			Value: "128Mi",
		},
		{
			Name:  "FUNC_PORT",
			Value: "8080",
		},
		{
			Name:  "NODE_PATH",
			Value: "$(KUBELESS_INSTALL_VOLUME)/node_modules",
		},
	}

	fnConfig := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "fn-config",
			Namespace: "default",
		},
		Data: map[string]string{
			"dockerRegistry":     "test",
			"serviceAccountName": "build-bot",
			"runtimes": `[
				{
					"ID": "nodejs8",
					"DockerFileName": "dockerfile-nodejs8",
				}
			]`,
		},
	}

	dockerFileConfigNodejs := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "dockerfile-nodejs8",
			Namespace: "default",
		},
		Data: map[string]string{
			"Dockerfile": `FROM kubeless/nodejs@sha256:5c3c21cf29231f25a0d7d2669c6f18c686894bf44e975fcbbbb420c6d045f7e7
				USER root
				RUN export KUBELESS_INSTALL_VOLUME='/kubeless' && \
					mkdir /kubeless && \
					cp /src/handler.js /kubeless && \
					cp /src/package.json /kubeless && \
					/kubeless-npm-install.sh
				USER 1000
			`,
		},
	}
	mgr, err := manager.New(cfg, manager.Options{})
	g.Expect(err).NotTo(gomega.HaveOccurred())
	c = mgr.GetClient()

	recFn, requests := SetupTestReconcile(newReconciler(mgr))
	g.Expect(add(mgr, recFn)).NotTo(gomega.HaveOccurred())

	stopMgr, mgrStopped := StartTestManager(mgr, g)

	defer func() {
		close(stopMgr)
		mgrStopped.Wait()
	}()

	// Create the Function object and expect the Reconcile and Deployment to be created
	err = c.Create(context.TODO(), dockerFileConfigNodejs)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}

	err = c.Create(context.TODO(), fnConfig)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}

	err = c.Create(context.TODO(), fnCreated)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}
	g.Expect(err).NotTo(gomega.HaveOccurred())

	g.Eventually(requests, timeout).Should(gomega.Receive(gomega.Equal(expectedRequest)))

	cm := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
	}

	service := &servingv1alpha1.Service{
		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
	}

	g.Eventually(func() error { return c.Get(context.TODO(), depKey, cm) }, timeout).
		Should(gomega.Succeed())

	g.Eventually(func() error { return c.Get(context.TODO(), depKey, service) }, timeout).
		Should(gomega.Succeed())
	g.Expect(service.Namespace).To(gomega.Equal("default"))

	g.Expect(service.Spec.RunLatest.Configuration.RevisionTemplate.Spec.Container.Env).To(gomega.Equal(expectedEnv))
	build := (*service.Spec.RunLatest.Configuration.Build)
	buildByte, err := build.MarshalJSON()
	if err != nil {
		t.Fatalf("Error while marshaling build object: %v", err)
	}
	var buildSpec buildv1alpha1.BuildSpec
	err = json.Unmarshal(buildByte, &buildSpec)
	if err != nil {
		t.Fatalf("Error while unmarshaling buildSpec: %v", err)
	}
	g.Expect(service.Spec.RunLatest.Configuration.RevisionTemplate.Spec.Container.Image).To(gomega.HavePrefix("test/default-foo"))
	g.Expect(len(buildSpec.Volumes)).To(gomega.Equal(2))
	g.Expect(buildSpec.ServiceAccountName).To(gomega.Equal("build-bot"))
	g.Expect(service.Spec.RunLatest.Configuration.RevisionTemplate.Spec.Container.Image).To(gomega.HavePrefix("test/default-foo"))

	fnFetched := &runtimev1alpha1.Function{}
	g.Expect(c.Get(context.TODO(), depKey, fnFetched)).NotTo(gomega.HaveOccurred())
	g.Expect(fnFetched.Spec).To(gomega.Equal(fnCreated.Spec))

	fnUpdated := fnFetched.DeepCopy()
	fnUpdated.Spec.Function = `main() {return "bla"}`
	fnUpdated.Spec.Deps = `dependencies`

	fnFetched = &runtimev1alpha1.Function{}
	g.Expect(c.Update(context.TODO(), fnUpdated)).NotTo(gomega.HaveOccurred())
	g.Expect(c.Get(context.TODO(), depKey, fnFetched)).NotTo(gomega.HaveOccurred())
	g.Expect(fnFetched.Spec).To(gomega.Equal(fnUpdated.Spec))

	g.Eventually(requests, timeout).Should(gomega.Receive(gomega.Equal(expectedRequest)))

	cmUpdated := &corev1.ConfigMap{}
	g.Eventually(func() string {
		c.Get(context.TODO(), depKey, cmUpdated)
		return cmUpdated.Data["handler.js"]
	}, timeout, 1*time.Second).Should(gomega.Equal(fnUpdated.Spec.Function))
	g.Eventually(func() string {
		c.Get(context.TODO(), depKey, cmUpdated)
		return cmUpdated.Data["package.json"]
	}, timeout, 1*time.Second).Should(gomega.Equal(`dependencies`))

	ksvcUpdated := &servingv1alpha1.Service{}
	g.Expect(c.Get(context.TODO(), depKey, ksvcUpdated)).NotTo(gomega.HaveOccurred())
	g.Expect(ksvcUpdated.Spec.RunLatest.Configuration.RevisionTemplate.Spec.Container.Image).To(gomega.Equal(fmt.Sprintf("%s-%s:%s", "default", "foo", cm.GetObjectMeta().GetResourceVersion())
}

func TestReconcileErrors(t *testing.T) {
	g := gomega.NewGomegaWithT(t)
	fnCreated := &runtimev1alpha1.Function{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "errortest",
			Namespace: "default",
		},
		Spec: runtimev1alpha1.FunctionSpec{
			Function:            "main() {asdfasdf}",
			FunctionContentType: "plaintext",
			Size:                "L",
			Runtime:             "nodejs6",
		},
	}

	fnConfig := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "fn-config",
			Namespace: "default",
		},
		Data: map[string]string{
			"dockerRegistry":     "test",
			"serviceAccountName": "build-bot",
			"runtimes": `[
				{
					"ID": "nodejs8",
					"DockerFileName": "dockerfile-nodejs8",
				}
			]`,
		},
	}

	dockerFileConfigNodejs := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "dockerfile-nodejs8",
			Namespace: "default",
		},
		Data: map[string]string{
			"Dockerfile": `FROM kubeless/nodejs@sha256:5c3c21cf29231f25a0d7d2669c6f18c686894bf44e975fcbbbb420c6d045f7e7
				USER root
				RUN export KUBELESS_INSTALL_VOLUME='/kubeless' && \
					mkdir /kubeless && \
					cp /src/handler.js /kubeless && \
					cp /src/package.json /kubeless && \
					/kubeless-npm-install.sh
				USER 1000
			`,
		},
	}

	mgr, err := manager.New(cfg, manager.Options{})
	g.Expect(err).NotTo(gomega.HaveOccurred())
	c = mgr.GetClient()

	// Create the Function object and expect the Reconcile and Deployment to be created
	err = c.Create(context.TODO(), dockerFileConfigNodejs)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}

	err = c.Create(context.TODO(), fnConfig)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}

	err = c.Create(context.TODO(), fnCreated)
	t.Logf("%+v", err)
	g.Expect(err).NotTo(gomega.HaveOccurred())
	request := reconcile.Request{
		NamespacedName: types.NamespacedName{
			Name:      "bla",
			Namespace: "blabla",
		},
	}
	s := scheme.Scheme
	s.AddKnownTypes(runtimev1alpha1.SchemeGroupVersion, fnCreated)
	r := &ReconcileFunction{Client: c, scheme: s}

	os.Setenv("CONTROLLER_CONFIGMAP", "bla-config")
	_, err = r.Reconcile(request)
	g.Expect(err).To(gomega.HaveOccurred())
	g.Expect(err.Error()).To(gomega.Equal(`ConfigMap "bla-config" not found`))

	os.Unsetenv("CONTROLLER_CONFIGMAP")
	os.Setenv("CONTROLLER_CONFIGMAP_NS", "stage")
	_, err = r.Reconcile(request)
	g.Expect(err).To(gomega.HaveOccurred())
	g.Expect(err.Error()).To(gomega.Equal(`ConfigMap "fn-config" not found`))

}
<commit_msg>add service update test<commit_after>/*
Copyright 2019 The Kyma Authors.

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

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

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

package function

import (
	"encoding/json"
	"fmt"
	"os"
	"testing"
	"time"

	buildv1alpha1 "github.com/knative/build/pkg/apis/build/v1alpha1"
	servingv1alpha1 "github.com/knative/serving/pkg/apis/serving/v1alpha1"
	runtimev1alpha1 "github.com/kyma-incubator/runtime/pkg/apis/runtime/v1alpha1"
	"github.com/onsi/gomega"
	"golang.org/x/net/context"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/kubernetes/scheme"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"
)

var c client.Client

var expectedRequest = reconcile.Request{NamespacedName: types.NamespacedName{Name: "foo", Namespace: "default"}}
var depKey = types.NamespacedName{Name: "foo", Namespace: "default"}

const timeout = time.Second * 10

func TestReconcile(t *testing.T) {
	g := gomega.NewGomegaWithT(t)
	fnCreated := &runtimev1alpha1.Function{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "foo",
			Namespace: "default",
		},
		Spec: runtimev1alpha1.FunctionSpec{
			Function:            "main() {asdfasdf}",
			FunctionContentType: "plaintext",
			Size:                "L",
			Runtime:             "nodejs6",
		},
	}

	expectedEnv := []corev1.EnvVar{
		{
			Name:  "FUNC_HANDLER",
			Value: "main",
		},
		{
			Name:  "MOD_NAME",
			Value: "handler",
		},
		{
			Name:  "FUNC_TIMEOUT",
			Value: "180",
		},
		{
			Name:  "FUNC_RUNTIME",
			Value: "nodejs8",
		},
		{
			Name:  "FUNC_MEMORY_LIMIT",
			Value: "128Mi",
		},
		{
			Name:  "FUNC_PORT",
			Value: "8080",
		},
		{
			Name:  "NODE_PATH",
			Value: "$(KUBELESS_INSTALL_VOLUME)/node_modules",
		},
	}

	fnConfig := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "fn-config",
			Namespace: "default",
		},
		Data: map[string]string{
			"dockerRegistry":     "test",
			"serviceAccountName": "build-bot",
			"runtimes": `[
				{
					"ID": "nodejs8",
					"DockerFileName": "dockerfile-nodejs8",
				}
			]`,
		},
	}

	dockerFileConfigNodejs := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "dockerfile-nodejs8",
			Namespace: "default",
		},
		Data: map[string]string{
			"Dockerfile": `FROM kubeless/nodejs@sha256:5c3c21cf29231f25a0d7d2669c6f18c686894bf44e975fcbbbb420c6d045f7e7
				USER root
				RUN export KUBELESS_INSTALL_VOLUME='/kubeless' && \
					mkdir /kubeless && \
					cp /src/handler.js /kubeless && \
					cp /src/package.json /kubeless && \
					/kubeless-npm-install.sh
				USER 1000
			`,
		},
	}
	mgr, err := manager.New(cfg, manager.Options{})
	g.Expect(err).NotTo(gomega.HaveOccurred())
	c = mgr.GetClient()

	recFn, requests := SetupTestReconcile(newReconciler(mgr))
	g.Expect(add(mgr, recFn)).NotTo(gomega.HaveOccurred())

	stopMgr, mgrStopped := StartTestManager(mgr, g)

	defer func() {
		close(stopMgr)
		mgrStopped.Wait()
	}()

	// Create the Function object and expect the Reconcile and Deployment to be created
	err = c.Create(context.TODO(), dockerFileConfigNodejs)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}

	err = c.Create(context.TODO(), fnConfig)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}

	err = c.Create(context.TODO(), fnCreated)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}
	g.Expect(err).NotTo(gomega.HaveOccurred())

	g.Eventually(requests, timeout).Should(gomega.Receive(gomega.Equal(expectedRequest)))

	cm := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
	}

	service := &servingv1alpha1.Service{
		ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "default"},
	}

	g.Eventually(func() error { return c.Get(context.TODO(), depKey, cm) }, timeout).
		Should(gomega.Succeed())

	g.Eventually(func() error { return c.Get(context.TODO(), depKey, service) }, timeout).
		Should(gomega.Succeed())
	g.Expect(service.Namespace).To(gomega.Equal("default"))

	g.Expect(service.Spec.RunLatest.Configuration.RevisionTemplate.Spec.Container.Env).To(gomega.Equal(expectedEnv))
	build := (*service.Spec.RunLatest.Configuration.Build)
	buildByte, err := build.MarshalJSON()
	if err != nil {
		t.Fatalf("Error while marshaling build object: %v", err)
	}
	var buildSpec buildv1alpha1.BuildSpec
	err = json.Unmarshal(buildByte, &buildSpec)
	if err != nil {
		t.Fatalf("Error while unmarshaling buildSpec: %v", err)
	}
	g.Expect(service.Spec.RunLatest.Configuration.RevisionTemplate.Spec.Container.Image).To(gomega.HavePrefix("test/default-foo"))
	g.Expect(len(buildSpec.Volumes)).To(gomega.Equal(2))
	g.Expect(buildSpec.ServiceAccountName).To(gomega.Equal("build-bot"))
	g.Expect(service.Spec.RunLatest.Configuration.RevisionTemplate.Spec.Container.Image).To(gomega.HavePrefix("test/default-foo"))

	fnFetched := &runtimev1alpha1.Function{}
	g.Expect(c.Get(context.TODO(), depKey, fnFetched)).NotTo(gomega.HaveOccurred())
	g.Expect(fnFetched.Spec).To(gomega.Equal(fnCreated.Spec))

	fnUpdated := fnFetched.DeepCopy()
	fnUpdated.Spec.Function = `main() {return "bla"}`
	fnUpdated.Spec.Deps = `dependencies`

	fnFetched = &runtimev1alpha1.Function{}
	g.Expect(c.Update(context.TODO(), fnUpdated)).NotTo(gomega.HaveOccurred())
	g.Expect(c.Get(context.TODO(), depKey, fnFetched)).NotTo(gomega.HaveOccurred())
	g.Expect(fnFetched.Spec).To(gomega.Equal(fnUpdated.Spec))

	g.Eventually(requests, timeout).Should(gomega.Receive(gomega.Equal(expectedRequest)))

	cmUpdated := &corev1.ConfigMap{}
	g.Eventually(func() string {
		c.Get(context.TODO(), depKey, cmUpdated)
		return cmUpdated.Data["handler.js"]
	}, timeout, 1*time.Second).Should(gomega.Equal(fnUpdated.Spec.Function))
	g.Eventually(func() string {
		c.Get(context.TODO(), depKey, cmUpdated)
		return cmUpdated.Data["package.json"]
	}, timeout, 1*time.Second).Should(gomega.Equal(`dependencies`))

	ksvcUpdated := &servingv1alpha1.Service{}
	g.Expect(c.Get(context.TODO(), depKey, ksvcUpdated)).NotTo(gomega.HaveOccurred())
	g.Expect(ksvcUpdated.Spec.RunLatest.Configuration.RevisionTemplate.Spec.Container.Image).
		To(gomega.Equal(fmt.Sprintf("test/%s-%s:%s", "default", "foo", cmUpdated.GetObjectMeta().GetResourceVersion())))
}

func TestReconcileErrors(t *testing.T) {
	g := gomega.NewGomegaWithT(t)
	fnCreated := &runtimev1alpha1.Function{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "errortest",
			Namespace: "default",
		},
		Spec: runtimev1alpha1.FunctionSpec{
			Function:            "main() {asdfasdf}",
			FunctionContentType: "plaintext",
			Size:                "L",
			Runtime:             "nodejs6",
		},
	}

	fnConfig := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "fn-config",
			Namespace: "default",
		},
		Data: map[string]string{
			"dockerRegistry":     "test",
			"serviceAccountName": "build-bot",
			"runtimes": `[
				{
					"ID": "nodejs8",
					"DockerFileName": "dockerfile-nodejs8",
				}
			]`,
		},
	}

	dockerFileConfigNodejs := &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "dockerfile-nodejs8",
			Namespace: "default",
		},
		Data: map[string]string{
			"Dockerfile": `FROM kubeless/nodejs@sha256:5c3c21cf29231f25a0d7d2669c6f18c686894bf44e975fcbbbb420c6d045f7e7
				USER root
				RUN export KUBELESS_INSTALL_VOLUME='/kubeless' && \
					mkdir /kubeless && \
					cp /src/handler.js /kubeless && \
					cp /src/package.json /kubeless && \
					/kubeless-npm-install.sh
				USER 1000
			`,
		},
	}

	mgr, err := manager.New(cfg, manager.Options{})
	g.Expect(err).NotTo(gomega.HaveOccurred())
	c = mgr.GetClient()

	// Create the Function object and expect the Reconcile and Deployment to be created
	err = c.Create(context.TODO(), dockerFileConfigNodejs)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}

	err = c.Create(context.TODO(), fnConfig)
	if apierrors.IsInvalid(err) {
		t.Logf("failed to create object, got an invalid object error: %v", err)
		return
	}

	err = c.Create(context.TODO(), fnCreated)
	t.Logf("%+v", err)
	g.Expect(err).NotTo(gomega.HaveOccurred())
	request := reconcile.Request{
		NamespacedName: types.NamespacedName{
			Name:      "bla",
			Namespace: "blabla",
		},
	}
	s := scheme.Scheme
	s.AddKnownTypes(runtimev1alpha1.SchemeGroupVersion, fnCreated)
	r := &ReconcileFunction{Client: c, scheme: s}

	os.Setenv("CONTROLLER_CONFIGMAP", "bla-config")
	_, err = r.Reconcile(request)
	g.Expect(err).To(gomega.HaveOccurred())
	g.Expect(err.Error()).To(gomega.Equal(`ConfigMap "bla-config" not found`))

	os.Unsetenv("CONTROLLER_CONFIGMAP")
	os.Setenv("CONTROLLER_CONFIGMAP_NS", "stage")
	_, err = r.Reconcile(request)
	g.Expect(err).To(gomega.HaveOccurred())
	g.Expect(err.Error()).To(gomega.Equal(`ConfigMap "fn-config" not found`))

}
<|endoftext|>


<commit_before>package azure

import (
	"bytes"
	"context"
	"errors"
	"net/http"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"

	"github.com/grafana/loki/pkg/storage/chunk/hedging"
)

type RoundTripperFunc func(*http.Request) (*http.Response, error)

func (fn RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return fn(req)
}

func Test_Hedging(t *testing.T) {
	metrics := NewBlobStorageMetrics()
	for _, tc := range []struct {
		name          string
		expectedCalls int32
		hedgeAt       time.Duration
		upTo          int
		do            func(c *BlobStorage)
	}{
		{
			"delete/put/list are not hedged",
			4, // Put makes an additional call to close the stream
			20 * time.Nanosecond,
			10,
			func(c *BlobStorage) {
				_ = c.DeleteObject(context.Background(), "foo")
				_, _, _ = c.List(context.Background(), "foo", "/")
				_ = c.PutObject(context.Background(), "foo", bytes.NewReader([]byte("bar")))
			},
		},
		{
			"gets are hedged",
			3,
			20 * time.Nanosecond,
			3,
			func(c *BlobStorage) {
				_, _, _ = c.GetObject(context.Background(), "foo")
			},
		},
		{
			"gets are not hedged when not configured",
			1,
			0,
			0,
			func(c *BlobStorage) {
				_, _, _ = c.GetObject(context.Background(), "foo")
			},
		},
	} {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			count := atomic.NewInt32(0)
			// hijack the client to count the number of calls
			defaultClientFactory = func() *http.Client {
				return &http.Client{
					Transport: RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
						count.Inc()
						time.Sleep(200 * time.Millisecond)
						return nil, errors.New("fo")
					}),
				}
			}

			c, err := NewBlobStorage(
				&BlobStorageConfig{
					AccountName: "account",
					Environment: azureGlobal,
					MaxRetries:  0,
				},
				metrics,
				hedging.Config{
					At:           tc.hedgeAt,
					UpTo:         tc.upTo,
					MaxPerSecond: 1000,
				})
			require.NoError(t, err)

			tc.do(c)
			require.Eventually(t, func() bool { return tc.expectedCalls == count.Load() }, time.Second, time.Millisecond)
		})
	}
}
<commit_msg>Fixes a flaky hedging test. (#5389)<commit_after>package azure

import (
	"bytes"
	"context"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"

	"github.com/grafana/loki/pkg/storage/chunk/hedging"
)

var metrics = NewBlobStorageMetrics()

type RoundTripperFunc func(*http.Request) (*http.Response, error)

func (fn RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return fn(req)
}

func Test_Hedging(t *testing.T) {
	for _, tc := range []struct {
		name          string
		expectedCalls int32
		hedgeAt       time.Duration
		upTo          int
		do            func(c *BlobStorage)
	}{
		{
			"delete/put/list are not hedged",
			3,
			20 * time.Nanosecond,
			10,
			func(c *BlobStorage) {
				_ = c.DeleteObject(context.Background(), "foo")
				_, _, _ = c.List(context.Background(), "foo", "/")
				_ = c.PutObject(context.Background(), "foo", bytes.NewReader([]byte("bar")))
			},
		},
		{
			"gets are hedged",
			3,
			20 * time.Nanosecond,
			3,
			func(c *BlobStorage) {
				_, _, _ = c.GetObject(context.Background(), "foo")
			},
		},
		{
			"gets are not hedged when not configured",
			1,
			0,
			0,
			func(c *BlobStorage) {
				_, _, _ = c.GetObject(context.Background(), "foo")
			},
		},
	} {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			count := atomic.NewInt32(0)
			// hijack the client to count the number of calls
			defaultClientFactory = func() *http.Client {
				return &http.Client{
					Transport: RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
						// blocklist is a call that can be fired by the SDK after PUT but is not guaranteed.
						if !strings.Contains(req.URL.String(), "blocklist") {
							count.Inc()
							time.Sleep(50 * time.Millisecond)
						}
						return nil, http.ErrNotSupported
					}),
				}
			}

			c, err := NewBlobStorage(
				&BlobStorageConfig{
					AccountName: "account",
					Environment: azureGlobal,
					MaxRetries:  0,
				},
				metrics,
				hedging.Config{
					At:           tc.hedgeAt,
					UpTo:         tc.upTo,
					MaxPerSecond: 1000,
				})
			require.NoError(t, err)

			tc.do(c)
			require.Eventually(t, func() bool { return tc.expectedCalls == count.Load() }, time.Second, time.Millisecond, "expected calls %d, got %d", tc.expectedCalls, count.Load())
		})
	}
}
<|endoftext|>


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

package gaesessions

import (
	"bytes"
	"encoding/base32"
	"encoding/gob"
	"net/http"
	"strings"
	"time"

	"appengine"
	"appengine/datastore"
	"appengine/memcache"

	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
)

// DatastoreStore -------------------------------------------------------------

// Session is used to load and save session data in the datastore.
type Session struct {
	Date  time.Time
	Value []byte
}

// NewDatastoreStore returns a new DatastoreStore.
//
// The kind argument is the kind name used to store the session data.
// If empty it will use "Session".
//
// See NewCookieStore() for a description of the other parameters.
func NewDatastoreStore(kind string, keyPairs ...[]byte) *DatastoreStore {
	if kind == "" {
		kind = "Session"
	}
	return &DatastoreStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		kind: kind,
	}
}

// DatastoreStore stores sessions in the App Engine datastore.
type DatastoreStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options // default configuration
	kind    string
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *DatastoreStore) Get(r *http.Request, name string) (*sessions.Session,
	error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *DatastoreStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if c, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, c.Value, &session.ID, s.Codecs...)
		if err == nil {
			err = s.load(r, session)
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *DatastoreStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID =
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	if err := s.save(r, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// save writes encoded session.Values to datastore.
func (s *DatastoreStore) save(r *http.Request,
	session *sessions.Session) error {
	if len(session.Values) == 0 {
		// Don't need to write anything.
		return nil
	}
	serialized, err := serialize(session.Values)
	if err != nil {
		return err
	}
	c := appengine.NewContext(r)
	k := datastore.NewKey(c, s.kind, session.ID, 0, nil)
	k, err = datastore.Put(c, k, &Session{
		Date:  time.Now(),
		Value: serialized,
	})
	if err != nil {
		return err
	}
	return nil
}

// load gets a value from datastore and decodes its content into
// session.Values.
func (s *DatastoreStore) load(r *http.Request,
	session *sessions.Session) error {
	c := appengine.NewContext(r)
	k := datastore.NewKey(c, s.kind, session.ID, 0, nil)
	entity := Session{}
	if err := datastore.Get(c, k, &entity); err != nil {
		return err
	}
	if err := deserialize(entity.Value, &session.Values); err != nil {
		return err
	}
	return nil
}

// MemcacheStore --------------------------------------------------------------

// NewMemcacheStore returns a new MemcacheStore.
//
// The keyPrefix argument is the prefix used for memcache keys. If empty it
// will use "gorilla.appengine.sessions.".
//
// See NewCookieStore() for a description of the other parameters.
func NewMemcacheStore(keyPrefix string, keyPairs ...[]byte) *MemcacheStore {
	if keyPrefix == "" {
		keyPrefix = "gorilla.appengine.sessions."
	}
	return &MemcacheStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		prefix: keyPrefix,
	}
}

// MemcacheStore stores sessions in the App Engine memcache.
type MemcacheStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options // default configuration
	prefix  string
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *MemcacheStore) Get(r *http.Request, name string) (*sessions.Session,
	error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *MemcacheStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if c, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, c.Value, &session.ID, s.Codecs...)
		if err == nil {
			err = s.load(r, session)
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *MemcacheStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID = s.prefix +
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	if err := s.save(r, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// save writes encoded session.Values to memcache.
func (s *MemcacheStore) save(r *http.Request,
	session *sessions.Session) error {
	if len(session.Values) == 0 {
		// Don't need to write anything.
		return nil
	}
	serialized, err := serialize(session.Values)
	if err != nil {
		return err
	}
	c := appengine.NewContext(r)
	var expiration time.Duration
	if session.Options.MaxAge > 0 {
		expiration = time.Duration(session.Options.MaxAge) * time.Second
	}
	c.Debugf("MemcacheStore.Save. session.ID=%s, expiration=%s",
		session.ID, expiration)
	err = memcache.Set(c, &memcache.Item{
		Key:        session.ID,
		Value:      serialized,
		Expiration: expiration,
	})
	if err != nil {
		return err
	}
	return nil
}

// load gets a value from memcache and decodes its content into session.Values.
func (s *MemcacheStore) load(r *http.Request,
	session *sessions.Session) error {
	item, err := memcache.Get(appengine.NewContext(r), session.ID)
	if err != nil {
		return err
	}
	if err := deserialize(item.Value, &session.Values); err != nil {
		return err
	}
	return nil
}

// Serialization --------------------------------------------------------------

// serialize encodes a value using gob.
func serialize(src interface{}) ([]byte, error) {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	if err := enc.Encode(src); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// deserialize decodes a value using gob.
func deserialize(src []byte, dst interface{}) error {
	dec := gob.NewDecoder(bytes.NewBuffer(src))
	if err := dec.Decode(dst); err != nil {
		return err
	}
	return nil
}
<commit_msg>Use taskqueue to delete expired session in datastore when MaxAge > 0.<commit_after>// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gaesessions

import (
	"bytes"
	"encoding/base32"
	"encoding/gob"
	"net/http"
	"strings"
	"time"

	"appengine"
	"appengine/datastore"
	"appengine/delay"
	"appengine/memcache"
	"appengine/taskqueue"

	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
)

// DatastoreStore -------------------------------------------------------------

// Session is used to load and save session data in the datastore.
type Session struct {
	Date       time.Time
	Expiration time.Duration
	Value      []byte
}

// NewDatastoreStore returns a new DatastoreStore.
//
// The kind argument is the kind name used to store the session data.
// If empty it will use "Session".
//
// See NewCookieStore() for a description of the other parameters.
func NewDatastoreStore(kind string, keyPairs ...[]byte) *DatastoreStore {
	if kind == "" {
		kind = "Session"
	}
	return &DatastoreStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		kind: kind,
	}
}

// DatastoreStore stores sessions in the App Engine datastore.
type DatastoreStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options // default configuration
	kind    string
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *DatastoreStore) Get(r *http.Request, name string) (*sessions.Session,
	error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *DatastoreStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if c, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, c.Value, &session.ID, s.Codecs...)
		if err == nil {
			err = s.load(r, session)
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *DatastoreStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID =
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	if err := s.save(r, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// save writes encoded session.Values to datastore.
func (s *DatastoreStore) save(r *http.Request,
	session *sessions.Session) error {
	if len(session.Values) == 0 {
		// Don't need to write anything.
		return nil
	}
	serialized, err := serialize(session.Values)
	if err != nil {
		return err
	}
	c := appengine.NewContext(r)
	k := datastore.NewKey(c, s.kind, session.ID, 0, nil)
	var expiration time.Duration
	if session.Options.MaxAge > 0 {
		expiration = time.Duration(session.Options.MaxAge) * time.Second
	}
	k, err = datastore.Put(c, k, &Session{
		Date:       time.Now(),
		Expiration: expiration,
		Value:      serialized,
	})
	if err != nil {
		return err
	}
	if session.Options.MaxAge > 0 {
		taskKey := strings.TrimRight(
			base32.StdEncoding.EncodeToString(
				securecookie.GenerateRandomKey(32)), "=")
		laterFunc := delay.Func(taskKey, expireFunc)
		task, err := laterFunc.Task(s.kind, session.ID)
		if err != nil {
			return err
		}
		task.Delay = expiration
		task, err = taskqueue.Add(c, task, "")
		if err != nil {
			return err
		}
	}
	return nil
}

// load gets a value from datastore and decodes its content into
// session.Values.
func (s *DatastoreStore) load(r *http.Request,
	session *sessions.Session) error {
	c := appengine.NewContext(r)
	k := datastore.NewKey(c, s.kind, session.ID, 0, nil)
	entity := Session{}
	if err := datastore.Get(c, k, &entity); err != nil {
		return err
	}
	if err := deserialize(entity.Value, &session.Values); err != nil {
		return err
	}
	return nil
}

func expireFunc(c appengine.Context, kind, sessionID string) error {
	c.Debugf("DatastoreStore expireFunc start session.ID=%s", sessionID)
	k := datastore.NewKey(c, kind, sessionID, 0, nil)
	entity := Session{}
	if err := datastore.Get(c, k, &entity); err != nil {
		c.Errorf("DatastoreStore expireFunc datastore.Get failed. session.ID=%s, err=%s", sessionID, err.Error())
		return err
	}
	session := sessions.Session{
		Values: make(map[interface{}]interface{}),
	}
	if err := deserialize(entity.Value, &session.Values); err != nil {
		c.Errorf("DatastoreStore expireFunc deserialize failed. session.ID=%s, err=%s", sessionID, err.Error())
		return err
	}
	now := time.Now()
	expirationTime := entity.Date.Add(entity.Expiration)
	if now.After(expirationTime) {
		err := datastore.Delete(c, k)
		if err != nil {
			c.Errorf("DatastoreStore expireFunc Delete session.ID=%s, now=%s, expirationTime=%s, err=%s", sessionID, now, expirationTime, err.Error())
			return err
		}
		c.Debugf("DatastoreStore expireFunc Delete done. session.ID=%s", sessionID)
	}
	return nil
}

// MemcacheStore --------------------------------------------------------------

// NewMemcacheStore returns a new MemcacheStore.
//
// The keyPrefix argument is the prefix used for memcache keys. If empty it
// will use "gorilla.appengine.sessions.".
//
// See NewCookieStore() for a description of the other parameters.
func NewMemcacheStore(keyPrefix string, keyPairs ...[]byte) *MemcacheStore {
	if keyPrefix == "" {
		keyPrefix = "gorilla.appengine.sessions."
	}
	return &MemcacheStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		prefix: keyPrefix,
	}
}

// MemcacheStore stores sessions in the App Engine memcache.
type MemcacheStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options // default configuration
	prefix  string
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *MemcacheStore) Get(r *http.Request, name string) (*sessions.Session,
	error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *MemcacheStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if c, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, c.Value, &session.ID, s.Codecs...)
		if err == nil {
			err = s.load(r, session)
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *MemcacheStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID = s.prefix +
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	if err := s.save(r, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// save writes encoded session.Values to memcache.
func (s *MemcacheStore) save(r *http.Request,
	session *sessions.Session) error {
	if len(session.Values) == 0 {
		// Don't need to write anything.
		return nil
	}
	serialized, err := serialize(session.Values)
	if err != nil {
		return err
	}
	c := appengine.NewContext(r)
	var expiration time.Duration
	if session.Options.MaxAge > 0 {
		expiration = time.Duration(session.Options.MaxAge) * time.Second
	}
	c.Debugf("MemcacheStore.Save. session.ID=%s, expiration=%s",
		session.ID, expiration)
	err = memcache.Set(c, &memcache.Item{
		Key:        session.ID,
		Value:      serialized,
		Expiration: expiration,
	})
	if err != nil {
		return err
	}
	return nil
}

// load gets a value from memcache and decodes its content into session.Values.
func (s *MemcacheStore) load(r *http.Request,
	session *sessions.Session) error {
	item, err := memcache.Get(appengine.NewContext(r), session.ID)
	if err != nil {
		return err
	}
	if err := deserialize(item.Value, &session.Values); err != nil {
		return err
	}
	return nil
}

// Serialization --------------------------------------------------------------

// serialize encodes a value using gob.
func serialize(src interface{}) ([]byte, error) {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	if err := enc.Encode(src); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// deserialize decodes a value using gob.
func deserialize(src []byte, dst interface{}) error {
	dec := gob.NewDecoder(bytes.NewBuffer(src))
	if err := dec.Decode(dst); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package goriak

import (
	"reflect"
	"sort"
	"testing"
)

func TestAutoMapSet(t *testing.T) {
	type ourTestType struct {
		Items *Set
	}

	testVal := ourTestType{}
	result, errset := bucket().Set(testVal).Run(con())

	if errset != nil {
		t.Error("Set:", errset)
	}

	// Get when empty
	var res ourTestType
	_, errget := bucket().Get(result.Key, &res).Run(con())

	if errget != nil {
		t.Error("Get:", errget)
	}

	res.Items.AddString("aaa")
	res.Items.AddString("bbb")
	res.Items.AddString("ccc")

	s := res.Items.Strings()

	if len(s) != 3 {
		t.Error("Unexpected length")
	}

	expected := []string{"aaa", "bbb", "ccc"}

	if !reflect.DeepEqual(expected, res.Items.Strings()) {
		t.Log(expected)
		t.Log(res.Items.Strings())
		t.Error("Not deep equal first")
	}

	err := res.Items.Exec(con())

	if err != nil {
		t.Error("Exec1: ", err)
	}

	// Get after save
	var res2 ourTestType
	_, errget = bucket().Get(result.Key, &res2).Run(con())

	if errget != nil {
		t.Error("Get:", errget)
	}

	if !reflect.DeepEqual(res2.Items.Value(), res.Items.Value()) {
		t.Error("Not deep equal after get")
	}

	// Remove
	err = res2.Items.RemoveString("bbb").Exec(con())

	if err != nil {
		t.Error("Exec2: ", err)
	}

	// Get after remove
	var res3 ourTestType
	_, errget = bucket().Get(result.Key, &res3).Run(con())

	if errget != nil {
		t.Error("Get:", errget)
	}

	if !reflect.DeepEqual(res3.Items.Value(), res2.Items.Value()) {
		t.Log(res3.Items.Strings())
		t.Log(res2.Items.Strings())
		t.Error("Not deep equal after get 2")
	}
}

func TestAutoMapSetAddRemove(t *testing.T) {
	set := NewSet()

	set.AddString("1")
	set.AddString("2")
	set.AddString("3")

	expected := []string{"1", "2", "3"}

	if !reflect.DeepEqual(expected, set.Strings()) {
		t.Log(expected)
		t.Log(set.Strings())
		t.Error("Unexpected 1")
	}

	set.RemoveString("2")

	expected = []string{"1", "3"}

	if !reflect.DeepEqual(expected, set.Strings()) {
		t.Log(expected)
		t.Log(set.Strings())
		t.Error("Unexpected 2")
	}

	expectedAdds := [][]byte{[]byte("1"), []byte("3")}

	if !reflect.DeepEqual(expectedAdds, set.adds) {
		t.Log(expectedAdds)
		t.Log(set.adds)
		t.Error("Unexpected adds 1")
	}

	expectedRemoves := [][]byte{[]byte("2")}

	if !reflect.DeepEqual(expectedRemoves, set.removes) {
		t.Log(expectedRemoves)
		t.Log(set.removes)
		t.Error("Unexpected removes 1")
	}

	set.RemoveString("4")

	expectedRemoves = [][]byte{[]byte("2"), []byte("4")}

	if !reflect.DeepEqual(expectedRemoves, set.removes) {
		t.Log(expectedRemoves)
		t.Log(set.removes)
		t.Error("Unexpected removes 2")
	}

	set.AddString("4")

	expectedAdds = [][]byte{[]byte("1"), []byte("3"), []byte("4")}

	if !reflect.DeepEqual(expectedAdds, set.adds) {
		t.Log(expectedAdds)
		t.Log(set.adds)
		t.Error("Unexpected adds 2")
	}

	expectedRemoves = [][]byte{[]byte("2")}

	if !reflect.DeepEqual(expectedRemoves, set.removes) {
		t.Log(expectedRemoves)
		t.Log(set.removes)
		t.Error("Unexpected removes 3")
	}
}

func TestAutoMapMultipleSet(t *testing.T) {
	set := NewSet()

	set.AddString("hello")
	set.AddString("hello")
	set.AddString("hello")
	set.AddString("hello")

	expectedVal := []string{"hello"}
	expectedAdds := [][]byte{[]byte("hello")}

	if !reflect.DeepEqual(expectedVal, set.Strings()) {
		t.Error("Unexpected value")
	}

	if !reflect.DeepEqual(expectedAdds, set.adds) {
		t.Error("Unexpected adds")
	}
}

func TestAutoMapSetAddRemoveSetMap(t *testing.T) {

	type ourTestType struct {
		Tags *Set

		Context []byte `goriak:"goriakcontext"`
	}

	testVal := ourTestType{
		Tags: NewSet(),
	}

	testVal.Tags.AddString("one")
	testVal.Tags.AddString("two")
	testVal.Tags.AddString("three")
	testVal.Tags.AddString("four")

	result, errset := bucket().Set(testVal).Run(con())

	if errset != nil {
		t.Error("Set 1: ", errset)
	}

	// Get it back
	var resVal ourTestType
	_, errget := bucket().Get(result.Key, &resVal).Run(con())

	if errget != nil {
		t.Error("Get: ", errget)
	}

	expected := []string{"one", "two", "three", "four"}
	sort.Strings(expected)

	val := resVal.Tags.Strings()
	sort.Strings(val)

	if !reflect.DeepEqual(expected, val) {
		t.Logf("Expected: %+v\n", expected)
		t.Logf("Got: %+v\n", val)
		t.Error("Unexpected value 1")
	}

	// Remove from fetched
	resVal.Tags.RemoveString("two")

	if errset != nil {
		t.Error("Set 2: ", resVal)
	}

	_, errset = bucket().Key(result.Key).Set(resVal).Run(con())

	if errset != nil {
		t.Error("Set 2: ", errset)
	}

	// Get it back
	var resVal2 ourTestType
	_, errget = bucket().Get(result.Key, &resVal2).Run(con())

	if errget != nil {
		t.Error("Get 2: ", errget)
	}

	expected = []string{"one", "three", "four"}
	sort.Strings(expected)

	val2 := resVal2.Tags.Strings()
	sort.Strings(val2)

	if !reflect.DeepEqual(expected, val2) {
		t.Logf("Expected: %+v\n", expected)
		t.Logf("Got: %+v\n", val2)
		t.Error("Unexpected value 2")
	}

}

func ExampleSet() {

	session, _ := Connect(ConnectOpts{
		Address: "127.0.0.1",
	})

	type Article struct {
		Tags *Set

		Context []byte `goriak:"goriakcontext"`
	}

	// Initializing a new Article and the Set within
	art := Article{
		Tags: NewSet(),
	}

	riakKey := "article-1"

	// Adding the tags "one" and "two"
	art.Tags.AddString("one")
	art.Tags.AddString("two")

	_, err := Bucket("bucket", "bucketType").Key(riakKey).Set(art).Run(session)

	if err != nil {
		// ..
	}

	// Retreiving from Riak
	var getArt Article
	_, err = Bucket("bucket", "bucketType").Get(riakKey, &getArt).Run(session)

	if err != nil {
		// ..
	}

	// Adding one extra tag.
	// Multiple AddString() and RemoveString() can be chained together before calling Exec().
	err = getArt.Tags.AddString("three").Exec(session)

	if err != nil {
		// ..
	}
}

func TestSetHas(t *testing.T) {
	s := NewSet()
	s.AddString("hello")
	s.AddString("it's")
	s.AddString("me")

	if !s.HasString("hello") {
		t.Error("Did not have hello")
	}

	if !s.HasString("it's") {
		t.Error("Did not have it's")
	}

	if s.HasString("i") {
		t.Error("Had i")
	}
}
<commit_msg>#26 - Added test for initializing Sets when using Set()<commit_after>package goriak

import (
	"reflect"
	"sort"
	"testing"
)

func TestAutoMapSet(t *testing.T) {
	type ourTestType struct {
		Items *Set
	}

	testVal := ourTestType{}
	result, errset := bucket().Set(testVal).Run(con())

	if errset != nil {
		t.Error("Set:", errset)
	}

	// Get when empty
	var res ourTestType
	_, errget := bucket().Get(result.Key, &res).Run(con())

	if errget != nil {
		t.Error("Get:", errget)
	}

	res.Items.AddString("aaa")
	res.Items.AddString("bbb")
	res.Items.AddString("ccc")

	s := res.Items.Strings()

	if len(s) != 3 {
		t.Error("Unexpected length")
	}

	expected := []string{"aaa", "bbb", "ccc"}

	if !reflect.DeepEqual(expected, res.Items.Strings()) {
		t.Log(expected)
		t.Log(res.Items.Strings())
		t.Error("Not deep equal first")
	}

	err := res.Items.Exec(con())

	if err != nil {
		t.Error("Exec1: ", err)
	}

	// Get after save
	var res2 ourTestType
	_, errget = bucket().Get(result.Key, &res2).Run(con())

	if errget != nil {
		t.Error("Get:", errget)
	}

	if !reflect.DeepEqual(res2.Items.Value(), res.Items.Value()) {
		t.Error("Not deep equal after get")
	}

	// Remove
	err = res2.Items.RemoveString("bbb").Exec(con())

	if err != nil {
		t.Error("Exec2: ", err)
	}

	// Get after remove
	var res3 ourTestType
	_, errget = bucket().Get(result.Key, &res3).Run(con())

	if errget != nil {
		t.Error("Get:", errget)
	}

	if !reflect.DeepEqual(res3.Items.Value(), res2.Items.Value()) {
		t.Log(res3.Items.Strings())
		t.Log(res2.Items.Strings())
		t.Error("Not deep equal after get 2")
	}
}

func TestAutoMapSetAddRemove(t *testing.T) {
	set := NewSet()

	set.AddString("1")
	set.AddString("2")
	set.AddString("3")

	expected := []string{"1", "2", "3"}

	if !reflect.DeepEqual(expected, set.Strings()) {
		t.Log(expected)
		t.Log(set.Strings())
		t.Error("Unexpected 1")
	}

	set.RemoveString("2")

	expected = []string{"1", "3"}

	if !reflect.DeepEqual(expected, set.Strings()) {
		t.Log(expected)
		t.Log(set.Strings())
		t.Error("Unexpected 2")
	}

	expectedAdds := [][]byte{[]byte("1"), []byte("3")}

	if !reflect.DeepEqual(expectedAdds, set.adds) {
		t.Log(expectedAdds)
		t.Log(set.adds)
		t.Error("Unexpected adds 1")
	}

	expectedRemoves := [][]byte{[]byte("2")}

	if !reflect.DeepEqual(expectedRemoves, set.removes) {
		t.Log(expectedRemoves)
		t.Log(set.removes)
		t.Error("Unexpected removes 1")
	}

	set.RemoveString("4")

	expectedRemoves = [][]byte{[]byte("2"), []byte("4")}

	if !reflect.DeepEqual(expectedRemoves, set.removes) {
		t.Log(expectedRemoves)
		t.Log(set.removes)
		t.Error("Unexpected removes 2")
	}

	set.AddString("4")

	expectedAdds = [][]byte{[]byte("1"), []byte("3"), []byte("4")}

	if !reflect.DeepEqual(expectedAdds, set.adds) {
		t.Log(expectedAdds)
		t.Log(set.adds)
		t.Error("Unexpected adds 2")
	}

	expectedRemoves = [][]byte{[]byte("2")}

	if !reflect.DeepEqual(expectedRemoves, set.removes) {
		t.Log(expectedRemoves)
		t.Log(set.removes)
		t.Error("Unexpected removes 3")
	}
}

func TestAutoMapMultipleSet(t *testing.T) {
	set := NewSet()

	set.AddString("hello")
	set.AddString("hello")
	set.AddString("hello")
	set.AddString("hello")

	expectedVal := []string{"hello"}
	expectedAdds := [][]byte{[]byte("hello")}

	if !reflect.DeepEqual(expectedVal, set.Strings()) {
		t.Error("Unexpected value")
	}

	if !reflect.DeepEqual(expectedAdds, set.adds) {
		t.Error("Unexpected adds")
	}
}

func TestAutoMapSetAddRemoveSetMap(t *testing.T) {

	type ourTestType struct {
		Tags *Set

		Context []byte `goriak:"goriakcontext"`
	}

	testVal := ourTestType{
		Tags: NewSet(),
	}

	testVal.Tags.AddString("one")
	testVal.Tags.AddString("two")
	testVal.Tags.AddString("three")
	testVal.Tags.AddString("four")

	result, errset := bucket().Set(testVal).Run(con())

	if errset != nil {
		t.Error("Set 1: ", errset)
	}

	// Get it back
	var resVal ourTestType
	_, errget := bucket().Get(result.Key, &resVal).Run(con())

	if errget != nil {
		t.Error("Get: ", errget)
	}

	expected := []string{"one", "two", "three", "four"}
	sort.Strings(expected)

	val := resVal.Tags.Strings()
	sort.Strings(val)

	if !reflect.DeepEqual(expected, val) {
		t.Logf("Expected: %+v\n", expected)
		t.Logf("Got: %+v\n", val)
		t.Error("Unexpected value 1")
	}

	// Remove from fetched
	resVal.Tags.RemoveString("two")

	if errset != nil {
		t.Error("Set 2: ", resVal)
	}

	_, errset = bucket().Key(result.Key).Set(resVal).Run(con())

	if errset != nil {
		t.Error("Set 2: ", errset)
	}

	// Get it back
	var resVal2 ourTestType
	_, errget = bucket().Get(result.Key, &resVal2).Run(con())

	if errget != nil {
		t.Error("Get 2: ", errget)
	}

	expected = []string{"one", "three", "four"}
	sort.Strings(expected)

	val2 := resVal2.Tags.Strings()
	sort.Strings(val2)

	if !reflect.DeepEqual(expected, val2) {
		t.Logf("Expected: %+v\n", expected)
		t.Logf("Got: %+v\n", val2)
		t.Error("Unexpected value 2")
	}

}

func ExampleSet() {

	session, _ := Connect(ConnectOpts{
		Address: "127.0.0.1",
	})

	type Article struct {
		Tags *Set

		Context []byte `goriak:"goriakcontext"`
	}

	// Initializing a new Article and the Set within
	art := Article{
		Tags: NewSet(),
	}

	riakKey := "article-1"

	// Adding the tags "one" and "two"
	art.Tags.AddString("one")
	art.Tags.AddString("two")

	_, err := Bucket("bucket", "bucketType").Key(riakKey).Set(art).Run(session)

	if err != nil {
		// ..
	}

	// Retreiving from Riak
	var getArt Article
	_, err = Bucket("bucket", "bucketType").Get(riakKey, &getArt).Run(session)

	if err != nil {
		// ..
	}

	// Adding one extra tag.
	// Multiple AddString() and RemoveString() can be chained together before calling Exec().
	err = getArt.Tags.AddString("three").Exec(session)

	if err != nil {
		// ..
	}
}

func TestSetHas(t *testing.T) {
	s := NewSet()
	s.AddString("hello")
	s.AddString("it's")
	s.AddString("me")

	if !s.HasString("hello") {
		t.Error("Did not have hello")
	}

	if !s.HasString("it's") {
		t.Error("Did not have it's")
	}

	if s.HasString("i") {
		t.Error("Had i")
	}
}

func TestSetInitializeSet(t *testing.T) {
	type testType struct {
		Foos *Set
	}

	// Pointer
	var testVal1 testType
	_, err := bucket().Set(&testVal1).Run(con())

	if err != nil {
		t.Error(err)
	}

	if testVal1.Foos == nil {
		t.Error("testVal1: Foos is not initialized")
	}

	// Non pointer
	var testVal2 testType
	_, err = bucket().Set(testVal2).Run(con())

	if err != nil {
		t.Error(err)
	}

	if testVal2.Foos != nil {
		t.Error("testVal2: Foos is not initialized")
	}
}
<|endoftext|>