<commit_before>package rizo

import (
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"time"
)

func check(err error) {
	if err != nil {
		panic(err)
	}
}

//RecordedRequest ...
type RecordedRequest struct {
	Request *http.Request
	body    string
}

//HTTPRequestPredicate ...
type HTTPRequestPredicate func(request RecordedRequest) bool

//HTTPResponseFactory ...
type HTTPResponseFactory func(writer http.ResponseWriter)

//UseWithPredicates ...
type UseWithPredicates struct {
	ResponseFactory   HTTPResponseFactory
	RequestPredicates []HTTPRequestPredicate
}

//RequestRecordingServer ...
type RequestRecordingServer struct {
	Requests []RecordedRequest
	port     int
	server   *httptest.Server
	use      []UseWithPredicates
}

//CreateRequestRecordingServer ...
func CreateRequestRecordingServer(port int) *RequestRecordingServer {
	return &RequestRecordingServer{
		Requests: []RecordedRequest{},
		port:     port,
		use:      []UseWithPredicates{},
	}
}

//CreateURL ...
func (instance *RequestRecordingServer) CreateURL(path string) string {
	return fmt.Sprintf("http://localhost:%d%s", instance.port, path)
}

func (instance *RequestRecordingServer) evaluatePredicates(recordedRequest RecordedRequest, w http.ResponseWriter) {
	for _, item := range instance.use {
		if item.RequestPredicates != nil {
			result := instance.Evaluate(recordedRequest, item.RequestPredicates...)
			if result {
				item.ResponseFactory(w)
				return
			}
		} else {
			item.ResponseFactory(w)
			return
		}
	}
	w.WriteHeader(http.StatusOK)
}

//Start ...
func (instance *RequestRecordingServer) Start() {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		check(err)
		recordedRequest := RecordedRequest{
			Request: r,
			body:    string(body),
		}
		instance.Requests = append(instance.Requests, recordedRequest)
		if instance.use != nil {
			instance.evaluatePredicates(recordedRequest, w)
		} else {
			w.WriteHeader(http.StatusOK)
		}
	})
	instance.server = httptest.NewUnstartedServer(handler)
	instance.server.Listener, err = net.Listen("tcp", ":"+strconv.Itoa(instance.port))
	if err != nil {
		panic(err)
	}
	instance.server.Start()
}

//Stop ...
func (instance *RequestRecordingServer) Stop() {
	if instance.server != nil {
		instance.server.Close()
		time.Sleep(1 * time.Microsecond)
	}
}

//Clear ...
func (instance *RequestRecordingServer) Clear() {
	instance.Requests = []RecordedRequest{}
	instance.use = []UseWithPredicates{}
}

//Evaluate ...
func (instance *RequestRecordingServer) Evaluate(request RecordedRequest, predicates ...HTTPRequestPredicate) bool {
	results := make([]bool, len(predicates))
	for index, predicate := range predicates {
		results[index] = predicate(request)
	}
	thing := true
	for _, result := range results {
		if !result {
			thing = false
			break
		}
	}
	return thing
}

//Find ...
func (instance *RequestRecordingServer) Find(predicates ...HTTPRequestPredicate) bool {
	for _, request := range instance.Requests {
		if instance.Evaluate(request) {
			return true
		}
	}
	return false
}

//Use ...
func (instance *RequestRecordingServer) Use(factory HTTPResponseFactory) *RequestRecordingServer {
	instance.use = append(instance.use, UseWithPredicates{
		ResponseFactory:   factory,
		RequestPredicates: []HTTPRequestPredicate{},
	})
	return instance
}

//For ...
func (instance *RequestRecordingServer) For(predicates ...HTTPRequestPredicate) {
	index := len(instance.use) - 1
	for _, item := range predicates {
		instance.use[index].RequestPredicates = append(instance.use[index].RequestPredicates, item)
	}
}

//RequestWithPath ...
func RequestWithPath(path string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.URL.Path == path
		if !result {
			log.Println(fmt.Sprintf("path does not equal %s it equals %s", path, r.Request.URL.Path))
		}
		return result
	})
}

//RequestWithMethod ...
func RequestWithMethod(method string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.Method == method
		if !result {
			log.Println("request method does not equal " + method)
		}
		return result
	})
}

//RequestWithHeader ...
func RequestWithHeader(key string, value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.Header.Get(key) == value
		if !result {
			log.Println(fmt.Sprintf("request method does not contain header with key %s and value %s actual %s", key, value, r.Request.Header.Get(key)))
		}
		return result
	})
}

//RequestWithBody ...
func RequestWithBody(value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := string(r.body) == value
		if !result {
			log.Println(fmt.Sprintf("request body does not equal %s it equals %s", value, r.body))
		}
		return result
	})
}

//RequestWithQuerystring ...
func RequestWithQuerystring(value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.URL.RawQuery == value
		if !result {
			log.Println("request query does not equal " + value + " | it equals " + r.Request.URL.RawQuery)
		}
		return result
	})
}
<commit_msg>Error if the listener cannot be setup<commit_after>package rizo

import (
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"time"
)

func check(err error) {
	if err != nil {
		panic(err)
	}
}

//RecordedRequest ...
type RecordedRequest struct {
	Request *http.Request
	body    string
}

//HTTPRequestPredicate ...
type HTTPRequestPredicate func(request RecordedRequest) bool

//HTTPResponseFactory ...
type HTTPResponseFactory func(writer http.ResponseWriter)

//UseWithPredicates ...
type UseWithPredicates struct {
	ResponseFactory   HTTPResponseFactory
	RequestPredicates []HTTPRequestPredicate
}

//RequestRecordingServer ...
type RequestRecordingServer struct {
	Requests []RecordedRequest
	port     int
	server   *httptest.Server
	use      []UseWithPredicates
}

//CreateRequestRecordingServer ...
func CreateRequestRecordingServer(port int) *RequestRecordingServer {
	return &RequestRecordingServer{
		Requests: []RecordedRequest{},
		port:     port,
		use:      []UseWithPredicates{},
	}
}

//CreateURL ...
func (instance *RequestRecordingServer) CreateURL(path string) string {
	return fmt.Sprintf("http://localhost:%d%s", instance.port, path)
}

func (instance *RequestRecordingServer) evaluatePredicates(recordedRequest RecordedRequest, w http.ResponseWriter) {
	for _, item := range instance.use {
		if item.RequestPredicates != nil {
			result := instance.Evaluate(recordedRequest, item.RequestPredicates...)
			if result {
				item.ResponseFactory(w)
				return
			}
		} else {
			item.ResponseFactory(w)
			return
		}
	}
	w.WriteHeader(http.StatusOK)
}

//Start ...
func (instance *RequestRecordingServer) Start() {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		check(err)
		recordedRequest := RecordedRequest{
			Request: r,
			body:    string(body),
		}
		instance.Requests = append(instance.Requests, recordedRequest)
		if instance.use != nil {
			instance.evaluatePredicates(recordedRequest, w)
		} else {
			w.WriteHeader(http.StatusOK)
		}
	})
	instance.server = httptest.NewUnstartedServer(handler)
	listener, err := net.Listen("tcp", ":"+strconv.Itoa(instance.port))
	if err != nil {
		panic(err)
	}
	instance.server.Listener = listener
	instance.server.Start()
}

//Stop ...
func (instance *RequestRecordingServer) Stop() {
	if instance.server != nil {
		instance.server.Close()
		time.Sleep(1 * time.Millisecond)
	}
}

//Clear ...
func (instance *RequestRecordingServer) Clear() {
	instance.Requests = []RecordedRequest{}
	instance.use = []UseWithPredicates{}
}

//Evaluate ...
func (instance *RequestRecordingServer) Evaluate(request RecordedRequest, predicates ...HTTPRequestPredicate) bool {
	results := make([]bool, len(predicates))
	for index, predicate := range predicates {
		results[index] = predicate(request)
	}
	thing := true
	for _, result := range results {
		if !result {
			thing = false
			break
		}
	}
	return thing
}

//Find ...
func (instance *RequestRecordingServer) Find(predicates ...HTTPRequestPredicate) bool {
	for _, request := range instance.Requests {
		if instance.Evaluate(request) {
			return true
		}
	}
	return false
}

//Use ...
func (instance *RequestRecordingServer) Use(factory HTTPResponseFactory) *RequestRecordingServer {
	instance.use = append(instance.use, UseWithPredicates{
		ResponseFactory:   factory,
		RequestPredicates: []HTTPRequestPredicate{},
	})
	return instance
}

//For ...
func (instance *RequestRecordingServer) For(predicates ...HTTPRequestPredicate) {
	index := len(instance.use) - 1
	for _, item := range predicates {
		instance.use[index].RequestPredicates = append(instance.use[index].RequestPredicates, item)
	}
}

//RequestWithPath ...
func RequestWithPath(path string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.URL.Path == path
		if !result {
			log.Println(fmt.Sprintf("path does not equal %s it equals %s", path, r.Request.URL.Path))
		}
		return result
	})
}

//RequestWithMethod ...
func RequestWithMethod(method string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.Method == method
		if !result {
			log.Println("request method does not equal " + method)
		}
		return result
	})
}

//RequestWithHeader ...
func RequestWithHeader(key string, value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.Header.Get(key) == value
		if !result {
			log.Println(fmt.Sprintf("request method does not contain header with key %s and value %s actual %s", key, value, r.Request.Header.Get(key)))
		}
		return result
	})
}

//RequestWithBody ...
func RequestWithBody(value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := string(r.body) == value
		if !result {
			log.Println(fmt.Sprintf("request body does not equal %s it equals %s", value, r.body))
		}
		return result
	})
}

//RequestWithQuerystring ...
func RequestWithQuerystring(value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.URL.RawQuery == value
		if !result {
			log.Println("request query does not equal " + value + " | it equals " + r.Request.URL.RawQuery)
		}
		return result
	})
}
<|endoftext|>


<commit_before>package rizo

import (
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"time"
)

func check(err error) {
	if err != nil {
		panic(err)
	}
}

//RecordedRequest ...
type RecordedRequest struct {
	Request *http.Request
	body    string
}

//HTTPRequestPredicate ...
type HTTPRequestPredicate func(request RecordedRequest) bool

//HTTPResponseFactory ...
type HTTPResponseFactory func(writer http.ResponseWriter)

//UseWithPredicates ...
type UseWithPredicates struct {
	ResponseFactory   HTTPResponseFactory
	RequestPredicates []HTTPRequestPredicate
}

//RequestRecordingServer ...
type RequestRecordingServer struct {
	Requests []RecordedRequest
	port     int
	server   *httptest.Server
	use      []UseWithPredicates
}

//CreateRequestRecordingServer ...
func CreateRequestRecordingServer(port int) *RequestRecordingServer {
	return &RequestRecordingServer{
		Requests: []RecordedRequest{},
		port:     port,
		use:      []UseWithPredicates{},
	}
}

func (instance *RequestRecordingServer) evaluatePredicates(recordedRequest RecordedRequest, w http.ResponseWriter) {
	for _, item := range instance.use {
		if item.RequestPredicates != nil {
			result := instance.Evaluate(recordedRequest, item.RequestPredicates...)
			if result {
				item.ResponseFactory(w)
				return
			}
		} else {
			item.ResponseFactory(w)
			return
		}
	}
	w.WriteHeader(http.StatusOK)
}

//Start ...
func (instance *RequestRecordingServer) Start() {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		check(err)
		recordedRequest := RecordedRequest{
			Request: r,
			body:    string(body),
		}
		instance.Requests = append(instance.Requests, recordedRequest)
		if instance.use != nil {
			instance.evaluatePredicates(recordedRequest, w)
		} else {
			w.WriteHeader(http.StatusOK)
		}
	})
	instance.server = httptest.NewUnstartedServer(handler)
	instance.server.Listener, _ = net.Listen("tcp", ":"+strconv.Itoa(instance.port))
	instance.server.Start()
}

//Stop ...
func (instance *RequestRecordingServer) Stop() {
	instance.server.Close()
	time.Sleep(1 * time.Microsecond)
}

//Clear ...
func (instance *RequestRecordingServer) Clear() {
	instance.Requests = []RecordedRequest{}
	instance.use = []UseWithPredicates{}
}

//Evaluate ...
func (instance *RequestRecordingServer) Evaluate(request RecordedRequest, predicates ...HTTPRequestPredicate) bool {
	results := make([]bool, len(predicates))
	for index, predicate := range predicates {
		results[index] = predicate(request)
	}
	thing := true
	for _, result := range results {
		if !result {
			thing = false
			break
		}
	}
	return thing
}

//Find ...
func (instance *RequestRecordingServer) Find(predicates ...HTTPRequestPredicate) bool {
	for _, request := range instance.Requests {
		if instance.Evaluate(request) {
			return true
		}
	}
	return false
}

//Use ...
func (instance *RequestRecordingServer) Use(factory HTTPResponseFactory) *RequestRecordingServer {
	instance.use = append(instance.use, UseWithPredicates{
		ResponseFactory:   factory,
		RequestPredicates: []HTTPRequestPredicate{},
	})
	return instance
}

//For ...
func (instance *RequestRecordingServer) For(predicates ...HTTPRequestPredicate) {
	index := len(instance.use) - 1
	for _, item := range predicates {
		instance.use[index].RequestPredicates = append(instance.use[index].RequestPredicates, item)
	}
}

//RequestWithPath ...
func RequestWithPath(path string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.URL.Path == path
		if !result {
			log.Println(fmt.Sprintf("path does not equal %s it equals %s", path, r.Request.URL.Path))
		}
		return result
	})
}

//RequestWithMethod ...
func RequestWithMethod(method string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.Method == method
		if !result {
			log.Println("request method does not equal " + method)
		}
		return result
	})
}

//RequestWithHeader ...
func RequestWithHeader(key string, value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.Header.Get(key) == value
		if !result {
			log.Println(fmt.Sprintf("request method does not contain header with key %s and value %s actual %s", key, value, r.Request.Header.Get(key)))
		}
		return result
	})
}

//RequestWithBody ...
func RequestWithBody(value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := string(r.body) == value
		if !result {
			log.Println(fmt.Sprintf("request body does not equal %s it equals %s", value, r.body))
		}
		return result
	})
}

//RequestWithQuerystring ...
func RequestWithQuerystring(value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.URL.RawQuery == value
		if !result {
			log.Println("request query does not equal " + value + " | it equals " + r.Request.URL.RawQuery)
		}
		return result
	})
}
<commit_msg>Adding a convenience method for constructing urls aimed at the Recording Server<commit_after>package rizo

import (
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"time"
)

func check(err error) {
	if err != nil {
		panic(err)
	}
}

//RecordedRequest ...
type RecordedRequest struct {
	Request *http.Request
	body    string
}

//HTTPRequestPredicate ...
type HTTPRequestPredicate func(request RecordedRequest) bool

//HTTPResponseFactory ...
type HTTPResponseFactory func(writer http.ResponseWriter)

//UseWithPredicates ...
type UseWithPredicates struct {
	ResponseFactory   HTTPResponseFactory
	RequestPredicates []HTTPRequestPredicate
}

//RequestRecordingServer ...
type RequestRecordingServer struct {
	Requests []RecordedRequest
	port     int
	server   *httptest.Server
	use      []UseWithPredicates
}

//CreateRequestRecordingServer ...
func CreateRequestRecordingServer(port int) *RequestRecordingServer {
	return &RequestRecordingServer{
		Requests: []RecordedRequest{},
		port:     port,
		use:      []UseWithPredicates{},
	}
}

//CreateURL ...
func (instance *RequestRecordingServer) CreateURL(path string) string {
	return fmt.Sprintf("http://localhost:%d%s", instance.port, path)
}

func (instance *RequestRecordingServer) evaluatePredicates(recordedRequest RecordedRequest, w http.ResponseWriter) {
	for _, item := range instance.use {
		if item.RequestPredicates != nil {
			result := instance.Evaluate(recordedRequest, item.RequestPredicates...)
			if result {
				item.ResponseFactory(w)
				return
			}
		} else {
			item.ResponseFactory(w)
			return
		}
	}
	w.WriteHeader(http.StatusOK)
}

//Start ...
func (instance *RequestRecordingServer) Start() {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		check(err)
		recordedRequest := RecordedRequest{
			Request: r,
			body:    string(body),
		}
		instance.Requests = append(instance.Requests, recordedRequest)
		if instance.use != nil {
			instance.evaluatePredicates(recordedRequest, w)
		} else {
			w.WriteHeader(http.StatusOK)
		}
	})
	instance.server = httptest.NewUnstartedServer(handler)
	instance.server.Listener, _ = net.Listen("tcp", ":"+strconv.Itoa(instance.port))
	instance.server.Start()
}

//Stop ...
func (instance *RequestRecordingServer) Stop() {
	instance.server.Close()
	time.Sleep(1 * time.Microsecond)
}

//Clear ...
func (instance *RequestRecordingServer) Clear() {
	instance.Requests = []RecordedRequest{}
	instance.use = []UseWithPredicates{}
}

//Evaluate ...
func (instance *RequestRecordingServer) Evaluate(request RecordedRequest, predicates ...HTTPRequestPredicate) bool {
	results := make([]bool, len(predicates))
	for index, predicate := range predicates {
		results[index] = predicate(request)
	}
	thing := true
	for _, result := range results {
		if !result {
			thing = false
			break
		}
	}
	return thing
}

//Find ...
func (instance *RequestRecordingServer) Find(predicates ...HTTPRequestPredicate) bool {
	for _, request := range instance.Requests {
		if instance.Evaluate(request) {
			return true
		}
	}
	return false
}

//Use ...
func (instance *RequestRecordingServer) Use(factory HTTPResponseFactory) *RequestRecordingServer {
	instance.use = append(instance.use, UseWithPredicates{
		ResponseFactory:   factory,
		RequestPredicates: []HTTPRequestPredicate{},
	})
	return instance
}

//For ...
func (instance *RequestRecordingServer) For(predicates ...HTTPRequestPredicate) {
	index := len(instance.use) - 1
	for _, item := range predicates {
		instance.use[index].RequestPredicates = append(instance.use[index].RequestPredicates, item)
	}
}

//RequestWithPath ...
func RequestWithPath(path string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.URL.Path == path
		if !result {
			log.Println(fmt.Sprintf("path does not equal %s it equals %s", path, r.Request.URL.Path))
		}
		return result
	})
}

//RequestWithMethod ...
func RequestWithMethod(method string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.Method == method
		if !result {
			log.Println("request method does not equal " + method)
		}
		return result
	})
}

//RequestWithHeader ...
func RequestWithHeader(key string, value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.Header.Get(key) == value
		if !result {
			log.Println(fmt.Sprintf("request method does not contain header with key %s and value %s actual %s", key, value, r.Request.Header.Get(key)))
		}
		return result
	})
}

//RequestWithBody ...
func RequestWithBody(value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := string(r.body) == value
		if !result {
			log.Println(fmt.Sprintf("request body does not equal %s it equals %s", value, r.body))
		}
		return result
	})
}

//RequestWithQuerystring ...
func RequestWithQuerystring(value string) HTTPRequestPredicate {
	return HTTPRequestPredicate(func(r RecordedRequest) bool {
		result := r.Request.URL.RawQuery == value
		if !result {
			log.Println("request query does not equal " + value + " | it equals " + r.Request.URL.RawQuery)
		}
		return result
	})
}
<|endoftext|>


<commit_before>package graceful

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"sync"
	"syscall"
	"testing"
	"time"
)

var (
	killTime    = 500 * time.Millisecond
	timeoutTime = 1000 * time.Millisecond
	waitTime    = 100 * time.Millisecond
)

func runQuery(t *testing.T, expected int, shouldErr bool, wg *sync.WaitGroup, once *sync.Once) {
	wg.Add(1)
	defer wg.Done()
	client := http.Client{}
	r, err := client.Get("http://localhost:3000")
	if shouldErr && err == nil {
		once.Do(func() {
			t.Fatal("Expected an error but none was encountered.")
		})
	} else if shouldErr && err != nil {
		if checkErr(t, err, once) {
			return
		}
	}
	if r != nil && r.StatusCode != expected {
		once.Do(func() {
			t.Fatalf("Incorrect status code on response. Expected %d. Got %d", expected, r.StatusCode)
		})
	} else if r == nil {
		once.Do(func() {
			t.Fatal("No response when a response was expected.")
		})
	}
}

func checkErr(t *testing.T, err error, once *sync.Once) bool {
	if err.(*url.Error).Err == io.EOF {
		return true
	}
	errno := err.(*url.Error).Err.(*net.OpError).Err.(syscall.Errno)
	if errno == syscall.ECONNREFUSED {
		return true
	} else if err != nil {
		once.Do(func() {
			t.Fatal("Error on Get:", err)
		})
	}
	return false
}

func createListener(sleep time.Duration) (*http.Server, net.Listener, error) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) {
		time.Sleep(sleep)
		rw.WriteHeader(http.StatusOK)
	})

	server := &http.Server{Addr: ":3000", Handler: mux}
	l, err := net.Listen("tcp", ":3000")
	if err != nil {
		fmt.Println(err)
	}
	return server, l, err
}

func runServer(timeout, sleep time.Duration, c chan os.Signal) error {
	server, l, err := createListener(sleep)
	if err != nil {
		return err
	}

	srv := &Server{Timeout: timeout, Server: server, interrupt: c}
	return srv.Serve(l)
}

func launchTestQueries(t *testing.T, wg *sync.WaitGroup, c chan os.Signal) {
	var once sync.Once
	for i := 0; i < 8; i++ {
		go runQuery(t, http.StatusOK, false, wg, &once)
	}

	time.Sleep(waitTime)
	c <- os.Interrupt
	time.Sleep(waitTime)

	for i := 0; i < 8; i++ {
		go runQuery(t, 0, true, wg, &once)
	}

	wg.Done()
}

func TestGracefulRun(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		runServer(killTime, killTime/2, c)
		wg.Done()
	}()

	wg.Add(1)
	go launchTestQueries(t, &wg, c)
	wg.Wait()
}

func TestGracefulRunTimesOut(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		runServer(killTime, killTime*10, c)
		wg.Done()
	}()

	var once sync.Once
	wg.Add(1)
	go func() {
		for i := 0; i < 8; i++ {
			go runQuery(t, 0, true, &wg, &once)
		}
		time.Sleep(waitTime)
		c <- os.Interrupt
		time.Sleep(waitTime)
		for i := 0; i < 8; i++ {
			go runQuery(t, 0, true, &wg, &once)
		}
		wg.Done()
	}()

	wg.Wait()

}

func TestGracefulRunDoesntTimeOut(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		runServer(0, killTime*2, c)
		wg.Done()
	}()

	wg.Add(1)
	go launchTestQueries(t, &wg, c)
	wg.Wait()
}

func TestGracefulRunNoRequests(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		runServer(0, killTime*2, c)
		wg.Done()
	}()

	c <- os.Interrupt

	wg.Wait()

}

func TestGracefulForwardsConnState(t *testing.T) {
	c := make(chan os.Signal, 1)
	states := make(map[http.ConnState]int)
	var stateLock sync.Mutex

	connState := func(conn net.Conn, state http.ConnState) {
		stateLock.Lock()
		states[state]++
		stateLock.Unlock()
	}

	var wg sync.WaitGroup
	wg.Add(1)

	expected := map[http.ConnState]int{
		http.StateNew:    8,
		http.StateActive: 8,
		http.StateClosed: 8,
	}

	go func() {
		server, l, _ := createListener(killTime / 2)
		srv := &Server{
			ConnState: connState,
			Timeout:   killTime,
			Server:    server,
			interrupt: c,
		}
		srv.Serve(l)

		wg.Done()
	}()

	wg.Add(1)
	go launchTestQueries(t, &wg, c)
	wg.Wait()

	stateLock.Lock()
	if !reflect.DeepEqual(states, expected) {
		t.Errorf("Incorrect connection state tracking.\n  actual: %v\nexpected: %v\n", states, expected)
	}
	stateLock.Unlock()
}

func TestGracefulExplicitStop(t *testing.T) {
	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	srv := &Server{Timeout: killTime, Server: server}

	go func() {
		go srv.Serve(l)
		time.Sleep(waitTime)
		srv.Stop(killTime)
	}()

	// block on the stopChan until the server has shut down
	select {
	case <-srv.StopChan():
	case <-time.After(timeoutTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}
}

func TestGracefulExplicitStopOverride(t *testing.T) {
	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	srv := &Server{Timeout: killTime, Server: server}

	go func() {
		go srv.Serve(l)
		time.Sleep(waitTime)
		srv.Stop(killTime / 2)
	}()

	// block on the stopChan until the server has shut down
	select {
	case <-srv.StopChan():
	case <-time.After(killTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}
}

func TestShutdownInitiatedCallback(t *testing.T) {
	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	called := make(chan struct{})
	cb := func() { close(called) }

	srv := &Server{Server: server, ShutdownInitiated: cb}

	go func() {
		go srv.Serve(l)
		time.Sleep(waitTime)
		srv.Stop(killTime)
	}()

	select {
	case <-called:
	case <-time.After(killTime):
		t.Fatal("Timed out while waiting for ShutdownInitiated callback to be called")
	}
}
func hijackingListener(srv *Server) (*http.Server, net.Listener, error) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) {
		conn, bufrw, err := rw.(http.Hijacker).Hijack()
		if err != nil {
			http.Error(rw, "webserver doesn't support hijacking", http.StatusInternalServerError)
			return
		}

		defer conn.Close()

		bufrw.WriteString("HTTP/1.1 200 OK\r\n\r\n")
		bufrw.Flush()
	})

	server := &http.Server{Addr: ":3000", Handler: mux}
	l, err := net.Listen("tcp", ":3000")
	return server, l, err
}

func TestNotifyClosed(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	srv := &Server{Timeout: killTime, interrupt: c}
	server, l, err := hijackingListener(srv)
	if err != nil {
		t.Fatal(err)
	}

	srv.Server = server

	go func() {
		srv.Serve(l)
		wg.Done()
	}()

	var once sync.Once
	for i := 0; i < 8; i++ {
		runQuery(t, http.StatusOK, false, &wg, &once)
	}

	srv.Stop(0)

	// block on the stopChan until the server has shut down
	select {
	case <-srv.StopChan():
	case <-time.After(timeoutTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}

	if len(srv.connections) > 0 {
		t.Fatal("hijacked connections should not be managed")
	}

}

func TestStopDeadlock(t *testing.T) {
	c := make(chan struct{})

	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	srv := &Server{Server: server, NoSignalHandling: true}

	go func() {
		time.Sleep(waitTime)
		srv.Serve(l)
	}()

	go func() {
		srv.Stop(0)
		close(c)
	}()

	select {
	case <-c:
	case <-time.After(timeoutTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}
}
<commit_msg>Add a test that reproduces a data race in Stop()<commit_after>package graceful

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"sync"
	"syscall"
	"testing"
	"time"
)

var (
	killTime    = 500 * time.Millisecond
	timeoutTime = 1000 * time.Millisecond
	waitTime    = 100 * time.Millisecond
)

func runQuery(t *testing.T, expected int, shouldErr bool, wg *sync.WaitGroup, once *sync.Once) {
	wg.Add(1)
	defer wg.Done()
	client := http.Client{}
	r, err := client.Get("http://localhost:9654")
	if shouldErr && err == nil {
		once.Do(func() {
			t.Fatal("Expected an error but none was encountered.")
		})
	} else if shouldErr && err != nil {
		if checkErr(t, err, once) {
			return
		}
	}
	if r != nil && r.StatusCode != expected {
		once.Do(func() {
			t.Fatalf("Incorrect status code on response. Expected %d. Got %d", expected, r.StatusCode)
		})
	} else if r == nil {
		once.Do(func() {
			t.Fatal("No response when a response was expected.")
		})
	}
}

func checkErr(t *testing.T, err error, once *sync.Once) bool {
	if err.(*url.Error).Err == io.EOF {
		return true
	}
	errno := err.(*url.Error).Err.(*net.OpError).Err.(syscall.Errno)
	if errno == syscall.ECONNREFUSED {
		return true
	} else if err != nil {
		once.Do(func() {
			t.Fatal("Error on Get:", err)
		})
	}
	return false
}

func createListener(sleep time.Duration) (*http.Server, net.Listener, error) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) {
		time.Sleep(sleep)
		rw.WriteHeader(http.StatusOK)
	})

	server := &http.Server{Addr: ":9654", Handler: mux}
	l, err := net.Listen("tcp", ":9654")
	if err != nil {
		fmt.Println(err)
	}
	return server, l, err
}

func runServer(timeout, sleep time.Duration, c chan os.Signal) error {
	server, l, err := createListener(sleep)
	if err != nil {
		return err
	}

	srv := &Server{Timeout: timeout, Server: server, interrupt: c}
	return srv.Serve(l)
}

func launchTestQueries(t *testing.T, wg *sync.WaitGroup, c chan os.Signal) {
	var once sync.Once
	for i := 0; i < 8; i++ {
		go runQuery(t, http.StatusOK, false, wg, &once)
	}

	time.Sleep(waitTime)
	c <- os.Interrupt
	time.Sleep(waitTime)

	for i := 0; i < 8; i++ {
		go runQuery(t, 0, true, wg, &once)
	}

	wg.Done()
}

func TestGracefulRun(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		runServer(killTime, killTime/2, c)
		wg.Done()
	}()

	wg.Add(1)
	go launchTestQueries(t, &wg, c)
	wg.Wait()
}

func TestGracefulRunTimesOut(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		runServer(killTime, killTime*10, c)
		wg.Done()
	}()

	var once sync.Once
	wg.Add(1)
	go func() {
		for i := 0; i < 8; i++ {
			go runQuery(t, 0, true, &wg, &once)
		}
		time.Sleep(waitTime)
		c <- os.Interrupt
		time.Sleep(waitTime)
		for i := 0; i < 8; i++ {
			go runQuery(t, 0, true, &wg, &once)
		}
		wg.Done()
	}()

	wg.Wait()

}

func TestGracefulRunDoesntTimeOut(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		runServer(0, killTime*2, c)
		wg.Done()
	}()

	wg.Add(1)
	go launchTestQueries(t, &wg, c)
	wg.Wait()
}

func TestGracefulRunNoRequests(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		runServer(0, killTime*2, c)
		wg.Done()
	}()

	c <- os.Interrupt

	wg.Wait()

}

func TestGracefulForwardsConnState(t *testing.T) {
	c := make(chan os.Signal, 1)
	states := make(map[http.ConnState]int)
	var stateLock sync.Mutex

	connState := func(conn net.Conn, state http.ConnState) {
		stateLock.Lock()
		states[state]++
		stateLock.Unlock()
	}

	var wg sync.WaitGroup
	wg.Add(1)

	expected := map[http.ConnState]int{
		http.StateNew:    8,
		http.StateActive: 8,
		http.StateClosed: 8,
	}

	go func() {
		server, l, _ := createListener(killTime / 2)
		srv := &Server{
			ConnState: connState,
			Timeout:   killTime,
			Server:    server,
			interrupt: c,
		}
		srv.Serve(l)

		wg.Done()
	}()

	wg.Add(1)
	go launchTestQueries(t, &wg, c)
	wg.Wait()

	stateLock.Lock()
	if !reflect.DeepEqual(states, expected) {
		t.Errorf("Incorrect connection state tracking.\n  actual: %v\nexpected: %v\n", states, expected)
	}
	stateLock.Unlock()
}

func TestGracefulExplicitStop(t *testing.T) {
	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	srv := &Server{Timeout: killTime, Server: server}

	go func() {
		go srv.Serve(l)
		time.Sleep(waitTime)
		srv.Stop(killTime)
	}()

	// block on the stopChan until the server has shut down
	select {
	case <-srv.StopChan():
	case <-time.After(timeoutTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}
}

func TestGracefulExplicitStopOverride(t *testing.T) {
	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	srv := &Server{Timeout: killTime, Server: server}

	go func() {
		go srv.Serve(l)
		time.Sleep(waitTime)
		srv.Stop(killTime / 2)
	}()

	// block on the stopChan until the server has shut down
	select {
	case <-srv.StopChan():
	case <-time.After(killTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}
}

func TestShutdownInitiatedCallback(t *testing.T) {
	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	called := make(chan struct{})
	cb := func() { close(called) }

	srv := &Server{Server: server, ShutdownInitiated: cb}

	go func() {
		go srv.Serve(l)
		time.Sleep(waitTime)
		srv.Stop(killTime)
	}()

	select {
	case <-called:
	case <-time.After(killTime):
		t.Fatal("Timed out while waiting for ShutdownInitiated callback to be called")
	}
}
func hijackingListener(srv *Server) (*http.Server, net.Listener, error) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) {
		conn, bufrw, err := rw.(http.Hijacker).Hijack()
		if err != nil {
			http.Error(rw, "webserver doesn't support hijacking", http.StatusInternalServerError)
			return
		}

		defer conn.Close()

		bufrw.WriteString("HTTP/1.1 200 OK\r\n\r\n")
		bufrw.Flush()
	})

	server := &http.Server{Addr: ":9654", Handler: mux}
	l, err := net.Listen("tcp", ":9654")
	return server, l, err
}

func TestNotifyClosed(t *testing.T) {
	c := make(chan os.Signal, 1)

	var wg sync.WaitGroup
	wg.Add(1)

	srv := &Server{Timeout: killTime, interrupt: c}
	server, l, err := hijackingListener(srv)
	if err != nil {
		t.Fatal(err)
	}

	srv.Server = server

	go func() {
		srv.Serve(l)
		wg.Done()
	}()

	var once sync.Once
	for i := 0; i < 8; i++ {
		runQuery(t, http.StatusOK, false, &wg, &once)
	}

	srv.Stop(0)

	// block on the stopChan until the server has shut down
	select {
	case <-srv.StopChan():
	case <-time.After(timeoutTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}

	if len(srv.connections) > 0 {
		t.Fatal("hijacked connections should not be managed")
	}

}

func TestStopDeadlock(t *testing.T) {
	c := make(chan struct{})

	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	srv := &Server{Server: server, NoSignalHandling: true}

	go func() {
		time.Sleep(waitTime)
		srv.Serve(l)
	}()

	go func() {
		srv.Stop(0)
		close(c)
	}()

	select {
	case <-c:
	case <-time.After(timeoutTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}
}

// Run with --race
func TestStopRace(t *testing.T) {
	server, l, err := createListener(1 * time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}

	srv := &Server{Timeout: killTime, Server: server}

	go func() {
		go srv.Serve(l)
		srv.Stop(killTime)
	}()
	srv.Stop(0)
	select {
	case <-srv.StopChan():
	case <-time.After(timeoutTime):
		t.Fatal("Timed out while waiting for explicit stop to complete")
	}
}
<|endoftext|>


<commit_before>// Package filecache provides simple file system cache
package filecache

import (
	"crypto/md5"
	"encoding/hex"
	"io/ioutil"
	"log"
	"os"
	"time"
)

// Type Handler holds the cache directory and cache
// life in minutes
type Handler struct {
	Dir  string
	Life float64
}

// Set saves content to the cache
func (h Handler) Set(key string, data []byte) error {
	return ioutil.WriteFile(h.filename(key), data, 0644)
}

// Get returns content from the cache
func (h Handler) Get(key string) []byte {
	file := h.filename(key)
	result, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}
	fi, err := os.Stat(file)
	panicon(err)
	age := time.Since(fi.ModTime()).Minutes()
	if age > h.Life {
		err := os.Remove(file)
		panicon(err)
		return nil
	}
	return result
}

// filename returns the full path to the cache file
func (h Handler) filename(key string) string {
	return h.dir() + "/" + hash(key)
}

// dir creates and returns a valid cache directory
// temp dir is returned if Handler.Dir is not set
func (h Handler) dir() string {
	if h.Dir == "" {
		h.Dir = os.TempDir()
		return h.Dir
	}

	yes, err := exists(h.Dir)
	panicon(err)
	if !yes {
		err := os.MkdirAll(h.Dir, 0666)
		panicon(err)
	}
	return h.Dir
}

// hash returns the checksum of a string
func hash(text string) string {
	hasher := md5.New()
	hasher.Write([]byte(text))
	return hex.EncodeToString(hasher.Sum(nil))
}

// exists returns whether the given file or directory exists or not
func exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}

func panicon(err error) {
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Filename function is now exported<commit_after>// Package filecache provides simple file system cache
package filecache

import (
	"crypto/md5"
	"encoding/hex"
	"io/ioutil"
	"log"
	"os"
	"time"
)

// Type Handler holds the cache directory and cache
// life in minutes
type Handler struct {
	Dir  string
	Life float64
}

// Set saves content to the cache
func (h Handler) Set(key string, data []byte) error {
	return ioutil.WriteFile(h.Filename(key), data, 0644)
}

// Get returns content from the cache
func (h Handler) Get(key string) []byte {
	file := h.Filename(key)
	result, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}
	fi, err := os.Stat(file)
	panicon(err)
	age := time.Since(fi.ModTime()).Minutes()
	if age > h.Life {
		err := os.Remove(file)
		panicon(err)
		return nil
	}
	return result
}

// Filename returns the full path to the cache file
func (h Handler) Filename(key string) string {
	return h.dir() + "/" + hash(key)
}

// dir creates and returns a valid cache directory
// temp dir is returned if Handler.Dir is not set
func (h Handler) dir() string {
	if h.Dir == "" {
		h.Dir = os.TempDir()
		return h.Dir
	}

	yes, err := exists(h.Dir)
	panicon(err)
	if !yes {
		err := os.MkdirAll(h.Dir, 0666)
		panicon(err)
	}
	return h.Dir
}

// hash returns the checksum of a string
func hash(text string) string {
	hasher := md5.New()
	hasher.Write([]byte(text))
	return hex.EncodeToString(hasher.Sum(nil))
}

// exists returns whether the given file or directory exists or not
func exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}

func panicon(err error) {
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package koding

import (
	"fmt"
	"koding/db/models"
	"koding/kites/kloud/contexthelper/session"
	"koding/kites/kloud/eventer"
	"koding/kites/kloud/klient"
	"koding/kites/kloud/kloud"
	"koding/kites/kloud/machinestate"
	"koding/kites/kloud/plans"
	"time"

	"golang.org/x/net/context"

	"github.com/koding/logging"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// Machine represents a single MongodDB document that represents a Koding
// Provider from the jMachines collection.
type Machine struct {
	Id          bson.ObjectId `bson:"_id" json:"-"`
	Label       string        `bson:"label"`
	Domain      string        `bson:"domain"`
	QueryString string        `bson:"queryString"`
	IpAddress   string        `bson:"ipAddress"`
	Assignee    struct {
		InProgress      bool      `bson:"inProgress"`
		AssignedAt      time.Time `bson:"assignedAt"`
		KlientMissingAt time.Time `bson:"klientMissingAt"`
	} `bson:"assignee"`
	Status struct {
		State      string    `bson:"state"`
		Reason     string    `bson:"reason"`
		ModifiedAt time.Time `bson:"modifiedAt"`
	} `bson:"status"`
	Provider   string    `bson:"provider"`
	Credential string    `bson:"credential"`
	CreatedAt  time.Time `bson:"createdAt"`
	Meta       struct {
		AlwaysOn     bool   `bson:"alwaysOn"`
		InstanceId   string `structs:"instanceId" bson:"instanceId"`
		InstanceType string `structs:"instance_type" bson:"instance_type"`
		InstanceName string `structs:"instanceName" bson:"instanceName"`
		Region       string `structs:"region" bson:"region"`
		StorageSize  int    `structs:"storage_size" bson:"storage_size"`
		SourceAmi    string `structs:"source_ami" bson:"source_ami"`
		SnapshotId   string `structs:"snapshotId" bson:"snapshotId"`
	} `bson:"meta"`
	Users  []models.Permissions `bson:"users"`
	Groups []models.Permissions `bson:"groups"`

	// internal fields, not availabile in MongoDB schema
	Username string                 `bson:"-"`
	User     *models.User           `bson:"-"`
	Payment  *plans.PaymentResponse `bson:"-"`
	Checker  plans.Checker          `bson:"-"`
	Session  *session.Session       `bson:"-"`
	Log      logging.Logger         `bson:"-"`
	locker   kloud.Locker           `bson:"-"`

	// cleanFuncs are a list of functions that are called when after a method
	// is finished
	cleanFuncs []func()
}

// runCleanupFunctions calls all cleanup functions and set the
// list to nil. Once called any other call will not have any
// effect.
func (m *Machine) runCleanupFunctions() {
	if m.cleanFuncs == nil {
		return
	}

	for _, fn := range m.cleanFuncs {
		fn()
	}

	m.cleanFuncs = nil
}

func (m *Machine) State() machinestate.State {
	return machinestate.States[m.Status.State]
}

func (m *Machine) PublicIpAddress() string {
	return m.IpAddress
}

// push pushes the given message to the eventer
func (m *Machine) push(msg string, percentage int, state machinestate.State) {
	if m.Session.Eventer != nil {
		m.Session.Eventer.Push(&eventer.Event{
			Message:    msg,
			Percentage: percentage,
			Status:     state,
		})
	}
}

// switchAWSRegion switches to the given AWS region. This should be only used when
// you know what to do, otherwiese never, never change the region of a machine.
func (m *Machine) switchAWSRegion(region string) error {
	m.Meta.InstanceId = "" // we neglect any previous instanceId
	m.QueryString = ""     //
	m.Meta.Region = "us-east-1"

	client, err := m.Session.AWSClients.Region("us-east-1")
	if err != nil {
		return err
	}
	m.Session.AWSClient.Client = client

	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"meta.instanceId": "",
				"queryString":     "",
				"meta.region":     "us-east-1",
			}},
		)
	})
}

// markAsNotInitialized marks the machine as NotInitialized by cleaning up all
// necessary fields and marking the VM as notinitialized so the User can build
// it again.
func (m *Machine) markAsNotInitialized() error {
	m.Log.Warning("Instance is not available. Marking it as NotInitialized")
	if err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"ipAddress":         "",
				"queryString":       "",
				"meta.instanceType": "",
				"meta.instanceName": "",
				"meta.instanceId":   "",
				"status.state":      machinestate.NotInitialized.String(),
				"status.modifiedAt": time.Now().UTC(),
				"status.reason":     "Machine is marked as NotInitialized",
			}},
		)
	}); err != nil {
		return err
	}

	m.IpAddress = ""
	m.QueryString = ""
	m.Meta.InstanceType = ""
	m.Meta.InstanceName = ""
	m.Meta.InstanceId = ""

	// so any State() method can return the correct status
	m.Status.State = machinestate.NotInitialized.String()
	return nil
}

func (m *Machine) markAsStopped() error {
	return m.markAsStoppedWithReason("Machine is stopped")
}

func (m *Machine) markAsStoppedWithReason(reason string) error {
	m.Log.Debug("Marking instance as stopped")
	if err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"ipAddress":         "",
				"status.state":      machinestate.Stopped.String(),
				"status.modifiedAt": time.Now().UTC(),
				"status.reason":     reason,
			}},
		)
	}); err != nil {
		return err
	}

	// so any State() method returns the correct status
	m.Status.State = machinestate.Stopped.String()
	m.IpAddress = ""
	return nil
}

func (m *Machine) updateStorageSize(size int) error {
	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{"meta.storage_size": size}},
		)
	})
}

func (m *Machine) isKlientReady() bool {
	m.Log.Debug("All finished, testing for klient connection IP [%s]", m.IpAddress)
	klientRef, err := klient.NewWithTimeout(m.Session.Kite, m.QueryString, time.Minute*5)
	if err != nil {
		m.Log.Warning("Connecting to remote Klient instance err: %s", err)
		return false
	}
	defer klientRef.Close()

	m.Log.Debug("Sending a ping message")
	if err := klientRef.Ping(); err != nil {
		m.Log.Debug("Sending a ping message err: %s", err)
		return false
	}

	return true
}

// Lock performs a Lock on this Machine
func (m *Machine) Lock() error {
	if !m.Id.Valid() {
		return kloud.NewError(kloud.ErrMachineIdMissing)
	}

	if m.locker == nil {
		return fmt.Errorf("Machine '%s' missing Locker", m.Id.Hex())
	}

	return m.locker.Lock(m.Id.Hex())
}

// Unlock performs an Unlock on this Machine instance
func (m *Machine) Unlock() error {
	if !m.Id.Valid() {
		return kloud.NewError(kloud.ErrMachineIdMissing)
	}

	if m.locker == nil {
		return fmt.Errorf("Machine '%s' missing Locker", m.Id.Hex())
	}

	// Unlock does not return an error
	m.locker.Unlock(m.Id.Hex())
	return nil
}

// klientIsNotMissing will unset the `assignee.klientMissingAt` value
// from the database, only if the Machine.Assignee.KlientMissingAt value
// has data. Therefor it is safe to call as frequently.
func (m *Machine) klientIsNotMissing() error {
	if m.Assignee.KlientMissingAt.IsZero() {
		return nil
	}

	m.Log.Debug("Clearing assignee.klientMissingAt")

	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$unset": bson.M{"assignee.klientMissingAt": ""}},
		)
	})
}

// stopIfKlientIsMissing will stop the current Machine X minutes after
// the `assignee.klientMissingAt` value. If the value does not exist in
// the databse, it will write it and return.
//
// Therefor, this method is expected be called as often as needed,
// and will shutdown the Machine if klient has been missing for too long.
func (m *Machine) stopIfKlientIsMissing(ctx context.Context) error {

	// If this is the first time Klient has been found missing,
	// set the missingat time and return
	if m.Assignee.KlientMissingAt.IsZero() {
		m.Log.Debug("Klient has been reported missing, recording this as the first time it went missing")

		return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
			return c.UpdateId(
				m.Id,
				bson.M{"$set": bson.M{"assignee.klientMissingAt": time.Now().UTC()}},
			)
		})
	}

	// If the klient has been missing less than X minutes, don't stop
	if time.Since(m.Assignee.KlientMissingAt) < time.Minute*5 {
		return nil
	}

	// lock so it doesn't interfere with others.
	err := m.Lock()
	defer m.Unlock()
	if err != nil {
		return err
	}

	// Hasta la vista, baby!
	m.Log.Info("======> STOP started (missing klient) <======")
	if err := m.Stop(ctx); err != nil {
		m.Log.Info("======> STOP aborted (missing klient: %s) <======", err)
		return err
	}
	m.Log.Info("======> STOP finished (missing klient) <======")

	// Clear the klientMissingAt field, or we risk Stopping the user's
	// machine next time they run it, without waiting the proper X minute
	// timeout.
	m.klientIsNotMissing()

	return nil
}
<commit_msg>kloud: Changed log wording<commit_after>package koding

import (
	"fmt"
	"koding/db/models"
	"koding/kites/kloud/contexthelper/session"
	"koding/kites/kloud/eventer"
	"koding/kites/kloud/klient"
	"koding/kites/kloud/kloud"
	"koding/kites/kloud/machinestate"
	"koding/kites/kloud/plans"
	"time"

	"golang.org/x/net/context"

	"github.com/koding/logging"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// Machine represents a single MongodDB document that represents a Koding
// Provider from the jMachines collection.
type Machine struct {
	Id          bson.ObjectId `bson:"_id" json:"-"`
	Label       string        `bson:"label"`
	Domain      string        `bson:"domain"`
	QueryString string        `bson:"queryString"`
	IpAddress   string        `bson:"ipAddress"`
	Assignee    struct {
		InProgress      bool      `bson:"inProgress"`
		AssignedAt      time.Time `bson:"assignedAt"`
		KlientMissingAt time.Time `bson:"klientMissingAt"`
	} `bson:"assignee"`
	Status struct {
		State      string    `bson:"state"`
		Reason     string    `bson:"reason"`
		ModifiedAt time.Time `bson:"modifiedAt"`
	} `bson:"status"`
	Provider   string    `bson:"provider"`
	Credential string    `bson:"credential"`
	CreatedAt  time.Time `bson:"createdAt"`
	Meta       struct {
		AlwaysOn     bool   `bson:"alwaysOn"`
		InstanceId   string `structs:"instanceId" bson:"instanceId"`
		InstanceType string `structs:"instance_type" bson:"instance_type"`
		InstanceName string `structs:"instanceName" bson:"instanceName"`
		Region       string `structs:"region" bson:"region"`
		StorageSize  int    `structs:"storage_size" bson:"storage_size"`
		SourceAmi    string `structs:"source_ami" bson:"source_ami"`
		SnapshotId   string `structs:"snapshotId" bson:"snapshotId"`
	} `bson:"meta"`
	Users  []models.Permissions `bson:"users"`
	Groups []models.Permissions `bson:"groups"`

	// internal fields, not availabile in MongoDB schema
	Username string                 `bson:"-"`
	User     *models.User           `bson:"-"`
	Payment  *plans.PaymentResponse `bson:"-"`
	Checker  plans.Checker          `bson:"-"`
	Session  *session.Session       `bson:"-"`
	Log      logging.Logger         `bson:"-"`
	locker   kloud.Locker           `bson:"-"`

	// cleanFuncs are a list of functions that are called when after a method
	// is finished
	cleanFuncs []func()
}

// runCleanupFunctions calls all cleanup functions and set the
// list to nil. Once called any other call will not have any
// effect.
func (m *Machine) runCleanupFunctions() {
	if m.cleanFuncs == nil {
		return
	}

	for _, fn := range m.cleanFuncs {
		fn()
	}

	m.cleanFuncs = nil
}

func (m *Machine) State() machinestate.State {
	return machinestate.States[m.Status.State]
}

func (m *Machine) PublicIpAddress() string {
	return m.IpAddress
}

// push pushes the given message to the eventer
func (m *Machine) push(msg string, percentage int, state machinestate.State) {
	if m.Session.Eventer != nil {
		m.Session.Eventer.Push(&eventer.Event{
			Message:    msg,
			Percentage: percentage,
			Status:     state,
		})
	}
}

// switchAWSRegion switches to the given AWS region. This should be only used when
// you know what to do, otherwiese never, never change the region of a machine.
func (m *Machine) switchAWSRegion(region string) error {
	m.Meta.InstanceId = "" // we neglect any previous instanceId
	m.QueryString = ""     //
	m.Meta.Region = "us-east-1"

	client, err := m.Session.AWSClients.Region("us-east-1")
	if err != nil {
		return err
	}
	m.Session.AWSClient.Client = client

	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"meta.instanceId": "",
				"queryString":     "",
				"meta.region":     "us-east-1",
			}},
		)
	})
}

// markAsNotInitialized marks the machine as NotInitialized by cleaning up all
// necessary fields and marking the VM as notinitialized so the User can build
// it again.
func (m *Machine) markAsNotInitialized() error {
	m.Log.Warning("Instance is not available. Marking it as NotInitialized")
	if err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"ipAddress":         "",
				"queryString":       "",
				"meta.instanceType": "",
				"meta.instanceName": "",
				"meta.instanceId":   "",
				"status.state":      machinestate.NotInitialized.String(),
				"status.modifiedAt": time.Now().UTC(),
				"status.reason":     "Machine is marked as NotInitialized",
			}},
		)
	}); err != nil {
		return err
	}

	m.IpAddress = ""
	m.QueryString = ""
	m.Meta.InstanceType = ""
	m.Meta.InstanceName = ""
	m.Meta.InstanceId = ""

	// so any State() method can return the correct status
	m.Status.State = machinestate.NotInitialized.String()
	return nil
}

func (m *Machine) markAsStopped() error {
	return m.markAsStoppedWithReason("Machine is stopped")
}

func (m *Machine) markAsStoppedWithReason(reason string) error {
	m.Log.Debug("Marking instance as stopped")
	if err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"ipAddress":         "",
				"status.state":      machinestate.Stopped.String(),
				"status.modifiedAt": time.Now().UTC(),
				"status.reason":     reason,
			}},
		)
	}); err != nil {
		return err
	}

	// so any State() method returns the correct status
	m.Status.State = machinestate.Stopped.String()
	m.IpAddress = ""
	return nil
}

func (m *Machine) updateStorageSize(size int) error {
	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{"meta.storage_size": size}},
		)
	})
}

func (m *Machine) isKlientReady() bool {
	m.Log.Debug("All finished, testing for klient connection IP [%s]", m.IpAddress)
	klientRef, err := klient.NewWithTimeout(m.Session.Kite, m.QueryString, time.Minute*5)
	if err != nil {
		m.Log.Warning("Connecting to remote Klient instance err: %s", err)
		return false
	}
	defer klientRef.Close()

	m.Log.Debug("Sending a ping message")
	if err := klientRef.Ping(); err != nil {
		m.Log.Debug("Sending a ping message err: %s", err)
		return false
	}

	return true
}

// Lock performs a Lock on this Machine
func (m *Machine) Lock() error {
	if !m.Id.Valid() {
		return kloud.NewError(kloud.ErrMachineIdMissing)
	}

	if m.locker == nil {
		return fmt.Errorf("Machine '%s' missing Locker", m.Id.Hex())
	}

	return m.locker.Lock(m.Id.Hex())
}

// Unlock performs an Unlock on this Machine instance
func (m *Machine) Unlock() error {
	if !m.Id.Valid() {
		return kloud.NewError(kloud.ErrMachineIdMissing)
	}

	if m.locker == nil {
		return fmt.Errorf("Machine '%s' missing Locker", m.Id.Hex())
	}

	// Unlock does not return an error
	m.locker.Unlock(m.Id.Hex())
	return nil
}

// klientIsNotMissing will unset the `assignee.klientMissingAt` value
// from the database, only if the Machine.Assignee.KlientMissingAt value
// has data. Therefor it is safe to call as frequently.
func (m *Machine) klientIsNotMissing() error {
	if m.Assignee.KlientMissingAt.IsZero() {
		return nil
	}

	m.Log.Debug("Clearing assignee.klientMissingAt")

	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$unset": bson.M{"assignee.klientMissingAt": ""}},
		)
	})
}

// stopIfKlientIsMissing will stop the current Machine X minutes after
// the `assignee.klientMissingAt` value. If the value does not exist in
// the databse, it will write it and return.
//
// Therefor, this method is expected be called as often as needed,
// and will shutdown the Machine if klient has been missing for too long.
func (m *Machine) stopIfKlientIsMissing(ctx context.Context) error {

	// If this is the first time Klient has been found missing,
	// set the missingat time and return
	if m.Assignee.KlientMissingAt.IsZero() {
		m.Log.Debug("Klient has been reported missing, recording this as the first time it went missing")

		return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
			return c.UpdateId(
				m.Id,
				bson.M{"$set": bson.M{"assignee.klientMissingAt": time.Now().UTC()}},
			)
		})
	}

	// If the klient has been missing less than X minutes, don't stop
	if time.Since(m.Assignee.KlientMissingAt) < time.Minute*5 {
		return nil
	}

	// lock so it doesn't interfere with others.
	err := m.Lock()
	defer m.Unlock()
	if err != nil {
		return err
	}

	// Hasta la vista, baby!
	m.Log.Info("======> STOP started (missing klient) <======")
	if err := m.Stop(ctx); err != nil {
		m.Log.Info("======> STOP failed (missing klient: %s) <======", err)
		return err
	}
	m.Log.Info("======> STOP finished (missing klient) <======")

	// Clear the klientMissingAt field, or we risk Stopping the user's
	// machine next time they run it, without waiting the proper X minute
	// timeout.
	m.klientIsNotMissing()

	return nil
}
<|endoftext|>


<commit_before>package koding

import (
	"fmt"
	"koding/db/models"
	"koding/kites/kloud/contexthelper/session"
	"koding/kites/kloud/eventer"
	"koding/kites/kloud/klient"
	"koding/kites/kloud/kloud"
	"koding/kites/kloud/machinestate"
	"koding/kites/kloud/plans"
	"time"

	"golang.org/x/net/context"

	"github.com/koding/logging"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// Machine represents a single MongodDB document that represents a Koding
// Provider from the jMachines collection.
type Machine struct {
	Id          bson.ObjectId `bson:"_id" json:"-"`
	Label       string        `bson:"label"`
	Domain      string        `bson:"domain"`
	QueryString string        `bson:"queryString"`
	IpAddress   string        `bson:"ipAddress"`
	Assignee    struct {
		InProgress      bool      `bson:"inProgress"`
		AssignedAt      time.Time `bson:"assignedAt"`
		KlientMissingAt time.Time `bson:"klientMissingAt"`
	} `bson:"assignee"`
	Status struct {
		State      string    `bson:"state"`
		Reason     string    `bson:"reason"`
		ModifiedAt time.Time `bson:"modifiedAt"`
	} `bson:"status"`
	Provider   string    `bson:"provider"`
	Credential string    `bson:"credential"`
	CreatedAt  time.Time `bson:"createdAt"`
	Meta       struct {
		AlwaysOn     bool   `bson:"alwaysOn"`
		InstanceId   string `structs:"instanceId" bson:"instanceId"`
		InstanceType string `structs:"instance_type" bson:"instance_type"`
		InstanceName string `structs:"instanceName" bson:"instanceName"`
		Region       string `structs:"region" bson:"region"`
		StorageSize  int    `structs:"storage_size" bson:"storage_size"`
		SourceAmi    string `structs:"source_ami" bson:"source_ami"`
		SnapshotId   string `structs:"snapshotId" bson:"snapshotId"`
	} `bson:"meta"`
	Users  []models.Permissions `bson:"users"`
	Groups []models.Permissions `bson:"groups"`

	// internal fields, not availabile in MongoDB schema
	Username string                 `bson:"-"`
	User     *models.User           `bson:"-"`
	Payment  *plans.PaymentResponse `bson:"-"`
	Checker  plans.Checker          `bson:"-"`
	Session  *session.Session       `bson:"-"`
	Log      logging.Logger         `bson:"-"`
	locker   kloud.Locker           `bson:"-"`

	// cleanFuncs are a list of functions that are called when after a method
	// is finished
	cleanFuncs []func()
}

// runCleanupFunctions calls all cleanup functions and set the
// list to nil. Once called any other call will not have any
// effect.
func (m *Machine) runCleanupFunctions() {
	if m.cleanFuncs == nil {
		return
	}

	for _, fn := range m.cleanFuncs {
		fn()
	}

	m.cleanFuncs = nil
}

func (m *Machine) State() machinestate.State {
	return machinestate.States[m.Status.State]
}

func (m *Machine) PublicIpAddress() string {
	return m.IpAddress
}

// push pushes the given message to the eventer
func (m *Machine) push(msg string, percentage int, state machinestate.State) {
	if m.Session.Eventer != nil {
		m.Session.Eventer.Push(&eventer.Event{
			Message:    msg,
			Percentage: percentage,
			Status:     state,
		})
	}
}

// switchAWSRegion switches to the given AWS region. This should be only used when
// you know what to do, otherwiese never, never change the region of a machine.
func (m *Machine) switchAWSRegion(region string) error {
	m.Meta.InstanceId = "" // we neglect any previous instanceId
	m.QueryString = ""     //
	m.Meta.Region = "us-east-1"

	client, err := m.Session.AWSClients.Region("us-east-1")
	if err != nil {
		return err
	}
	m.Session.AWSClient.Client = client

	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"meta.instanceId": "",
				"queryString":     "",
				"meta.region":     "us-east-1",
			}},
		)
	})
}

// markAsNotInitialized marks the machine as NotInitialized by cleaning up all
// necessary fields and marking the VM as notinitialized so the User can build
// it again.
func (m *Machine) markAsNotInitialized() error {
	m.Log.Warning("Instance is not available. Marking it as NotInitialized")
	if err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"ipAddress":         "",
				"queryString":       "",
				"meta.instanceType": "",
				"meta.instanceName": "",
				"meta.instanceId":   "",
				"status.state":      machinestate.NotInitialized.String(),
				"status.modifiedAt": time.Now().UTC(),
				"status.reason":     "Machine is marked as NotInitialized",
			}},
		)
	}); err != nil {
		return err
	}

	m.IpAddress = ""
	m.QueryString = ""
	m.Meta.InstanceType = ""
	m.Meta.InstanceName = ""
	m.Meta.InstanceId = ""

	// so any State() method can return the correct status
	m.Status.State = machinestate.NotInitialized.String()
	return nil
}

func (m *Machine) markAsStopped() error {
	return m.markAsStoppedWithReason("Machine is stopped")
}

func (m *Machine) markAsStoppedWithReason(reason string) error {
	m.Log.Debug("Marking instance as stopped")
	if err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"ipAddress":         "",
				"status.state":      machinestate.Stopped.String(),
				"status.modifiedAt": time.Now().UTC(),
				"status.reason":     reason,
			}},
		)
	}); err != nil {
		return err
	}

	// so any State() method returns the correct status
	m.Status.State = machinestate.Stopped.String()
	m.IpAddress = ""
	return nil
}

func (m *Machine) updateStorageSize(size int) error {
	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{"meta.storage_size": size}},
		)
	})
}

func (m *Machine) isKlientReady() bool {
	m.Log.Debug("All finished, testing for klient connection IP [%s]", m.IpAddress)
	klientRef, err := klient.NewWithTimeout(m.Session.Kite, m.QueryString, time.Minute*5)
	if err != nil {
		m.Log.Warning("Connecting to remote Klient instance err: %s", err)
		return false
	}
	defer klientRef.Close()

	m.Log.Debug("Sending a ping message")
	if err := klientRef.Ping(); err != nil {
		m.Log.Debug("Sending a ping message err: %s", err)
		return false
	}

	return true
}

// Lock performs a Lock on this Machine
func (m *Machine) Lock() error {
	if !m.Id.Valid() {
		return kloud.NewError(kloud.ErrMachineIdMissing)
	}

	if m.locker == nil {
		return fmt.Errorf("Machine '%s' missing Locker", m.Id.Hex())
	}

	return m.locker.Lock(m.Id.Hex())
}

// Unlock performs an Unlock on this Machine instance
func (m *Machine) Unlock() error {
	if !m.Id.Valid() {
		return kloud.NewError(kloud.ErrMachineIdMissing)
	}

	if m.locker == nil {
		return fmt.Errorf("Machine '%s' missing Locker", m.Id.Hex())
	}

	// Unlock does not return an error
	m.locker.Unlock(m.Id.Hex())
	return nil
}

// klientIsNotMissing will unset the `assignee.klientMissingAt` value
// from the database, only if the Machine.Assignee.KlientMissingAt value
// has data. Therefor it is safe to call as frequently.
func (m *Machine) klientIsNotMissing() error {
	if m.Assignee.KlientMissingAt.IsZero() {
		return nil
	}

	m.Log.Debug("Clearing assignee.klientMissingAt")

	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$unset": bson.M{"assignee.klientMissingAt": ""}},
		)
	})
}

// stopIfKlientIsMissing will stop the current Machine X minutes after
// the `assignee.klientMissingAt` value. If the value does not exist in
// the databse, it will write it and return.
//
// Therefor, this method is expected be called as often as needed,
// and will shutdown the Machine if klient has been missing for too long.
func (m *Machine) stopIfKlientIsMissing(ctx context.Context) error {

	// If this is the first time Klient has been found missing,
	// set the missingat time and return
	if m.Assignee.KlientMissingAt.IsZero() {
		m.Log.Debug("Klient has been reported missing, recording this as the first time it went missing")

		return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
			return c.UpdateId(
				m.Id,
				bson.M{"$set": bson.M{"assignee.klientMissingAt": time.Now().UTC()}},
			)
		})
	}

	// If the klient has been missing less than X minutes, don't stop
	if time.Since(m.Assignee.KlientMissingAt) < time.Minute*5 {
		return nil
	}

	// lock so it doesn't interfere with others.
	err := m.Lock()
	defer m.Unlock()
	if err != nil {
		return err
	}

	// Hasta la vista, baby!
	m.Log.Info("======> STOP started (missing klient) <======")
	if err := m.Stop(ctx); err != nil {
		m.Log.Info("======> STOP failed (missing klient: %s) <======", err)
		return err
	}
	m.Log.Info("======> STOP finished (missing klient) <======")

	// Clear the klientMissingAt field, or we risk Stopping the user's
	// machine next time they run it, without waiting the proper X minute
	// timeout.
	m.klientIsNotMissing()

	return nil
}
<commit_msg>kloud: Deferred klientIsNotMissing call<commit_after>package koding

import (
	"fmt"
	"koding/db/models"
	"koding/kites/kloud/contexthelper/session"
	"koding/kites/kloud/eventer"
	"koding/kites/kloud/klient"
	"koding/kites/kloud/kloud"
	"koding/kites/kloud/machinestate"
	"koding/kites/kloud/plans"
	"time"

	"golang.org/x/net/context"

	"github.com/koding/logging"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// Machine represents a single MongodDB document that represents a Koding
// Provider from the jMachines collection.
type Machine struct {
	Id          bson.ObjectId `bson:"_id" json:"-"`
	Label       string        `bson:"label"`
	Domain      string        `bson:"domain"`
	QueryString string        `bson:"queryString"`
	IpAddress   string        `bson:"ipAddress"`
	Assignee    struct {
		InProgress      bool      `bson:"inProgress"`
		AssignedAt      time.Time `bson:"assignedAt"`
		KlientMissingAt time.Time `bson:"klientMissingAt"`
	} `bson:"assignee"`
	Status struct {
		State      string    `bson:"state"`
		Reason     string    `bson:"reason"`
		ModifiedAt time.Time `bson:"modifiedAt"`
	} `bson:"status"`
	Provider   string    `bson:"provider"`
	Credential string    `bson:"credential"`
	CreatedAt  time.Time `bson:"createdAt"`
	Meta       struct {
		AlwaysOn     bool   `bson:"alwaysOn"`
		InstanceId   string `structs:"instanceId" bson:"instanceId"`
		InstanceType string `structs:"instance_type" bson:"instance_type"`
		InstanceName string `structs:"instanceName" bson:"instanceName"`
		Region       string `structs:"region" bson:"region"`
		StorageSize  int    `structs:"storage_size" bson:"storage_size"`
		SourceAmi    string `structs:"source_ami" bson:"source_ami"`
		SnapshotId   string `structs:"snapshotId" bson:"snapshotId"`
	} `bson:"meta"`
	Users  []models.Permissions `bson:"users"`
	Groups []models.Permissions `bson:"groups"`

	// internal fields, not availabile in MongoDB schema
	Username string                 `bson:"-"`
	User     *models.User           `bson:"-"`
	Payment  *plans.PaymentResponse `bson:"-"`
	Checker  plans.Checker          `bson:"-"`
	Session  *session.Session       `bson:"-"`
	Log      logging.Logger         `bson:"-"`
	locker   kloud.Locker           `bson:"-"`

	// cleanFuncs are a list of functions that are called when after a method
	// is finished
	cleanFuncs []func()
}

// runCleanupFunctions calls all cleanup functions and set the
// list to nil. Once called any other call will not have any
// effect.
func (m *Machine) runCleanupFunctions() {
	if m.cleanFuncs == nil {
		return
	}

	for _, fn := range m.cleanFuncs {
		fn()
	}

	m.cleanFuncs = nil
}

func (m *Machine) State() machinestate.State {
	return machinestate.States[m.Status.State]
}

func (m *Machine) PublicIpAddress() string {
	return m.IpAddress
}

// push pushes the given message to the eventer
func (m *Machine) push(msg string, percentage int, state machinestate.State) {
	if m.Session.Eventer != nil {
		m.Session.Eventer.Push(&eventer.Event{
			Message:    msg,
			Percentage: percentage,
			Status:     state,
		})
	}
}

// switchAWSRegion switches to the given AWS region. This should be only used when
// you know what to do, otherwiese never, never change the region of a machine.
func (m *Machine) switchAWSRegion(region string) error {
	m.Meta.InstanceId = "" // we neglect any previous instanceId
	m.QueryString = ""     //
	m.Meta.Region = "us-east-1"

	client, err := m.Session.AWSClients.Region("us-east-1")
	if err != nil {
		return err
	}
	m.Session.AWSClient.Client = client

	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"meta.instanceId": "",
				"queryString":     "",
				"meta.region":     "us-east-1",
			}},
		)
	})
}

// markAsNotInitialized marks the machine as NotInitialized by cleaning up all
// necessary fields and marking the VM as notinitialized so the User can build
// it again.
func (m *Machine) markAsNotInitialized() error {
	m.Log.Warning("Instance is not available. Marking it as NotInitialized")
	if err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"ipAddress":         "",
				"queryString":       "",
				"meta.instanceType": "",
				"meta.instanceName": "",
				"meta.instanceId":   "",
				"status.state":      machinestate.NotInitialized.String(),
				"status.modifiedAt": time.Now().UTC(),
				"status.reason":     "Machine is marked as NotInitialized",
			}},
		)
	}); err != nil {
		return err
	}

	m.IpAddress = ""
	m.QueryString = ""
	m.Meta.InstanceType = ""
	m.Meta.InstanceName = ""
	m.Meta.InstanceId = ""

	// so any State() method can return the correct status
	m.Status.State = machinestate.NotInitialized.String()
	return nil
}

func (m *Machine) markAsStopped() error {
	return m.markAsStoppedWithReason("Machine is stopped")
}

func (m *Machine) markAsStoppedWithReason(reason string) error {
	m.Log.Debug("Marking instance as stopped")
	if err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{
				"ipAddress":         "",
				"status.state":      machinestate.Stopped.String(),
				"status.modifiedAt": time.Now().UTC(),
				"status.reason":     reason,
			}},
		)
	}); err != nil {
		return err
	}

	// so any State() method returns the correct status
	m.Status.State = machinestate.Stopped.String()
	m.IpAddress = ""
	return nil
}

func (m *Machine) updateStorageSize(size int) error {
	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$set": bson.M{"meta.storage_size": size}},
		)
	})
}

func (m *Machine) isKlientReady() bool {
	m.Log.Debug("All finished, testing for klient connection IP [%s]", m.IpAddress)
	klientRef, err := klient.NewWithTimeout(m.Session.Kite, m.QueryString, time.Minute*5)
	if err != nil {
		m.Log.Warning("Connecting to remote Klient instance err: %s", err)
		return false
	}
	defer klientRef.Close()

	m.Log.Debug("Sending a ping message")
	if err := klientRef.Ping(); err != nil {
		m.Log.Debug("Sending a ping message err: %s", err)
		return false
	}

	return true
}

// Lock performs a Lock on this Machine
func (m *Machine) Lock() error {
	if !m.Id.Valid() {
		return kloud.NewError(kloud.ErrMachineIdMissing)
	}

	if m.locker == nil {
		return fmt.Errorf("Machine '%s' missing Locker", m.Id.Hex())
	}

	return m.locker.Lock(m.Id.Hex())
}

// Unlock performs an Unlock on this Machine instance
func (m *Machine) Unlock() error {
	if !m.Id.Valid() {
		return kloud.NewError(kloud.ErrMachineIdMissing)
	}

	if m.locker == nil {
		return fmt.Errorf("Machine '%s' missing Locker", m.Id.Hex())
	}

	// Unlock does not return an error
	m.locker.Unlock(m.Id.Hex())
	return nil
}

// klientIsNotMissing will unset the `assignee.klientMissingAt` value
// from the database, only if the Machine.Assignee.KlientMissingAt value
// has data. Therefor it is safe to call as frequently.
func (m *Machine) klientIsNotMissing() error {
	if m.Assignee.KlientMissingAt.IsZero() {
		return nil
	}

	m.Log.Debug("Clearing assignee.klientMissingAt")

	return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.UpdateId(
			m.Id,
			bson.M{"$unset": bson.M{"assignee.klientMissingAt": ""}},
		)
	})
}

// stopIfKlientIsMissing will stop the current Machine X minutes after
// the `assignee.klientMissingAt` value. If the value does not exist in
// the databse, it will write it and return.
//
// Therefor, this method is expected be called as often as needed,
// and will shutdown the Machine if klient has been missing for too long.
func (m *Machine) stopIfKlientIsMissing(ctx context.Context) error {

	// If this is the first time Klient has been found missing,
	// set the missingat time and return
	if m.Assignee.KlientMissingAt.IsZero() {
		m.Log.Debug("Klient has been reported missing, recording this as the first time it went missing")

		return m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
			return c.UpdateId(
				m.Id,
				bson.M{"$set": bson.M{"assignee.klientMissingAt": time.Now().UTC()}},
			)
		})
	}

	// If the klient has been missing less than X minutes, don't stop
	if time.Since(m.Assignee.KlientMissingAt) < time.Minute*5 {
		return nil
	}

	// lock so it doesn't interfere with others.
	err := m.Lock()

	defer func(m *Machine) {
		err := m.Unlock()
		if err != nil {
			m.Log.Error("Defer Error: Unlocking machine failed, " + err.Error())
		}
	}(m)

	// Check for a Lock error
	if err != nil {
		return err
	}

	// Clear the klientMissingAt field, or we risk Stopping the user's
	// machine next time they run it, without waiting the proper X minute
	// timeout.
	defer func(m *Machine) {
		err := m.klientIsNotMissing()
		if err != nil {
			m.Log.Error("Defer Error: Call to klientIsNotMissing failed, " + err.Error())
		}
	}(m)

	// Hasta la vista, baby!
	m.Log.Info("======> STOP started (missing klient) <======")
	if err := m.Stop(ctx); err != nil {
		m.Log.Info("======> STOP failed (missing klient: %s) <======", err)
		return err
	}
	m.Log.Info("======> STOP finished (missing klient) <======")

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"socialapi/workers/common/runner"
	"socialapi/workers/helper"
	"socialapi/workers/realtime/dispatcher/dispatcher"
	"socialapi/workers/realtime/models"
)

const Name = "Dispatcher"

func main() {
	r := runner.New(Name)
	if err := r.Init(); err != nil {
		fmt.Println(err)
		return
	}

	// create a realtime service provider instance.
	pubnub := models.NewPubNub(r.Conf.GateKeeper.Pubnub, r.Log)
	defer pubnub.Close()

	// later on broker support must be removed
	rmq := helper.NewRabbitMQ(r.Conf, r.Log)
	rmqConn, err := rmq.Connect("NewDispatcher")
	if err != nil {
		panic(err)
	}
	defer rmqConn.Conn().Close()

	broker := models.NewBroker(rmqConn, r.Log)

	c := dispatcher.NewController(rmqConn, pubnub, broker)

	r.SetContext(c)
	r.ListenFor("dispatcher_channel_updated", (*dispatcher.Controller).UpdateChannel)
	r.ListenFor("dispatcher_message_updated", (*dispatcher.Controller).UpdateMessage)
	r.ListenFor("dispatcher_notify_user", (*dispatcher.Controller).NotifyUser)
	r.ListenFor("api.channel_message_created", (*dispatcher.Controller).GrantMessagePublicAccess)
	r.ListenFor("event.channel_participant_removed_from_channel", (*dispatcher.Controller).RevokeChannelAccess)
	r.Listen()

	r.Wait()
}
<commit_msg>dispatcher: open new RMQ connection for sending message to broker<commit_after>package main

import (
	"fmt"
	"socialapi/workers/common/runner"
	"socialapi/workers/helper"
	"socialapi/workers/realtime/dispatcher/dispatcher"
	"socialapi/workers/realtime/models"
)

const Name = "Dispatcher"

func main() {
	r := runner.New(Name)
	if err := r.Init(); err != nil {
		fmt.Println(err)
		return
	}

	// create a realtime service provider instance.
	pubnub := models.NewPubNub(r.Conf.GateKeeper.Pubnub, r.Log)
	defer pubnub.Close()

	// later on broker support must be removed
	rmq := helper.NewRabbitMQ(r.Conf, r.Log)
	rmqConn, err := rmq.Connect("NewDispatcher")
	if err != nil {
		panic(err)
	}
	defer rmqConn.Conn().Close()

	// When we use the same RMQ connection for both, we received
	// 'Exception (504) Reason: "CHANNEL_ERROR - unexpected method in connection state running"'
	// error at some point. It needs debugging.
	rmqBroker := helper.NewRabbitMQ(r.Conf, r.Log)
	rmqBrokerConn, err := rmqBroker.Connect("NewBrokerDispatcher")
	if err != nil {
		panic(err)
	}
	defer rmqBrokerConn.Conn().Close()

	broker := models.NewBroker(rmqBrokerConn, r.Log)

	c := dispatcher.NewController(rmqConn, pubnub, broker)

	r.SetContext(c)
	r.ListenFor("dispatcher_channel_updated", (*dispatcher.Controller).UpdateChannel)
	r.ListenFor("dispatcher_message_updated", (*dispatcher.Controller).UpdateMessage)
	r.ListenFor("dispatcher_notify_user", (*dispatcher.Controller).NotifyUser)
	r.ListenFor("api.channel_message_created", (*dispatcher.Controller).GrantMessagePublicAccess)
	r.ListenFor("event.channel_participant_removed_from_channel", (*dispatcher.Controller).RevokeChannelAccess)
	r.Listen()

	r.Wait()
}
<|endoftext|>


<commit_before>// Copyright 2013 Alex Browne.  All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

// File options_test.go tests the different options that may be provided
// in model type declarations using struct tags

package zoom

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"reflect"
	"testing"
)

func TestInvalidOptionThrowsError(t *testing.T) {
	testingSetUp()
	testingTearDown()

	type invalid struct {
		Attr string `zoom:"index,poop"`
		DefaultData
	}
	Register(&invalid{})
	defer Unregister(&invalid{})
	if err := Register(&invalid{}); err == nil {
		t.Error("Expected error when registering struct with invalid tag")
	}
}

func TestIndexedPrimativesModelSpec(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	ms, found := modelSpecs["indexedPrimativesModel"]
	if !found {
		t.Error("Could not find modelSpec for indexedPrimativesModel")
	}
	numFields := ms.modelType.Elem().NumField()
	if len(ms.primativeIndexes) != numFields-1 {
		t.Errorf("Expected %d primative fields to be indexed but got %d", numFields-1, len(ms.primativeIndexes))
	} else {
		for i := 0; i < numFields; i++ {
			field := ms.modelType.Elem().Field(i)
			index := ms.primativeIndexes[field.Name]
			switch {
			case typeIsNumeric(field.Type):
				if index.iType != indexNumeric {
					t.Errorf("Expected iType to be numeric (%d) but got: %d", indexNumeric, index.iType)
				}
			case typeIsString(field.Type):
				if index.iType != indexAlpha {
					t.Errorf("Expected iType to be alpha (%d) but got: %d", indexAlpha, index.iType)
				}
			case typeIsBool(field.Type):
				if index.iType != indexBoolean {
					t.Errorf("Expected iType to be bool (%d) but got: %d", indexBoolean, index.iType)
				}
			}
		}
	}
}

func TestIndexedPointersModelSpec(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	ms, found := modelSpecs["indexedPointersModel"]
	if !found {
		t.Error("Could not find modelSpec for indexedPointersModel")
	}

	numFields := ms.modelType.Elem().NumField()
	if len(ms.pointerIndexes) != numFields-1 {
		t.Errorf("Expected %d primative fields to be indexed but got %d", numFields-1, len(ms.pointerIndexes))
	} else {
		// iterate throough each field and validate the index struct inside the modelSpec
		for i := 0; i < numFields; i++ {
			field := ms.modelType.Elem().Field(i)
			index := ms.pointerIndexes[field.Name]
			switch {
			case typeIsNumeric(field.Type):
				if index.iType != indexNumeric {
					t.Errorf("Expected iType to be numeric (%d) but got: %d", indexNumeric, index.iType)
				}
			case typeIsString(field.Type):
				if index.iType != indexAlpha {
					t.Errorf("Expected iType to be alpha (%d) but got: %d", indexAlpha, index.iType)
				}
			case typeIsBool(field.Type):
				fmt.Println("checking that boolean shit")
				if index.iType != indexBoolean {
					t.Errorf("Expected iType to be bool (%d) but got: %d", indexBoolean, index.iType)
				}
			}
		}
	}
}

func TestSaveIndexedPrimativesModel(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	ms, err := newIndexedPrimativesModels(1)
	if err != nil {
		t.Error(err)
	}
	m := ms[0]

	if err := Save(m); err != nil {
		t.Error(err)
	}

	conn := GetConn()
	defer conn.Close()

	spec, found := modelSpecs["indexedPrimativesModel"]
	if !found {
		t.Error("Could not find modelSpec for indexedPrimativesModel")
	}
	numFields := spec.modelType.Elem().NumField()

	// iterate through each field using reflection and validate that the index was set properly
	for i := 0; i < numFields; i++ {
		field := spec.modelType.Elem().Field(i)
		if field.Anonymous {
			continue // skip embedded structs
		}
		val := reflect.ValueOf(m).Elem().FieldByName(field.Name)
		switch {
		case typeIsNumeric(field.Type):
			validateNumericIndex(t, "indexedPrimativesModel", m.Id, field.Name, val, conn)
		case typeIsString(field.Type):
			validateAlphaIndex(t, "indexedPrimativesModel", m.Id, field.Name, val.String(), conn)
		case typeIsBool(field.Type):
			validateBooleanIndex(t, "indexedPrimativesModel", m.Id, field.Name, val.Bool(), conn)
		default:
			t.Errorf("Unexpected type %s in struct for %s", field.Type.String(), "indexedPrimativesModel")
		}
	}
}

func TestSaveIndexedPointersModel(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	ms, err := newIndexedPointersModels(1)
	if err != nil {
		t.Error(err)
	}
	m := ms[0]

	if err := Save(m); err != nil {
		t.Error(err)
	}

	conn := GetConn()
	defer conn.Close()

	spec, found := modelSpecs["indexedPointersModel"]
	if !found {
		t.Error("Could not find modelSpec for indexedPointersModel")
	}
	numFields := spec.modelType.Elem().NumField()

	// iterate through each field using reflection and validate that the index was set properly
	for i := 0; i < numFields; i++ {
		field := spec.modelType.Elem().Field(i)
		if field.Anonymous {
			continue // skip embedded structs
		}
		ptr := reflect.ValueOf(m).Elem().FieldByName(field.Name)
		val := ptr.Elem()
		switch {
		case typeIsNumeric(field.Type.Elem()):
			validateNumericIndex(t, "indexedPointersModel", m.Id, field.Name, val, conn)
		case typeIsString(field.Type.Elem()):
			validateAlphaIndex(t, "indexedPointersModel", m.Id, field.Name, val.String(), conn)
		case typeIsBool(field.Type.Elem()):
			validateBooleanIndex(t, "indexedPointersModel", m.Id, field.Name, val.Bool(), conn)
		default:
			t.Errorf("Unexpected type %s in struct for %s", field.Type.String(), "indexedPrimativesModel")
		}
	}
}

// validate index on a numeric type
func validateNumericIndex(t *testing.T, modelName string, modelId string, fieldName string, fieldValue reflect.Value, conn redis.Conn) {
	indexKey := modelName + ":" + fieldName
	score, err := convertNumericToFloat64(fieldValue)
	if err != nil {
		t.Error(err)
	}
	results, err := redis.Strings(conn.Do("ZRANGEBYSCORE", indexKey, score, score))
	if err != nil {
		t.Error(err)
	}
	if len(results) == 0 {
		t.Errorf("numeric index was not set.\nExpected %s:%s to have score %f\n", modelName, modelId, score)
	}
}

// validate index on a string type
func validateAlphaIndex(t *testing.T, modelName string, modelId string, fieldName string, fieldValue string, conn redis.Conn) {
	indexKey := modelName + ":" + fieldName
	memberKey := fieldValue + " " + modelId
	_, err := redis.Int(conn.Do("ZRANK", indexKey, memberKey))
	if err != nil {
		t.Errorf("alpha index was not set\nExpected to find member %s\n%s", memberKey, err)
	}
}

// validate index on a boolean type
func validateBooleanIndex(t *testing.T, modelName string, modelId string, fieldName string, fieldValue bool, conn redis.Conn) {
	var indexKey string
	if fieldValue == true {
		indexKey = modelName + ":" + fieldName + ":true"
	} else {
		indexKey = modelName + ":" + fieldName + ":false"
	}
	if found, err := redis.Bool(conn.Do("SISMEMBER", indexKey, modelId)); err != nil {
		t.Error(err)
	} else if !found {
		t.Errorf("bool index was not set\nExpected to find member %s in set %s\n", modelId, indexKey)
	}
}
<commit_msg>Add test for custom redis name and redis ignore<commit_after>// Copyright 2013 Alex Browne.  All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

// File options_test.go tests the different options that may be provided
// in model type declarations using struct tags

package zoom

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"reflect"
	"testing"
)

func TestRedisIgnoreOption(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	type ignored struct {
		Attr string `redis:"-"`
		DefaultData
	}
	Register(&ignored{})
	defer Unregister(&ignored{})

	// check the spec
	spec, found := modelSpecs["ignored"]
	if !found {
		t.Error("Could not find spec for model of type ignored")
	}
	if len(spec.primatives) != 0 {
		t.Errorf("Expected no primatives in model spec. Got %d", len(spec.primatives))
	}

	// save a new model and check redis
	m := &ignored{
		Attr: "this should be ignored",
	}
	if err := Save(m); err != nil {
		t.Error(err)
	}

	conn := GetConn()
	defer conn.Close()
	key := "ignored:" + m.Id
	gotAttr, err := redis.String(conn.Do("HGET", key, "Attr"))
	if err != nil && err != redis.ErrNil {
		t.Error(err)
	}
	if gotAttr != "" {
		t.Errorf("Expected empty attr but got: %s", gotAttr)
	}
}

func TestRedisNameOption(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	type customized struct {
		Attr string `redis:"a"`
		DefaultData
	}
	Register(&customized{})
	defer Unregister(&customized{})

	// check the spec
	spec, found := modelSpecs["customized"]
	if !found {
		t.Error("Could not find spec for model of type customized")
	}
	if spec.primatives["Attr"].redisName != "a" {
		t.Errorf("Redis name was incorrect.\nExpected: 'a'\nGot: %s\n", spec.primatives["Attr"].redisName)
	}

	// save a new model and check redis
	m := &customized{
		Attr: "test",
	}
	if err := Save(m); err != nil {
		t.Error(err)
	}

	conn := GetConn()
	defer conn.Close()
	key := "customized:" + m.Id
	gotAttr, err := redis.String(conn.Do("HGET", key, "a"))
	if err != nil {
		t.Error(err)
	}
	if gotAttr != "test" {
		t.Errorf("Expected 'test' but got: %s", gotAttr)
	}
}

func TestInvalidOptionThrowsError(t *testing.T) {
	testingSetUp()
	testingTearDown()

	type invalid struct {
		Attr string `zoom:"index,poop"`
		DefaultData
	}
	if err := Register(&invalid{}); err == nil {
		t.Error("Expected error when registering struct with invalid tag")
	}
	Unregister(&invalid{})
}

func TestIndexedPrimativesModelSpec(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	ms, found := modelSpecs["indexedPrimativesModel"]
	if !found {
		t.Error("Could not find modelSpec for indexedPrimativesModel")
	}
	numFields := ms.modelType.Elem().NumField()
	if len(ms.primativeIndexes) != numFields-1 {
		t.Errorf("Expected %d primative fields to be indexed but got %d", numFields-1, len(ms.primativeIndexes))
	} else {
		for i := 0; i < numFields; i++ {
			field := ms.modelType.Elem().Field(i)
			index := ms.primativeIndexes[field.Name]
			switch {
			case typeIsNumeric(field.Type):
				if index.iType != indexNumeric {
					t.Errorf("Expected iType to be numeric (%d) but got: %d", indexNumeric, index.iType)
				}
			case typeIsString(field.Type):
				if index.iType != indexAlpha {
					t.Errorf("Expected iType to be alpha (%d) but got: %d", indexAlpha, index.iType)
				}
			case typeIsBool(field.Type):
				if index.iType != indexBoolean {
					t.Errorf("Expected iType to be bool (%d) but got: %d", indexBoolean, index.iType)
				}
			}
		}
	}
}

func TestIndexedPointersModelSpec(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	ms, found := modelSpecs["indexedPointersModel"]
	if !found {
		t.Error("Could not find modelSpec for indexedPointersModel")
	}

	numFields := ms.modelType.Elem().NumField()
	if len(ms.pointerIndexes) != numFields-1 {
		t.Errorf("Expected %d primative fields to be indexed but got %d", numFields-1, len(ms.pointerIndexes))
	} else {
		// iterate throough each field and validate the index struct inside the modelSpec
		for i := 0; i < numFields; i++ {
			field := ms.modelType.Elem().Field(i)
			index := ms.pointerIndexes[field.Name]
			switch {
			case typeIsNumeric(field.Type):
				if index.iType != indexNumeric {
					t.Errorf("Expected iType to be numeric (%d) but got: %d", indexNumeric, index.iType)
				}
			case typeIsString(field.Type):
				if index.iType != indexAlpha {
					t.Errorf("Expected iType to be alpha (%d) but got: %d", indexAlpha, index.iType)
				}
			case typeIsBool(field.Type):
				fmt.Println("checking that boolean shit")
				if index.iType != indexBoolean {
					t.Errorf("Expected iType to be bool (%d) but got: %d", indexBoolean, index.iType)
				}
			}
		}
	}
}

func TestSaveIndexedPrimativesModel(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	ms, err := newIndexedPrimativesModels(1)
	if err != nil {
		t.Error(err)
	}
	m := ms[0]

	if err := Save(m); err != nil {
		t.Error(err)
	}

	conn := GetConn()
	defer conn.Close()

	spec, found := modelSpecs["indexedPrimativesModel"]
	if !found {
		t.Error("Could not find modelSpec for indexedPrimativesModel")
	}
	numFields := spec.modelType.Elem().NumField()

	// iterate through each field using reflection and validate that the index was set properly
	for i := 0; i < numFields; i++ {
		field := spec.modelType.Elem().Field(i)
		if field.Anonymous {
			continue // skip embedded structs
		}
		val := reflect.ValueOf(m).Elem().FieldByName(field.Name)
		switch {
		case typeIsNumeric(field.Type):
			validateNumericIndex(t, "indexedPrimativesModel", m.Id, field.Name, val, conn)
		case typeIsString(field.Type):
			validateAlphaIndex(t, "indexedPrimativesModel", m.Id, field.Name, val.String(), conn)
		case typeIsBool(field.Type):
			validateBooleanIndex(t, "indexedPrimativesModel", m.Id, field.Name, val.Bool(), conn)
		default:
			t.Errorf("Unexpected type %s in struct for %s", field.Type.String(), "indexedPrimativesModel")
		}
	}
}

func TestSaveIndexedPointersModel(t *testing.T) {
	testingSetUp()
	defer testingTearDown()

	ms, err := newIndexedPointersModels(1)
	if err != nil {
		t.Error(err)
	}
	m := ms[0]

	if err := Save(m); err != nil {
		t.Error(err)
	}

	conn := GetConn()
	defer conn.Close()

	spec, found := modelSpecs["indexedPointersModel"]
	if !found {
		t.Error("Could not find modelSpec for indexedPointersModel")
	}
	numFields := spec.modelType.Elem().NumField()

	// iterate through each field using reflection and validate that the index was set properly
	for i := 0; i < numFields; i++ {
		field := spec.modelType.Elem().Field(i)
		if field.Anonymous {
			continue // skip embedded structs
		}
		ptr := reflect.ValueOf(m).Elem().FieldByName(field.Name)
		val := ptr.Elem()
		switch {
		case typeIsNumeric(field.Type.Elem()):
			validateNumericIndex(t, "indexedPointersModel", m.Id, field.Name, val, conn)
		case typeIsString(field.Type.Elem()):
			validateAlphaIndex(t, "indexedPointersModel", m.Id, field.Name, val.String(), conn)
		case typeIsBool(field.Type.Elem()):
			validateBooleanIndex(t, "indexedPointersModel", m.Id, field.Name, val.Bool(), conn)
		default:
			t.Errorf("Unexpected type %s in struct for %s", field.Type.String(), "indexedPrimativesModel")
		}
	}
}

// validate index on a numeric type
func validateNumericIndex(t *testing.T, modelName string, modelId string, fieldName string, fieldValue reflect.Value, conn redis.Conn) {
	indexKey := modelName + ":" + fieldName
	score, err := convertNumericToFloat64(fieldValue)
	if err != nil {
		t.Error(err)
	}
	results, err := redis.Strings(conn.Do("ZRANGEBYSCORE", indexKey, score, score))
	if err != nil {
		t.Error(err)
	}
	if len(results) == 0 {
		t.Errorf("numeric index was not set.\nExpected %s:%s to have score %f\n", modelName, modelId, score)
	}
}

// validate index on a string type
func validateAlphaIndex(t *testing.T, modelName string, modelId string, fieldName string, fieldValue string, conn redis.Conn) {
	indexKey := modelName + ":" + fieldName
	memberKey := fieldValue + " " + modelId
	_, err := redis.Int(conn.Do("ZRANK", indexKey, memberKey))
	if err != nil {
		t.Errorf("alpha index was not set\nExpected to find member %s\n%s", memberKey, err)
	}
}

// validate index on a boolean type
func validateBooleanIndex(t *testing.T, modelName string, modelId string, fieldName string, fieldValue bool, conn redis.Conn) {
	var indexKey string
	if fieldValue == true {
		indexKey = modelName + ":" + fieldName + ":true"
	} else {
		indexKey = modelName + ":" + fieldName + ":false"
	}
	if found, err := redis.Bool(conn.Do("SISMEMBER", indexKey, modelId)); err != nil {
		t.Error(err)
	} else if !found {
		t.Errorf("bool index was not set\nExpected to find member %s in set %s\n", modelId, indexKey)
	}
}
<|endoftext|>