<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/elsonwu/goio"
	"github.com/gin-gonic/gin"
	"github.com/golang/glog"
)

var flagHost = flag.String("host", "127.0.0.1:9999", "the default server host")
var flagSSLHost = flag.String("sslhost", "", "the server host for https, it will override the host setting")
var flagAllowOrigin = flag.String("alloworigin", "*", "the host allow to cross site ajax")
var flagDebug = flag.Bool("debug", false, "enable debug mode or not")
var flagClientLifeCycle = flag.Int64("lifecycle", 30, "how many seconds of the client life cycle")
var flagClientMessageTimeout = flag.Int("messagetimeout", 15, "how many seconds of the client keep waiting for new messages")
var flagEnableHttps = flag.Bool("enablehttps", false, "enable https or not")
var flagDisableHttp = flag.Bool("disablehttp", false, "disable http and use https only")
var flagCertFile = flag.String("certfile", "", "certificate file path")
var flagKeyFile = flag.String("keyfile", "", "private file path")

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()

	goio.Rooms()
	goio.Users()
	goio.Clients()

	if *flagDebug {
		go func() {
			for {
				time.Sleep(3 * time.Second)
				log.Printf("rooms: %d, users: %d, clients: %d \n", goio.Rooms().Count(), goio.Users().Count(), goio.Clients().Count())
			}
		}()
	}

	goio.LifeCycle = *flagClientLifeCycle

	m := gin.New()
	m.Use(gin.Recovery())
	if *flagDebug {
		m.Use(gin.Logger())
	}

	m.Use(func(ctx *gin.Context) {
		ctx.Header("Content-Type", "text/plain; charset=utf-8")
		ctx.Header("Access-Control-Allow-Credentials", "true")
		ctx.Header("Access-Control-Allow-Methods", "GET,POST")
		ctx.Header("Powered-By", "Goio")

		if "" != *flagAllowOrigin {
			allowOrigins := strings.Split(*flagAllowOrigin, ",")
			for _, allowOrigin := range allowOrigins {
				ctx.Writer.Header().Add("Access-Control-Allow-Origin", allowOrigin)
			}
		}
	})

	if *flagDebug {
		m.GET("/test", func(ctx *gin.Context) {
			var userId string
			var roomId string
			for i := 0; i < 1000; i += 1 {
				userId = strconv.Itoa(i % 100)
				user := goio.Users().MustGet(userId)
				if user == nil {
					fmt.Printf("user %#v\n", user)
					return
				}

				client := goio.NewClient(user)

				roomId = strconv.Itoa(i % 100)
				room := goio.Rooms().MustGet(roomId)

				goio.AddRoomUser(room, client.User)
			}
		})
	}

	m.GET("/count", func(ctx *gin.Context) {
		ctx.String(200, fmt.Sprintf("rooms: %d, users: %d, clients: %d \n", goio.Rooms().Count(), goio.Users().Count(), goio.Clients().Count()))
	})

	// get user ids array from the given room
	m.GET("/room/users/:room_id", func(ctx *gin.Context) {
		roomId := ctx.Param("room_id")

		room := goio.Rooms().Get(roomId)
		if room == nil {
			ctx.String(200, "")
			return
		}

		ctx.String(200, strings.Join(room.UserIds(), ","))
	})

	m.GET("/user/data/:user_id/:key", func(ctx *gin.Context) {
		userId := ctx.Param("user_id")
		key := ctx.Param("key")

		user := goio.Users().Get(userId)
		if user == nil {
			ctx.String(200, "")
			return
		}

		ctx.String(200, user.GetData(key))
	})

	m.POST("/user/data/:client_id/:key", func(ctx *gin.Context) {
		val, err := ioutil.ReadAll(ctx.Request.Body)
		ctx.Request.Body.Close()

		if err != nil {
			ctx.String(500, err.Error())
			return
		}

		clientId := ctx.Param("client_id")
		key := ctx.Param("key")

		clt := goio.Clients().Get(clientId)
		if clt != nil && clt.User != nil {
			select {
			case clt.User.AddData <- goio.UserData{Key: key, Val: string(val)}:
			case <-time.After(time.Second):
			}
		}

		ctx.String(204, "")
	})

	m.POST("/online_status", func(ctx *gin.Context) {
		val, err := ioutil.ReadAll(ctx.Request.Body)
		ctx.Request.Body.Close()

		if err != nil {
			glog.V(1).Infof("online_status error %s\n", err.Error())
			ctx.String(500, err.Error())
			return
		}

		userIds := strings.Split(string(val), ",")
		glog.V(1).Infof("checking userIds:%s\n", string(val))

		status := ""
		for _, userId := range userIds {
			if nil == goio.Users().Get(userId) {
				status += "0,"
			} else {
				status += "1,"
			}
		}

		ctx.String(200, status)
	})

	m.POST("/client/:user_id", func(ctx *gin.Context) {
		userId := ctx.Param("user_id")
		user := goio.Users().MustGet(userId)
		clt := goio.NewClient(user)

		ctx.String(200, clt.Id)
	})

	m.GET("/kill_client/:client_id", func(ctx *gin.Context) {
		clt := goio.Clients().Get(ctx.Param("client_id"))
		if clt != nil {
			goio.DelUserClt(clt.User, clt)
		}

		ctx.String(204, "")
	})

	m.GET("/message/:client_id", func(ctx *gin.Context) {
		id := ctx.Param("client_id")
		clt := goio.Clients().Get(id)

		if clt == nil {
			ctx.String(404, fmt.Sprintf("Client %s does not exist\n", id))
			return
		}

		for i, wait := 0, 1; i*wait <= *flagClientMessageTimeout; i++ {

			msgs := clt.ReadMessages()
			if len(msgs) == 0 {
				time.Sleep(time.Duration(wait) * time.Second)
				continue
			}

			m, err := json.Marshal(msgs)
			if err != nil {
				ctx.String(500, err.Error())
				return
			}

			ctx.String(200, string(m))
			return
		}

		ctx.String(200, "")
	})

	m.POST("/message/:client_id", func(ctx *gin.Context) {
		id := ctx.Param("client_id")
		clt := goio.Clients().Get(id)
		if clt == nil {
			ctx.String(403, fmt.Sprintf("Client %s does not exist\n", id))
			return
		}

		if clt.User == nil {
			ctx.String(403, fmt.Sprintf("Client %s does not connect with any user\n", id))
			return
		}

		msg := &goio.Message{}
		ctx.BindJSON(msg)

		msg.CallerId = clt.User.Id
		msg.ClientId = clt.Id

		go func(msg *goio.Message, clt *goio.Client) {
			if *flagDebug {
				log.Printf("post msg: %#v\n", *msg)
			}

			glog.V(1).Infoln("user send msg - begin")

			switch msg.EventName {
			case "join":
				glog.V(1).Infof("msg event - join\n")
				if msg.RoomId != "" {
					r := goio.Rooms().MustGet(msg.RoomId)
					goio.AddRoomUser(r, clt.User)
					r.AddMessage(msg)
				}

			case "leave":
				glog.V(1).Infof("msg event - leave\n")
				if msg.RoomId != "" {
					r := goio.Rooms().Get(msg.RoomId)
					if r != nil {
						goio.DelRoomUser(r, clt.User)
						r.AddMessage(msg)
					}
				}
			case "broadcast":
				glog.V(1).Infoln("msg event - broadcast\n")

				if msg.RoomId != "" {
					r := goio.Rooms().Get(msg.RoomId)
					if r != nil {
						r.AddMessage(msg)
					}
				} else {
					for _, r := range clt.User.Rooms() {
						r.AddMessage(msg)
					}
				}

			default:
				glog.V(1).Infoln("unknown user msg")
			}

			glog.V(1).Infoln("user send msg - done")

		}(msg, clt)

		ctx.String(204, "")
	})

	m.OPTIONS("/*path", func(ctx *gin.Context) {})

	host := *flagHost
	if !*flagEnableHttps && *flagDisableHttp {
		log.Fatalln("You cannot disable http but not enable https in the same time")
	}

	//Prevent exiting
	ch := make(chan bool)

	if !*flagDisableHttp {
		go func() {
			log.Println("Serve at " + host + " - http")
			if err := http.ListenAndServe(host, m); err != nil {
				log.Println(err)
			}
		}()
	}

	if *flagEnableHttps {
		go func() {
			if *flagSSLHost != "" {
				host = *flagSSLHost
			}

			log.Println("Serve at " + host + " - https")
			if err := http.ListenAndServeTLS(host, *flagCertFile, *flagKeyFile, m); err != nil {
				log.Println(err)
			}
		}()
	}

	<-ch
}
<commit_msg>close each request after completed<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/elsonwu/goio"
	"github.com/gin-gonic/gin"
	"github.com/golang/glog"
)

var flagHost = flag.String("host", "127.0.0.1:9999", "the default server host")
var flagSSLHost = flag.String("sslhost", "", "the server host for https, it will override the host setting")
var flagAllowOrigin = flag.String("alloworigin", "*", "the host allow to cross site ajax")
var flagDebug = flag.Bool("debug", false, "enable debug mode or not")
var flagClientLifeCycle = flag.Int64("lifecycle", 30, "how many seconds of the client life cycle")
var flagClientMessageTimeout = flag.Int("messagetimeout", 15, "how many seconds of the client keep waiting for new messages")
var flagEnableHttps = flag.Bool("enablehttps", false, "enable https or not")
var flagDisableHttp = flag.Bool("disablehttp", false, "disable http and use https only")
var flagCertFile = flag.String("certfile", "", "certificate file path")
var flagKeyFile = flag.String("keyfile", "", "private file path")

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()

	goio.Rooms()
	goio.Users()
	goio.Clients()

	if *flagDebug {
		go func() {
			for {
				time.Sleep(3 * time.Second)
				log.Printf("rooms: %d, users: %d, clients: %d \n", goio.Rooms().Count(), goio.Users().Count(), goio.Clients().Count())
			}
		}()
	}

	goio.LifeCycle = *flagClientLifeCycle

	m := gin.New()
	m.Use(gin.Recovery())
	if *flagDebug {
		m.Use(gin.Logger())
	}

	m.Use(func(ctx *gin.Context) {
		defer ctx.Request.Body.Close()
		ctx.Next()
	})

	m.Use(func(ctx *gin.Context) {
		ctx.Header("Content-Type", "text/plain; charset=utf-8")
		ctx.Header("Access-Control-Allow-Credentials", "true")
		ctx.Header("Access-Control-Allow-Methods", "GET,POST")
		ctx.Header("Powered-By", "Goio")

		if "" != *flagAllowOrigin {
			allowOrigins := strings.Split(*flagAllowOrigin, ",")
			for _, allowOrigin := range allowOrigins {
				ctx.Writer.Header().Add("Access-Control-Allow-Origin", allowOrigin)
			}
		}
	})

	if *flagDebug {
		m.GET("/test", func(ctx *gin.Context) {
			var userId string
			var roomId string
			for i := 0; i < 1000; i += 1 {
				userId = strconv.Itoa(i % 100)
				user := goio.Users().MustGet(userId)
				if user == nil {
					fmt.Printf("user %#v\n", user)
					return
				}

				client := goio.NewClient(user)

				roomId = strconv.Itoa(i % 100)
				room := goio.Rooms().MustGet(roomId)

				goio.AddRoomUser(room, client.User)
			}
		})
	}

	m.GET("/count", func(ctx *gin.Context) {
		ctx.String(200, fmt.Sprintf("rooms: %d, users: %d, clients: %d \n", goio.Rooms().Count(), goio.Users().Count(), goio.Clients().Count()))
	})

	// get user ids array from the given room
	m.GET("/room/users/:room_id", func(ctx *gin.Context) {
		roomId := ctx.Param("room_id")

		room := goio.Rooms().Get(roomId)
		if room == nil {
			ctx.String(200, "")
			return
		}

		ctx.String(200, strings.Join(room.UserIds(), ","))
	})

	m.GET("/user/data/:user_id/:key", func(ctx *gin.Context) {
		userId := ctx.Param("user_id")
		key := ctx.Param("key")

		user := goio.Users().Get(userId)
		if user == nil {
			ctx.String(200, "")
			return
		}

		ctx.String(200, user.GetData(key))
	})

	m.POST("/user/data/:client_id/:key", func(ctx *gin.Context) {
		val, err := ioutil.ReadAll(ctx.Request.Body)
		ctx.Request.Body.Close()

		if err != nil {
			ctx.String(500, err.Error())
			return
		}

		clientId := ctx.Param("client_id")
		key := ctx.Param("key")

		clt := goio.Clients().Get(clientId)
		if clt != nil && clt.User != nil {
			select {
			case clt.User.AddData <- goio.UserData{Key: key, Val: string(val)}:
			case <-time.After(time.Second):
			}
		}

		ctx.String(204, "")
	})

	m.POST("/online_status", func(ctx *gin.Context) {
		val, err := ioutil.ReadAll(ctx.Request.Body)
		ctx.Request.Body.Close()

		if err != nil {
			glog.V(1).Infof("online_status error %s\n", err.Error())
			ctx.String(500, err.Error())
			return
		}

		userIds := strings.Split(string(val), ",")
		glog.V(1).Infof("checking userIds:%s\n", string(val))

		status := ""
		for _, userId := range userIds {
			if nil == goio.Users().Get(userId) {
				status += "0,"
			} else {
				status += "1,"
			}
		}

		ctx.String(200, status)
	})

	m.POST("/client/:user_id", func(ctx *gin.Context) {
		userId := ctx.Param("user_id")
		user := goio.Users().MustGet(userId)
		clt := goio.NewClient(user)

		ctx.String(200, clt.Id)
	})

	m.GET("/kill_client/:client_id", func(ctx *gin.Context) {
		clt := goio.Clients().Get(ctx.Param("client_id"))
		if clt != nil {
			goio.DelUserClt(clt.User, clt)
		}

		ctx.String(204, "")
	})

	m.GET("/message/:client_id", func(ctx *gin.Context) {
		id := ctx.Param("client_id")
		clt := goio.Clients().Get(id)

		if clt == nil {
			ctx.String(404, fmt.Sprintf("Client %s does not exist\n", id))
			return
		}

		for i, wait := 0, 1; i*wait <= *flagClientMessageTimeout; i++ {

			msgs := clt.ReadMessages()
			if len(msgs) == 0 {
				time.Sleep(time.Duration(wait) * time.Second)
				continue
			}

			m, err := json.Marshal(msgs)
			if err != nil {
				ctx.String(500, err.Error())
				return
			}

			ctx.String(200, string(m))
			return
		}

		ctx.String(200, "")
	})

	m.POST("/message/:client_id", func(ctx *gin.Context) {
		id := ctx.Param("client_id")
		clt := goio.Clients().Get(id)
		if clt == nil {
			ctx.String(403, fmt.Sprintf("Client %s does not exist\n", id))
			return
		}

		if clt.User == nil {
			ctx.String(403, fmt.Sprintf("Client %s does not connect with any user\n", id))
			return
		}

		msg := &goio.Message{}
		ctx.BindJSON(msg)

		msg.CallerId = clt.User.Id
		msg.ClientId = clt.Id

		go func(msg *goio.Message, clt *goio.Client) {
			if *flagDebug {
				log.Printf("post msg: %#v\n", *msg)
			}

			glog.V(1).Infoln("user send msg - begin")

			switch msg.EventName {
			case "join":
				glog.V(1).Infof("msg event - join\n")
				if msg.RoomId != "" {
					r := goio.Rooms().MustGet(msg.RoomId)
					goio.AddRoomUser(r, clt.User)
					r.AddMessage(msg)
				}

			case "leave":
				glog.V(1).Infof("msg event - leave\n")
				if msg.RoomId != "" {
					r := goio.Rooms().Get(msg.RoomId)
					if r != nil {
						goio.DelRoomUser(r, clt.User)
						r.AddMessage(msg)
					}
				}
			case "broadcast":
				glog.V(1).Infoln("msg event - broadcast\n")

				if msg.RoomId != "" {
					r := goio.Rooms().Get(msg.RoomId)
					if r != nil {
						r.AddMessage(msg)
					}
				} else {
					for _, r := range clt.User.Rooms() {
						r.AddMessage(msg)
					}
				}

			default:
				glog.V(1).Infoln("unknown user msg")
			}

			glog.V(1).Infoln("user send msg - done")

		}(msg, clt)

		ctx.String(204, "")
	})

	m.OPTIONS("/*path", func(ctx *gin.Context) {})

	host := *flagHost
	if !*flagEnableHttps && *flagDisableHttp {
		log.Fatalln("You cannot disable http but not enable https in the same time")
	}

	//Prevent exiting
	ch := make(chan bool)

	if !*flagDisableHttp {
		go func() {
			log.Println("Serve at " + host + " - http")
			if err := http.ListenAndServe(host, m); err != nil {
				log.Println(err)
			}
		}()
	}

	if *flagEnableHttps {
		go func() {
			if *flagSSLHost != "" {
				host = *flagSSLHost
			}

			log.Println("Serve at " + host + " - https")
			if err := http.ListenAndServeTLS(host, *flagCertFile, *flagKeyFile, m); err != nil {
				log.Println(err)
			}
		}()
	}

	<-ch
}
<|endoftext|>


<commit_before>package app

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/google/gnxi/utils/xpath"
	"github.com/manifoldco/promptui"
	"github.com/openconfig/goyang/pkg/yang"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

type pathGenOpts struct {
	search     bool
	withDescr  bool
	withTypes  bool
	withPrefix bool
	pathType   string
	stateOnly  bool
	configOnly bool
	json       bool
}

type generatedPath struct {
	Path        string `json:"path,omitempty"`
	Type        string `json:"type,omitempty"`
	Description string `json:"description,omitempty"`
	Default     string `json:"default,omitempty"`
	IsState     bool   `json:"is-state,omitempty"`
	Namespace   string `json:"namespace,omitempty"`
}

func (a *App) PathCmdRun(d, f, e []string, pgo pathGenOpts) error {
	err := a.GenerateYangSchema(d, f, e)
	if err != nil {
		return err
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	out := make(chan *generatedPath)
	gpaths := make([]*generatedPath, 0)

	go func(ctx context.Context, out chan *generatedPath) {
		for {
			select {
			case m, ok := <-out:
				if !ok {
					return
				}
				gpaths = append(gpaths, m)
			case <-ctx.Done():
				return
			}
		}
	}(ctx, out)

	collected := make([]*yang.Entry, 0, 256)
	for _, entry := range a.SchemaTree.Dir {
		collected = append(collected, collectSchemaNodes(entry, true)...)
	}
	for _, entry := range collected {
		if !pgo.stateOnly && !pgo.configOnly || pgo.stateOnly && pgo.configOnly {
			out <- a.generatePath(entry, pgo.withPrefix, pgo.pathType)
			continue
		}
		state := isState(entry)
		if state && pgo.stateOnly {
			out <- a.generatePath(entry, pgo.withPrefix, pgo.pathType)
			continue
		}
		if !state && pgo.configOnly {
			out <- a.generatePath(entry, pgo.withPrefix, pgo.pathType)
			continue
		}
	}
	close(out)
	sort.Slice(gpaths, func(i, j int) bool {
		return gpaths[i].Path < gpaths[j].Path
	})
	if pgo.json {
		b, err := json.MarshalIndent(gpaths, "", "  ")
		if err != nil {
			return err
		}
		fmt.Fprintln(os.Stdout, string(b))
		return nil
	}

	if len(gpaths) == 0 {
		return errors.New("no results found")
	}

	// regular print
	if !pgo.search {
		sb := new(strings.Builder)
		for _, gp := range gpaths {
			sb.Reset()
			sb.WriteString(gp.Path)
			if pgo.withTypes {
				sb.WriteString("\t(type=")
				sb.WriteString(gp.Type)
				sb.WriteString(")")
			}
			if pgo.withDescr {
				sb.WriteString("\n")
				sb.WriteString(indent("\t", gp.Description))
			}
			fmt.Fprintln(os.Stdout, sb.String())
		}
		return nil
	}
	// search
	paths := make([]string, 0, len(gpaths))
	for _, gp := range gpaths {
		paths = append(paths, gp.Path)
	}
	p := promptui.Select{
		Label:        "select path",
		Items:        paths,
		Size:         10,
		Stdout:       os.Stdout,
		HideSelected: true,
		Searcher: func(input string, index int) bool {
			kws := strings.Split(input, " ")
			result := true
			for _, kw := range kws {
				if strings.HasPrefix(kw, "!") {
					kw = strings.TrimLeft(kw, "!")
					if kw == "" {
						continue
					}
					result = result && !strings.Contains(paths[index], kw)
				} else {
					result = result && strings.Contains(paths[index], kw)
				}
			}
			return result
		},
		Keys: &promptui.SelectKeys{
			Prev:     promptui.Key{Code: promptui.KeyPrev, Display: promptui.KeyPrevDisplay},
			Next:     promptui.Key{Code: promptui.KeyNext, Display: promptui.KeyNextDisplay},
			PageUp:   promptui.Key{Code: promptui.KeyBackward, Display: promptui.KeyBackwardDisplay},
			PageDown: promptui.Key{Code: promptui.KeyForward, Display: promptui.KeyForwardDisplay},
			Search:   promptui.Key{Code: ':', Display: ":"},
		},
	}
	index, selected, err := p.Run()
	if err != nil {
		return err
	}
	fmt.Println(selected)
	fmt.Println(a.generateTypeInfo(collected[index]))

	return nil
}

func (a *App) PathPreRunE(cmd *cobra.Command, args []string) error {
	a.Config.SetLocalFlagsFromFile(cmd)
	if a.Config.PathSearch && a.Config.PathWithDescr {
		return errors.New("flags --search and --descr cannot be used together")
	}
	if a.Config.LocalFlags.PathPathType != "xpath" && a.Config.LocalFlags.PathPathType != "gnmi" {
		return errors.New("path-type must be one of 'xpath' or 'gnmi'")
	}
	return a.yangFilesPreProcessing()
}

func (a *App) PathRunE(cmd *cobra.Command, args []string) error {
	return a.PathCmdRun(
		a.Config.GlobalFlags.Dir,
		a.Config.GlobalFlags.File,
		a.Config.GlobalFlags.Exclude,
		pathGenOpts{
			search:     a.Config.LocalFlags.PathSearch,
			withDescr:  a.Config.LocalFlags.PathWithDescr,
			withTypes:  a.Config.LocalFlags.PathWithTypes,
			withPrefix: a.Config.LocalFlags.PathWithPrefix,
			pathType:   a.Config.LocalFlags.PathPathType,
			stateOnly:  a.Config.LocalFlags.PathState,
			configOnly: a.Config.LocalFlags.PathConfig,
		},
	)
}

func (a *App) InitPathFlags(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&a.Config.LocalFlags.PathPathType, "path-type", "", "xpath", "path type xpath or gnmi")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathWithDescr, "descr", "", false, "print leaf description")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathWithPrefix, "with-prefix", "", false, "include module/submodule prefix in path elements")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathWithTypes, "types", "", false, "print leaf type")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathSearch, "search", "", false, "search through path list")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathState, "state-only", "", false, "generate paths only for YANG leafs representing state data")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathConfig, "config-only", "", false, "generate paths only for YANG leafs representing config data")
	cmd.LocalFlags().VisitAll(func(flag *pflag.Flag) {
		a.Config.FileConfig.BindPFlag(fmt.Sprintf("%s-%s", cmd.Name(), flag.Name), flag)
	})
}

func collectSchemaNodes(e *yang.Entry, leafOnly bool) []*yang.Entry {
	if e == nil {
		return []*yang.Entry{}
	}
	collected := make([]*yang.Entry, 0, 128)
	for _, child := range e.Dir {
		collected = append(collected,
			collectSchemaNodes(child, leafOnly)...)
	}
	if e.Parent != nil {
		switch {
		case e.Dir == nil && e.ListAttr != nil: // leaf-list
			fallthrough
		case e.Dir == nil: // leaf
			collected = append(collected, e)
		case e.ListAttr != nil: // list
			fallthrough
		default: // container
			if !leafOnly {
				collected = append(collected, e)
			}
		}
	}
	return collected
}

func (a *App) generatePath(entry *yang.Entry, prefixTagging bool, pType string) *generatedPath {
	gp := new(generatedPath)
	for e := entry; e != nil && e.Parent != nil; e = e.Parent {
		if e.IsCase() || e.IsChoice() {
			continue
		}
		elementName := e.Name
		if prefixTagging && e.Prefix != nil {
			elementName = fmt.Sprintf("%s:%s", e.Prefix.Name, elementName)
		}
		if e.Key != "" {
			for _, k := range strings.Fields(e.Key) {
				if prefixTagging && e.Prefix != nil {
					k = fmt.Sprintf("%s:%s", e.Prefix.Name, k)
				}
				elementName = fmt.Sprintf("%s[%s=*]", elementName, k)
			}
		}
		gp.Path = fmt.Sprintf("/%s%s", elementName, gp.Path)
	}

	gp.Description = entry.Description
	gp.Type = entry.Type.Name
	gp.Default = entry.DefaultValue()
	gp.IsState = isState(entry)
	gp.Namespace = entry.Namespace().NName()
	if pType == "gnmi" {
		gnmiPath, err := xpath.ToGNMIPath(gp.Path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "path: %s could not be changed to gnmi format: %v\n", gp.Path, err)
		}
		gp.Path = gnmiPath.String()
	}
	return gp
}

func (a *App) generateTypeInfo(e *yang.Entry) string {
	if e == nil || e.Type == nil {
		return "unknown type"
	}
	t := e.Type
	rstr := fmt.Sprintf("- type: %s", t.Kind)
	switch t.Kind {
	case yang.Ybits:
		data := getAnnotation(e, "bits")
		if data != nil {
			rstr += fmt.Sprintf(" %v", data)
		}
	case yang.Yenum:
		data := getAnnotation(e, "enum")
		if data != nil {
			rstr += fmt.Sprintf(" %v", data)
		}
	case yang.Yleafref:
		rstr += fmt.Sprintf(" %q", t.Path)
	case yang.Yidentityref:
		rstr += fmt.Sprintf(" %q", t.IdentityBase.Name)
		if a.Config.LocalFlags.PathWithPrefix {
			data := getAnnotation(e, "prefix-qualified-identities")
			if data != nil {
				rstr += fmt.Sprintf(" %v", data)
			}
		} else {
			identities := make([]string, 0, 64)
			for i := range t.IdentityBase.Values {
				identities = append(identities, t.IdentityBase.Values[i].Name)
			}
			rstr += fmt.Sprintf(" %v", identities)
		}

	case yang.Yunion:
		unionlist := make([]string, 0, len(t.Type))
		for i := range t.Type {
			unionlist = append(unionlist, t.Type[i].Name)
		}
		rstr += fmt.Sprintf(" %v", unionlist)
	default:
	}
	rstr += "\n"

	if t.Root != nil {
		data := getAnnotation(e, "root.type")
		if data != nil && t.Kind.String() != data.(string) {
			rstr += fmt.Sprintf("- root.type: %v\n", data)
		}
	}
	if t.Units != "" {
		rstr += fmt.Sprintf("- units: %s\n", t.Units)
	}
	if t.Default != "" {
		rstr += fmt.Sprintf("- default: %q\n", t.Default)
	}
	if t.FractionDigits != 0 {
		rstr += fmt.Sprintf("- fraction-digits: %d\n", t.FractionDigits)
	}
	if len(t.Length) > 0 {
		rstr += fmt.Sprintf("- length: %s\n", t.Length)
	}
	if t.Kind == yang.YinstanceIdentifier && !t.OptionalInstance {
		rstr += "- required\n"
	}

	if len(t.Pattern) > 0 {
		rstr += fmt.Sprintf("- pattern: %s\n", strings.Join(t.Pattern, "|"))
	}
	b := yang.BaseTypedefs[t.Kind.String()].YangType
	if len(t.Range) > 0 && !t.Range.Equal(b.Range) {
		rstr += fmt.Sprintf("- range: %s\n", t.Range)
	}
	return rstr
}

func getAnnotation(entry *yang.Entry, name string) interface{} {
	if entry.Annotation != nil {
		data, ok := entry.Annotation[name]
		if ok {
			return data
		}
	}
	return nil
}

func isState(e *yang.Entry) bool {
	if e.Config == yang.TSFalse {
		return true
	}
	if e.Parent != nil {
		return isState(e.Parent)
	}
	return false
}
<commit_msg>collapse prefixes in generated paths<commit_after>package app

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/google/gnxi/utils/xpath"
	"github.com/karimra/gnmic/utils"
	"github.com/manifoldco/promptui"
	"github.com/openconfig/goyang/pkg/yang"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

type pathGenOpts struct {
	search     bool
	withDescr  bool
	withTypes  bool
	withPrefix bool
	pathType   string
	stateOnly  bool
	configOnly bool
	json       bool
}

type generatedPath struct {
	Path        string `json:"path,omitempty"`
	Type        string `json:"type,omitempty"`
	Description string `json:"description,omitempty"`
	Default     string `json:"default,omitempty"`
	IsState     bool   `json:"is-state,omitempty"`
	Namespace   string `json:"namespace,omitempty"`
}

func (a *App) PathCmdRun(d, f, e []string, pgo pathGenOpts) error {
	err := a.GenerateYangSchema(d, f, e)
	if err != nil {
		return err
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	out := make(chan *generatedPath)
	gpaths := make([]*generatedPath, 0)

	go func(ctx context.Context, out chan *generatedPath) {
		for {
			select {
			case m, ok := <-out:
				if !ok {
					return
				}
				gpaths = append(gpaths, m)
			case <-ctx.Done():
				return
			}
		}
	}(ctx, out)

	collected := make([]*yang.Entry, 0, 256)
	for _, entry := range a.SchemaTree.Dir {
		collected = append(collected, collectSchemaNodes(entry, true)...)
	}
	for _, entry := range collected {
		if !pgo.stateOnly && !pgo.configOnly || pgo.stateOnly && pgo.configOnly {
			out <- a.generatePath(entry, pgo.withPrefix, pgo.pathType)
			continue
		}
		state := isState(entry)
		if state && pgo.stateOnly {
			out <- a.generatePath(entry, pgo.withPrefix, pgo.pathType)
			continue
		}
		if !state && pgo.configOnly {
			out <- a.generatePath(entry, pgo.withPrefix, pgo.pathType)
			continue
		}
	}
	close(out)
	sort.Slice(gpaths, func(i, j int) bool {
		return gpaths[i].Path < gpaths[j].Path
	})
	if pgo.withPrefix {
		for _, gp := range gpaths {
			gp.Path = collapsePrefixes(gp.Path)
		}
	}
	if pgo.json {
		b, err := json.MarshalIndent(gpaths, "", "  ")
		if err != nil {
			return err
		}
		fmt.Fprintln(os.Stdout, string(b))
		return nil
	}

	if len(gpaths) == 0 {
		return errors.New("no results found")
	}

	// regular print
	if !pgo.search {
		sb := new(strings.Builder)
		for _, gp := range gpaths {
			sb.Reset()
			sb.WriteString(gp.Path)
			if pgo.withTypes {
				sb.WriteString("\t(type=")
				sb.WriteString(gp.Type)
				sb.WriteString(")")
			}
			if pgo.withDescr {
				sb.WriteString("\n")
				sb.WriteString(indent("\t", gp.Description))
			}
			fmt.Fprintln(os.Stdout, sb.String())
		}
		return nil
	}
	// search
	paths := make([]string, 0, len(gpaths))
	for _, gp := range gpaths {
		paths = append(paths, gp.Path)
	}
	p := promptui.Select{
		Label:        "select path",
		Items:        paths,
		Size:         10,
		Stdout:       os.Stdout,
		HideSelected: true,
		Searcher: func(input string, index int) bool {
			kws := strings.Split(input, " ")
			result := true
			for _, kw := range kws {
				if strings.HasPrefix(kw, "!") {
					kw = strings.TrimLeft(kw, "!")
					if kw == "" {
						continue
					}
					result = result && !strings.Contains(paths[index], kw)
				} else {
					result = result && strings.Contains(paths[index], kw)
				}
			}
			return result
		},
		Keys: &promptui.SelectKeys{
			Prev:     promptui.Key{Code: promptui.KeyPrev, Display: promptui.KeyPrevDisplay},
			Next:     promptui.Key{Code: promptui.KeyNext, Display: promptui.KeyNextDisplay},
			PageUp:   promptui.Key{Code: promptui.KeyBackward, Display: promptui.KeyBackwardDisplay},
			PageDown: promptui.Key{Code: promptui.KeyForward, Display: promptui.KeyForwardDisplay},
			Search:   promptui.Key{Code: ':', Display: ":"},
		},
	}
	index, selected, err := p.Run()
	if err != nil {
		return err
	}
	fmt.Println(selected)
	fmt.Println(a.generateTypeInfo(collected[index]))

	return nil
}

func (a *App) PathPreRunE(cmd *cobra.Command, args []string) error {
	a.Config.SetLocalFlagsFromFile(cmd)
	if a.Config.PathSearch && a.Config.PathWithDescr {
		return errors.New("flags --search and --descr cannot be used together")
	}
	if a.Config.LocalFlags.PathPathType != "xpath" && a.Config.LocalFlags.PathPathType != "gnmi" {
		return errors.New("path-type must be one of 'xpath' or 'gnmi'")
	}
	return a.yangFilesPreProcessing()
}

func (a *App) PathRunE(cmd *cobra.Command, args []string) error {
	return a.PathCmdRun(
		a.Config.GlobalFlags.Dir,
		a.Config.GlobalFlags.File,
		a.Config.GlobalFlags.Exclude,
		pathGenOpts{
			search:     a.Config.LocalFlags.PathSearch,
			withDescr:  a.Config.LocalFlags.PathWithDescr,
			withTypes:  a.Config.LocalFlags.PathWithTypes,
			withPrefix: a.Config.LocalFlags.PathWithPrefix,
			pathType:   a.Config.LocalFlags.PathPathType,
			stateOnly:  a.Config.LocalFlags.PathState,
			configOnly: a.Config.LocalFlags.PathConfig,
		},
	)
}

func (a *App) InitPathFlags(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&a.Config.LocalFlags.PathPathType, "path-type", "", "xpath", "path type xpath or gnmi")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathWithDescr, "descr", "", false, "print leaf description")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathWithPrefix, "with-prefix", "", false, "include module/submodule prefix in path elements")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathWithTypes, "types", "", false, "print leaf type")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathSearch, "search", "", false, "search through path list")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathState, "state-only", "", false, "generate paths only for YANG leafs representing state data")
	cmd.Flags().BoolVarP(&a.Config.LocalFlags.PathConfig, "config-only", "", false, "generate paths only for YANG leafs representing config data")
	cmd.LocalFlags().VisitAll(func(flag *pflag.Flag) {
		a.Config.FileConfig.BindPFlag(fmt.Sprintf("%s-%s", cmd.Name(), flag.Name), flag)
	})
}

func collectSchemaNodes(e *yang.Entry, leafOnly bool) []*yang.Entry {
	if e == nil {
		return []*yang.Entry{}
	}
	collected := make([]*yang.Entry, 0, 128)
	for _, child := range e.Dir {
		collected = append(collected,
			collectSchemaNodes(child, leafOnly)...)
	}
	if e.Parent != nil {
		switch {
		case e.Dir == nil && e.ListAttr != nil: // leaf-list
			fallthrough
		case e.Dir == nil: // leaf
			collected = append(collected, e)
		case e.ListAttr != nil: // list
			fallthrough
		default: // container
			if !leafOnly {
				collected = append(collected, e)
			}
		}
	}
	return collected
}

func (a *App) generatePath(entry *yang.Entry, prefixTagging bool, pType string) *generatedPath {
	gp := new(generatedPath)
	for e := entry; e != nil && e.Parent != nil; e = e.Parent {
		if e.IsCase() || e.IsChoice() {
			continue
		}
		elementName := e.Name
		if prefixTagging && e.Prefix != nil {
			elementName = fmt.Sprintf("%s:%s", e.Prefix.Name, elementName)
		}
		if e.Key != "" {
			for _, k := range strings.Fields(e.Key) {
				if prefixTagging && e.Prefix != nil {
					k = fmt.Sprintf("%s:%s", e.Prefix.Name, k)
				}
				elementName = fmt.Sprintf("%s[%s=*]", elementName, k)
			}
		}
		gp.Path = fmt.Sprintf("/%s%s", elementName, gp.Path)
	}

	gp.Description = entry.Description
	gp.Type = entry.Type.Name
	gp.Default = entry.DefaultValue()
	gp.IsState = isState(entry)
	gp.Namespace = entry.Namespace().NName()
	if pType == "gnmi" {
		gnmiPath, err := xpath.ToGNMIPath(gp.Path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "path: %s could not be changed to gnmi format: %v\n", gp.Path, err)
		}
		gp.Path = gnmiPath.String()
	}
	return gp
}

func (a *App) generateTypeInfo(e *yang.Entry) string {
	if e == nil || e.Type == nil {
		return "unknown type"
	}
	t := e.Type
	rstr := fmt.Sprintf("- type: %s", t.Kind)
	switch t.Kind {
	case yang.Ybits:
		data := getAnnotation(e, "bits")
		if data != nil {
			rstr += fmt.Sprintf(" %v", data)
		}
	case yang.Yenum:
		data := getAnnotation(e, "enum")
		if data != nil {
			rstr += fmt.Sprintf(" %v", data)
		}
	case yang.Yleafref:
		rstr += fmt.Sprintf(" %q", t.Path)
	case yang.Yidentityref:
		rstr += fmt.Sprintf(" %q", t.IdentityBase.Name)
		if a.Config.LocalFlags.PathWithPrefix {
			data := getAnnotation(e, "prefix-qualified-identities")
			if data != nil {
				rstr += fmt.Sprintf(" %v", data)
			}
		} else {
			identities := make([]string, 0, 64)
			for i := range t.IdentityBase.Values {
				identities = append(identities, t.IdentityBase.Values[i].Name)
			}
			rstr += fmt.Sprintf(" %v", identities)
		}

	case yang.Yunion:
		unionlist := make([]string, 0, len(t.Type))
		for i := range t.Type {
			unionlist = append(unionlist, t.Type[i].Name)
		}
		rstr += fmt.Sprintf(" %v", unionlist)
	default:
	}
	rstr += "\n"

	if t.Root != nil {
		data := getAnnotation(e, "root.type")
		if data != nil && t.Kind.String() != data.(string) {
			rstr += fmt.Sprintf("- root.type: %v\n", data)
		}
	}
	if t.Units != "" {
		rstr += fmt.Sprintf("- units: %s\n", t.Units)
	}
	if t.Default != "" {
		rstr += fmt.Sprintf("- default: %q\n", t.Default)
	}
	if t.FractionDigits != 0 {
		rstr += fmt.Sprintf("- fraction-digits: %d\n", t.FractionDigits)
	}
	if len(t.Length) > 0 {
		rstr += fmt.Sprintf("- length: %s\n", t.Length)
	}
	if t.Kind == yang.YinstanceIdentifier && !t.OptionalInstance {
		rstr += "- required\n"
	}

	if len(t.Pattern) > 0 {
		rstr += fmt.Sprintf("- pattern: %s\n", strings.Join(t.Pattern, "|"))
	}
	b := yang.BaseTypedefs[t.Kind.String()].YangType
	if len(t.Range) > 0 && !t.Range.Equal(b.Range) {
		rstr += fmt.Sprintf("- range: %s\n", t.Range)
	}
	return rstr
}

func getAnnotation(entry *yang.Entry, name string) interface{} {
	if entry.Annotation != nil {
		data, ok := entry.Annotation[name]
		if ok {
			return data
		}
	}
	return nil
}

func isState(e *yang.Entry) bool {
	if e.Config == yang.TSFalse {
		return true
	}
	if e.Parent != nil {
		return isState(e.Parent)
	}
	return false
}

// collapsePrefixes removes prefixes from path element names and keys
func collapsePrefixes(p string) string {
	gp, err := utils.ParsePath(p)
	if err != nil {
		return p
	}
	parentPrefix := ""
	for _, pe := range gp.Elem {
		currentPrefix, name := getPrefixElem(pe.Name)
		if parentPrefix == "" {
			// first elem
			parentPrefix = currentPrefix
		} else if currentPrefix == parentPrefix {
			pe.Name = name
		}
		for k, v := range pe.Key {
			kp, kn := getPrefixElem(k)
			if kp == parentPrefix {
				pe.Key[kn] = v
				delete(pe.Key, k)
			}
		}
	}
	return fmt.Sprintf("/%s", utils.GnmiPathToXPath(gp, false))
}

// takes a path element name or a key name
// and returns the prefix and name
func getPrefixElem(pe string) (string, string) {
	if pe == "" {
		return "", ""
	}
	pes := strings.SplitN(pe, ":", 2)
	if len(pes) > 1 {
		return pes[0], pes[1]
	}
	return "", pes[0]
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
)

func fileContains (filename string, contents string) bool{
	b, _ := ioutil.ReadFile(filename)
	return string(b) == contents
}

func TestReplaceInPath(t *testing.T) {
	c := &Substitute{
		searchPattern:  "foo",
		replacePattern: "bar",
		paths:          []string{"test1"},
	}
	ioutil.WriteFile("test1", []byte("foo"), 0644)
	defer os.Remove("test1")
	ioutil.WriteFile("test2", []byte("foo"), 0644)
	defer os.Remove("test2")

	exec.Command("git", "add", "test1", "test2").Run()
	defer exec.Command("git", "reset", "HEAD", "test1", "test2").Run()

	c.Run()
	if fileContains("test1", "foo") {
		t.Error("pattern not replaced in target")
	}
	if fileContains("test2", "bar") {
		t.Error("pattern replaced in non-target")
	}
}

func TestIgnoresFilesNotInGit(t *testing.T) {
	t.Skip("god i'm tired")
}

func TestFancyPatterns(t *testing.T) {
	t.Skip("I need to think of a good example first.")
}
<commit_msg>`go fmt`<commit_after>package main

import (
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
)

func fileContains(filename string, contents string) bool {
	b, _ := ioutil.ReadFile(filename)
	return string(b) == contents
}

func TestReplaceInPath(t *testing.T) {
	c := &Substitute{
		searchPattern:  "foo",
		replacePattern: "bar",
		paths:          []string{"test1"},
	}
	ioutil.WriteFile("test1", []byte("foo"), 0644)
	defer os.Remove("test1")
	ioutil.WriteFile("test2", []byte("foo"), 0644)
	defer os.Remove("test2")

	exec.Command("git", "add", "test1", "test2").Run()
	defer exec.Command("git", "reset", "HEAD", "test1", "test2").Run()

	c.Run()
	if fileContains("test1", "foo") {
		t.Error("pattern not replaced in target")
	}
	if fileContains("test2", "bar") {
		t.Error("pattern replaced in non-target")
	}
}

func TestIgnoresFilesNotInGit(t *testing.T) {
	t.Skip("god i'm tired")
}

func TestFancyPatterns(t *testing.T) {
	t.Skip("I need to think of a good example first.")
}
<|endoftext|>


<commit_before>package yandexexport

var CloudInitScript string = `#!/usr/bin/env bash
GetMetadata () {
  echo "$(curl -f -H "Metadata-Flavor: Google" http://169.254.169.254/computeMetadata/v1/instance/attributes/$1 2> /dev/null)"
}

GetInstanceId () {
  echo "$(curl -f -H "Metadata-Flavor: Google" http://169.254.169.254/computeMetadata/v1/instance/id 2> /dev/null)"
}

GetServiceAccountId () {
  yc compute instance get ${INSTANCE_ID} | grep service_account | cut -f2 -d' '
}

InstallYc () {
  curl -s https://storage.yandexcloud.net/yandexcloud-yc/install.sh | sudo bash -s -- -n -i /usr/local
}

InstallAwsCli () {
  curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
  unzip -o awscliv2.zip > /dev/null
  sudo ./aws/install
}

InstallPackages () {
  sudo apt-get update -qq && sudo apt-get install -y unzip jq qemu-utils
}

InstallTools () {
  InstallPackages
  InstallYc
  InstallAwsCli
}

IMAGE_ID=$(GetMetadata image_id)
INSTANCE_ID=$(GetInstanceId)
DISKNAME=${INSTANCE_ID}-toexport
PATHS=$(GetMetadata paths)
ZONE=$(GetMetadata zone)

Exit () {
  for i in ${PATHS}; do
    LOGDEST="${i}.exporter.log"
    echo "Uploading exporter log to ${LOGDEST}..."
    aws s3 --region ru-central1 --endpoint-url=https://storage.yandexcloud.net cp /var/log/syslog ${LOGDEST}
  done

  if [ $1 -ne 0 ]; then
	echo "Set metadata key 'cloud-init-status' to 'cloud-init-error' value"
    if ! yc compute instance update ${INSTANCE_ID} --metadata cloud-init-status=cloud-init-error ; then
	  echo "Failed to update metadata key 'cloud-init-status'."
	  exit 111
	fi
  fi

  exit $1
}

InstallTools

echo "####### Export configuration #######"
echo "Image ID - ${IMAGE_ID}"
echo "Instance ID - ${INSTANCE_ID}"
echo "Instance zone - ${ZONE}"
echo "Disk name - ${DISKNAME}"
echo "Export paths - ${PATHS}"
echo "####################################"

echo "Detect Service Account ID..."
SERVICE_ACCOUNT_ID=$(GetServiceAccountId)
echo "Use Service Account ID: ${SERVICE_ACCOUNT_ID}"

echo "Create static access key..."
SEC_json=$(yc iam access-key create --service-account-id ${SERVICE_ACCOUNT_ID} \
    --description "this key is for export image to storage" --format json)

if [ $? -ne 0 ]; then
  echo "Failed to create static access key."
  Exit 1
fi

echo "Setup env variables to access storage..."
eval "$(jq -r '@sh "export YC_SK_ID=\(.access_key.id); export AWS_ACCESS_KEY_ID=\(.access_key.key_id); export AWS_SECRET_ACCESS_KEY=\(.secret)"' <<<${SEC_json}  )"

echo "Check access to storage..."
if ! aws s3 --region ru-central1 --endpoint-url=https://storage.yandexcloud.net ls > /dev/null ; then
  echo "Failed to access storage."
  Exit 1
fi

echo "Creating disk from image to be exported..."
if ! yc compute disk create --name ${DISKNAME} --source-image-id ${IMAGE_ID} --zone ${ZONE}; then
  echo "Failed to create disk."
  Exit 1
fi

echo "Attaching disk..."
if ! yc compute instance attach-disk ${INSTANCE_ID} --disk-name ${DISKNAME} --device-name doexport --auto-delete ; then
  echo "Failed to attach disk."
  Exit 1
fi

echo "Dumping disk..."
if ! qemu-img convert -O qcow2 -o cluster_size=2M /dev/disk/by-id/virtio-doexport disk.qcow2 ; then
  echo "Failed to dump disk to qcow2 image."
  Exit 1
fi

echo "Detaching disk..."
if ! yc compute instance detach-disk ${INSTANCE_ID}  --disk-name ${DISKNAME} ; then
  echo "Failed to detach disk."
fi

FAIL=0
echo "Deleting disk..."
if ! yc compute disk delete --name ${DISKNAME} ; then
  echo "Failed to delete disk."
  FAIL=1
fi
for i in ${PATHS}; do
  echo "Uploading qcow2 disk image to ${i}..."
  if ! aws s3 --region ru-central1 --endpoint-url=https://storage.yandexcloud.net cp disk.qcow2 ${i}; then
    echo "Failed to upload image to ${i}."
    FAIL=1
  fi
done

echo "Delete static access key..."
if ! yc iam access-key delete ${YC_SK_ID} ; then
  echo "Failed to delete static access key."
  FAIL=1
fi

echo "Set metadata key 'cloud-init-status' to 'cloud-init-done' value"
if ! yc compute instance update ${INSTANCE_ID} --metadata cloud-init-status=cloud-init-done ; then
  echo "Failed to update metadata key to 'cloud-init-status'."
  Exit 1
fi

Exit ${FAIL}`
<commit_msg>verifying the access to a specific bucket<commit_after>package yandexexport

var CloudInitScript string = `#!/usr/bin/env bash
GetMetadata () {
  echo "$(curl -f -H "Metadata-Flavor: Google" http://169.254.169.254/computeMetadata/v1/instance/attributes/$1 2> /dev/null)"
}

GetInstanceId () {
  echo "$(curl -f -H "Metadata-Flavor: Google" http://169.254.169.254/computeMetadata/v1/instance/id 2> /dev/null)"
}

GetServiceAccountId () {
  yc compute instance get ${INSTANCE_ID} | grep service_account | cut -f2 -d' '
}

InstallYc () {
  curl -s https://storage.yandexcloud.net/yandexcloud-yc/install.sh | sudo bash -s -- -n -i /usr/local
}

InstallAwsCli () {
  curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
  unzip -o awscliv2.zip > /dev/null
  sudo ./aws/install
}

InstallPackages () {
  sudo apt-get update -qq && sudo apt-get install -y unzip jq qemu-utils
}

InstallTools () {
  InstallPackages
  InstallYc
  InstallAwsCli
}

IMAGE_ID=$(GetMetadata image_id)
INSTANCE_ID=$(GetInstanceId)
DISKNAME=${INSTANCE_ID}-toexport
PATHS=$(GetMetadata paths)
ZONE=$(GetMetadata zone)

Exit () {
  for i in ${PATHS}; do
    LOGDEST="${i}.exporter.log"
    echo "Uploading exporter log to ${LOGDEST}..."
    aws s3 --region ru-central1 --endpoint-url=https://storage.yandexcloud.net cp /var/log/syslog ${LOGDEST}
  done

  if [ $1 -ne 0 ]; then
	echo "Set metadata key 'cloud-init-status' to 'cloud-init-error' value"
    if ! yc compute instance update ${INSTANCE_ID} --metadata cloud-init-status=cloud-init-error ; then
	  echo "Failed to update metadata key 'cloud-init-status'."
	  exit 111
	fi
  fi

  exit $1
}

InstallTools

echo "####### Export configuration #######"
echo "Image ID - ${IMAGE_ID}"
echo "Instance ID - ${INSTANCE_ID}"
echo "Instance zone - ${ZONE}"
echo "Disk name - ${DISKNAME}"
echo "Export paths - ${PATHS}"
echo "####################################"

echo "Detect Service Account ID..."
SERVICE_ACCOUNT_ID=$(GetServiceAccountId)
echo "Use Service Account ID: ${SERVICE_ACCOUNT_ID}"

echo "Create static access key..."
SEC_json=$(yc iam access-key create --service-account-id ${SERVICE_ACCOUNT_ID} \
    --description "this key is for export image to storage" --format json)

if [ $? -ne 0 ]; then
  echo "Failed to create static access key."
  Exit 1
fi

echo "Setup env variables to access storage..."
eval "$(jq -r '@sh "export YC_SK_ID=\(.access_key.id); export AWS_ACCESS_KEY_ID=\(.access_key.key_id); export AWS_SECRET_ACCESS_KEY=\(.secret)"' <<<${SEC_json}  )"

for i in ${PATHS}; do
  bucket=$(echo ${i} | sed 's/\(s3:\/\/[^\/]*\).*/\1/')
  echo "Check access to storage: '${bucket}'..."
  if ! aws s3 --region ru-central1 --endpoint-url=https://storage.yandexcloud.net ls ${bucket} > /dev/null ; then
    echo "Failed to access storage: '${bucket}'."
    Exit 1
  fi
done

echo "Creating disk from image to be exported..."
if ! yc compute disk create --name ${DISKNAME} --source-image-id ${IMAGE_ID} --zone ${ZONE}; then
  echo "Failed to create disk."
  Exit 1
fi

echo "Attaching disk..."
if ! yc compute instance attach-disk ${INSTANCE_ID} --disk-name ${DISKNAME} --device-name doexport --auto-delete ; then
  echo "Failed to attach disk."
  Exit 1
fi

echo "Dumping disk..."
if ! qemu-img convert -O qcow2 -o cluster_size=2M /dev/disk/by-id/virtio-doexport disk.qcow2 ; then
  echo "Failed to dump disk to qcow2 image."
  Exit 1
fi

echo "Detaching disk..."
if ! yc compute instance detach-disk ${INSTANCE_ID}  --disk-name ${DISKNAME} ; then
  echo "Failed to detach disk."
fi

FAIL=0
echo "Deleting disk..."
if ! yc compute disk delete --name ${DISKNAME} ; then
  echo "Failed to delete disk."
  FAIL=1
fi
for i in ${PATHS}; do
  echo "Uploading qcow2 disk image to ${i}..."
  if ! aws s3 --region ru-central1 --endpoint-url=https://storage.yandexcloud.net cp disk.qcow2 ${i}; then
    echo "Failed to upload image to ${i}."
    FAIL=1
  fi
done

echo "Delete static access key..."
if ! yc iam access-key delete ${YC_SK_ID} ; then
  echo "Failed to delete static access key."
  FAIL=1
fi

echo "Set metadata key 'cloud-init-status' to 'cloud-init-done' value"
if ! yc compute instance update ${INSTANCE_ID} --metadata cloud-init-status=cloud-init-done ; then
  echo "Failed to update metadata key to 'cloud-init-status'."
  Exit 1
fi

Exit ${FAIL}`
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sync"
)

const DEFAULT_CONF_PATH string = "/etc/cloudpelican_supervisor.conf"

type Conf struct {
	data    map[string]string
	dataMux sync.RWMutex
}

func (c *Conf) Get(k string) string {
	// Locking in base function
	return c.GetOrDefault(k, "")
}

func (c *Conf) Set(k string, v string) {
	c.dataMux.Lock()
	defer c.dataMux.Unlock()
	c.data[k] = v
	log.Printf("Set conf %s=%s", k, v)
}

func (c *Conf) Save() bool {
	c.dataMux.RLock()
	defer c.dataMux.RUnlock()
	b, je := json.Marshal(c.data)
	if je != nil {
		log.Printf("Failed saving conf: %s", je)
		return false
	}
	if len(confPath) < 1 {
		confPath = DEFAULT_CONF_PATH
	}
	log.Printf("Writing conf to %s", confPath)
	we := ioutil.WriteFile(confPath, b, 0600)
	if we != nil {
		log.Printf("Failed saving conf: %s", we)
		return false
	}
	return true
}

func (c *Conf) GetNotEmpty(k string) string {
	// Locking in base function
	val := c.GetOrDefault(k, "")
	if len(val) < 1 {
		panic(fmt.Sprintf("Value %s empty", k))
	}
	return val
}

func (c *Conf) GetOrDefault(k string, d string) string {
	c.dataMux.RLock()
	defer c.dataMux.RUnlock()
	if len(c.data[k]) == 0 {
		return d
	}
	return c.data[k]
}

func newConf(path string) *Conf {
	c := &Conf{}

	// Default conf read?
	if len(path) < 1 {
		if _, err := os.Stat(DEFAULT_CONF_PATH); err != nil {
			if os.IsExist(err) {
				path = DEFAULT_CONF_PATH
				confPath = path
			}
		}
	}

	if len(path) > 0 {
		// Load file
		b, e := ioutil.ReadFile(path)
		if e != nil {
			log.Fatal(fmt.Sprintf("Failed to load conf: %s", e))
		}

		// Parse JSON
		var data map[string]string
		je := json.Unmarshal(b, &data)
		if je != nil {
			log.Fatal(fmt.Sprintf("Failed to parse conf: %s", je))
		}
		c.data = data
	} else {
		c.data = make(map[string]string)
	}
	return c
}
<commit_msg>Log and fix bug<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sync"
)

const DEFAULT_CONF_PATH string = "/etc/cloudpelican_supervisor.conf"

type Conf struct {
	data    map[string]string
	dataMux sync.RWMutex
}

func (c *Conf) Get(k string) string {
	// Locking in base function
	return c.GetOrDefault(k, "")
}

func (c *Conf) Set(k string, v string) {
	c.dataMux.Lock()
	defer c.dataMux.Unlock()
	c.data[k] = v
	log.Printf("Set conf %s=%s", k, v)
}

func (c *Conf) Save() bool {
	c.dataMux.RLock()
	defer c.dataMux.RUnlock()
	b, je := json.Marshal(c.data)
	if je != nil {
		log.Printf("Failed saving conf: %s", je)
		return false
	}
	if len(confPath) < 1 {
		confPath = DEFAULT_CONF_PATH
	}
	log.Printf("Writing conf to %s", confPath)
	we := ioutil.WriteFile(confPath, b, 0600)
	if we != nil {
		log.Printf("Failed saving conf: %s", we)
		return false
	}
	return true
}

func (c *Conf) GetNotEmpty(k string) string {
	// Locking in base function
	val := c.GetOrDefault(k, "")
	if len(val) < 1 {
		panic(fmt.Sprintf("Value %s empty", k))
	}
	return val
}

func (c *Conf) GetOrDefault(k string, d string) string {
	c.dataMux.RLock()
	defer c.dataMux.RUnlock()
	if len(c.data[k]) == 0 {
		return d
	}
	return c.data[k]
}

func newConf(path string) *Conf {
	c := &Conf{}

	// Default conf read?
	if len(path) < 1 {
		if _, err := os.Stat(DEFAULT_CONF_PATH); err == nil {
			path = DEFAULT_CONF_PATH
			confPath = path
			log.Printf("Default conf from %s", DEFAULT_CONF_PATH)
		}
	}

	if len(path) > 0 {
		// Load file
		b, e := ioutil.ReadFile(path)
		if e != nil {
			log.Fatal(fmt.Sprintf("Failed to load conf: %s", e))
		}

		// Parse JSON
		var data map[string]string
		je := json.Unmarshal(b, &data)
		if je != nil {
			log.Fatal(fmt.Sprintf("Failed to parse conf: %s", je))
		}
		c.data = data
	} else {
		c.data = make(map[string]string)
	}
	return c
}
<|endoftext|>


<commit_before>package immortal

import (
	//	"fmt"
	"os"
	"os/signal"
	"syscall"
	"testing"
)

type myFork struct{}

func (self myFork) Fork() {
	return
}

func TestHelperProcess(t *testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGHUP, syscall.SIGUSR1)
	select {
	case s := <-c:
		if s != syscall.SIGHUP {
			expect(t, syscall.SIGHUP, s)
		}
	}
}

func TestSupervisor(t *testing.T) {
	os.Setenv("GO_WANT_HELPER_PROCESS", "1")
	cfg := &Config{
		command: []string{"go", "test", "-run", "TestHelperProcess"},
	}
	d := Daemon{
		Config: cfg,
		Control: &Control{
			fifo:  make(chan Return),
			quit:  make(chan struct{}),
			state: make(chan error),
		},
		Forker: &myFork{},
		Logger: &LogWriter{
			logger: NewLogger(cfg),
		},
	}
	d.Run()
	//for {
	//select {
	//case s := <-d.Control.state:
	//t.Logf("%#v", s)

	//default:
	//fmt.Println(&d)
	//}
	//}
	sup := new(Sup)
	Supervise(sup, &d)
	sup.HandleSignals("h", &d)
}
<commit_msg>	modified:   supervisor_test.go<commit_after>package immortal

import (
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"testing"
)

type myFork struct{}

func (self myFork) Fork() {
	return
}

func TestHelperProcess(t *testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGHUP, syscall.SIGUSR1)
	select {
	case s := <-c:
		if s != syscall.SIGHUP {
			expect(t, syscall.SIGHUP, s)
		}
	}
}

func TestSupervisor(t *testing.T) {
	os.Setenv("GO_WANT_HELPER_PROCESS", "1")
	cfg := &Config{
		command: []string{"go", "test", "-run", "TestHelperProcess"},
	}
	d := &Daemon{
		Config: cfg,
		Control: &Control{
			fifo:  make(chan Return),
			quit:  make(chan struct{}),
			state: make(chan error),
		},
		Forker: &myFork{},
		Logger: &LogWriter{
			logger: NewLogger(cfg),
		},
	}
	d.Run()
	for {
		select {
		case s := <-d.Control.state:
			t.Logf("%#v", s)

		case fifo := <-d.Control.fifo:
			if fifo.err != nil {
				t.Error(fifo.err)
			}
		default:
			fmt.Printf("%#v", d.process.Pid)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/gorilla/mux"
	"io/ioutil"
	"net/http"
	"strings"
	"wayang/wayang"
)

type Config struct {
	HTTPPrefix      string `json:"http_prefix"`
	Database        string `json:"db"`
	DatabaseAddress string `json:"db_addr"`
	StaticConf      string `json:"static_conf"`
}

var (
	db wayang.DataStore
)

func main() {

	configPath := flag.String("conf", "./config.json", "Path to configuration file")
	listenAddr := flag.String("laddr", ":8000", "Address to listen on")
	flag.Parse()

	confData, err := ioutil.ReadFile(*configPath)
	if err != nil {
		fmt.Println("Failed to read configuration file:", *configPath)
		return
	}
	config := Config{}
	config.HTTPPrefix = ""

	err = json.Unmarshal(confData, &config)
	if err != nil {
		fmt.Println("Failed to parse configuration file:", *configPath)
		return
	}

	r := mux.NewRouter()
	r.StrictSlash(true)

	if config.Database == "mongodb" {

		db, err = wayang.NewMongoStore(config.DatabaseAddress)
		if err != nil {
			fmt.Println(err)
			return
		}

		r.HandleFunc(config.HTTPPrefix+"/", indexPage).Methods("GET")
		r.HandleFunc(config.HTTPPrefix+"/", newMock).Methods("POST")
		r.HandleFunc(config.HTTPPrefix+"/{id:[0-9a-z]+}", optionsHandlerRoot).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/{id:[0-9a-z]+}/{endpoint:[a-zA-Z0-9/]+}", optionsHandler).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/{id:[0-9a-z]+}", mockRespondRoot).Methods("GET", "POST", "PUT", "PATCH", "DELETE")
		r.HandleFunc(config.HTTPPrefix+"/{id:[0-9a-z]+}/{endpoint:[a-zA-Z0-9/]+}", mockRespond).Methods("GET", "POST", "PUT", "PATCH", "DELETE")

	} else if config.Database == "static" {

		mock := wayang.Mock{}
		mockData, err := ioutil.ReadFile(config.StaticConf)
		if err != nil {
			fmt.Println("Failed to read static configuration file:", config.StaticConf)
			return
		}
		err = json.Unmarshal(mockData, &mock)
		if err != nil {
			fmt.Println("Failed to parse static configuration file:", config.StaticConf)
			return
		}

		db, err = wayang.NewStaticStore(mock)
		if err != nil {
			fmt.Println(err)
			return
		}

		r.HandleFunc(config.HTTPPrefix+"/", optionsHandlerRoot).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/{endpoint:[a-zA-Z0-9/]+}", optionsHandler).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/", mockRespondRoot).Methods("GET", "POST", "PUT", "PATCH", "DELETE")
		r.HandleFunc(config.HTTPPrefix+"/{endpoint:[a-zA-Z0-9/]+}", mockRespond).Methods("GET", "POST", "PUT", "PATCH", "DELETE")

	}
	defer db.Close()

	http.Handle("/", r)
	http.ListenAndServe(*listenAddr, nil)

}

func accessControlAllow(rw http.ResponseWriter, req *http.Request) {
	origin := req.Header.Get("Origin")
	if origin == "" {
		rw.Header().Set("Access-Control-Allow-Origin", "*")
	} else {
		rw.Header().Set("Access-Control-Allow-Origin", origin)
	}
}

func indexPage(rw http.ResponseWriter, req *http.Request) {
	accessControlAllow(rw, req)
	helpMessage := `Hi!
To create an endpoint, do a POST request to the current URL.
The request should contain JSON of the following format:
{
	"/": {
		"GET": {
			"some_key": "some_val"
		},
		"POST": {
			"some_key": {
				"some_other_key": "some_val"
			}
		}
	},
	"/other_endpoint": {
		"DELETE": {
			"key": "value"
		}
	}
}`
	rw.Write([]byte(helpMessage))
}

type NewMockResponse struct {
	Status string `json:"status"`
	Detail string `json:"detail,omitempty"`
	URL    string `json:"url,omitempty"`
}

func newMock(rw http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	accessControlAllow(rw, req)
	rw.Header().Set("Content-Type", "application/json")
	resp := NewMockResponse{}

	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		resp.URL = ""
		resp.Status = http.StatusText(http.StatusInternalServerError)
		resp.Detail = err.Error()
		rv, _ := json.Marshal(resp)
		http.Error(rw, string(rv), http.StatusInternalServerError)
		return
	}
	mock := wayang.Mock{}
	err = json.Unmarshal(body, &mock)
	if err != nil {
		resp.URL = ""
		resp.Status = http.StatusText(http.StatusBadRequest)
		resp.Detail = err.Error()
		rv, _ := json.Marshal(resp)
		http.Error(rw, string(rv), http.StatusBadRequest)
		return
	}

	// Ensure that all routes begin with a /
	for k, v := range mock {
		if k[:1] != "/" {
			mock["/"+k] = v
			delete(mock, k)
		}
	}

	// Ensure that all requests are in the right format
	for k, v := range mock {
		for kk, vv := range v {
			KK := strings.ToUpper(kk)
			if KK != "GET" &&
				KK != "POST" &&
				KK != "PUT" &&
				KK != "PATCH" &&
				KK != "DELETE" {
				resp.URL = ""
				resp.Status = http.StatusText(http.StatusBadRequest)
				resp.Detail = "Invalid key " + kk + " in JSON."
				rv, _ := json.Marshal(resp)
				http.Error(rw, string(rv), http.StatusBadRequest)
				return
			}
			if kk != KK {
				mock[k][KK] = vv
				delete(mock[k], kk)
			}
		}
	}

	fmt.Println(mock)

	id, err := db.NewMock(mock)
	if err != nil {
		resp.URL = ""
		resp.Status = http.StatusText(http.StatusInternalServerError)
		resp.Detail = err.Error()
		rv, _ := json.Marshal(resp)
		http.Error(rw, string(rv), http.StatusInternalServerError)
		return
	} else {
		resp.Status = "200 OK"
	}
	resp.URL = "http://" + req.Host + "/" + id

	rv, _ := json.Marshal(resp)

	rw.Write(rv)
}

func optionsHandlerRoot(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	endpoint := "/"

	accessControlAllow(rw, req)
	processOptionsResponse(rw, req, id, endpoint)
}

func optionsHandler(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	endpoint := "/" + vars["endpoint"]

	accessControlAllow(rw, req)
	processOptionsResponse(rw, req, id, endpoint)
}

func processOptionsResponse(rw http.ResponseWriter, req *http.Request, id string, endpoint string) {

	defer req.Body.Close()

	accessControlAllow(rw, req)
	ep, err := db.GetEndpoint(id, endpoint)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	var allowedMethods []string
	allowedMethods = append(allowedMethods, "OPTIONS")
	for k, _ := range ep {
		allowedMethods = append(allowedMethods, k)
	}

	csv := strings.Join(allowedMethods, ",")
	rw.Header().Set("Allow", csv)
	rw.Header().Set("Access-Control-Allow-Methods", csv)

}

func mockRespondRoot(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	endpoint := "/"

	accessControlAllow(rw, req)
	processEndpointResponse(rw, req, id, endpoint)
}

func mockRespond(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	endpoint := "/" + vars["endpoint"]

	accessControlAllow(rw, req)
	processEndpointResponse(rw, req, id, endpoint)
}

func processEndpointResponse(rw http.ResponseWriter, req *http.Request, id string, endpoint string) {

	defer req.Body.Close()

	accessControlAllow(rw, req)
	ep, err := db.GetEndpoint(id, endpoint)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	rw.Header().Set("Content-Type", "application/json")
	mockResponse := ep[req.Method]

	if len(mockResponse) == 0 {
		http.Error(rw, http.StatusText(http.StatusNotImplemented), http.StatusNotImplemented)
		return
	}

	rv, _ := json.Marshal(mockResponse)
	rw.Write([]byte(rv))

}
<commit_msg>Change static_conf to db_addr, add OPTIONS to /.<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/gorilla/mux"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"wayang/wayang"
)

type Config struct {
	HTTPPrefix      string `json:"http_prefix"`
	Database        string `json:"db"`
	DatabaseAddress string `json:"db_addr"`
}

var (
	db wayang.DataStore
)

func main() {

	configPath := flag.String("conf", "./config.json", "Path to configuration file")
	listenAddr := flag.String("laddr", ":8000", "Address to listen on")
	flag.Parse()

	confData, err := ioutil.ReadFile(*configPath)
	if err != nil {
		fmt.Println("Failed to read configuration file:", *configPath)
		return
	}
	config := Config{}
	config.HTTPPrefix = ""

	err = json.Unmarshal(confData, &config)
	if err != nil {
		fmt.Println("Failed to parse configuration file:", *configPath)
		return
	}

	r := mux.NewRouter()
	r.StrictSlash(true)

	log.Println("Using", config.Database, "@", config.DatabaseAddress)

	if config.Database == "mongodb" {

		db, err = wayang.NewMongoStore(config.DatabaseAddress)
		if err != nil {
			fmt.Println(err)
			return
		}

		r.HandleFunc(config.HTTPPrefix+"/", rootGet).Methods("GET")
		r.HandleFunc(config.HTTPPrefix+"/", rootPost).Methods("POST")
		r.HandleFunc(config.HTTPPrefix+"/", rootOptions).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/{id:[0-9a-z]+}", optionsHandlerRoot).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/{id:[0-9a-z]+}/{endpoint:[a-zA-Z0-9/]+}", optionsHandler).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/{id:[0-9a-z]+}", mockRespondRoot).Methods("GET", "POST", "PUT", "PATCH", "DELETE")
		r.HandleFunc(config.HTTPPrefix+"/{id:[0-9a-z]+}/{endpoint:[a-zA-Z0-9/]+}", mockRespond).Methods("GET", "POST", "PUT", "PATCH", "DELETE")

	} else if config.Database == "static" {

		mock := wayang.Mock{}
		mockData, err := ioutil.ReadFile(config.DatabaseAddress)
		if err != nil {
			fmt.Println("Failed to read static configuration file:", config.DatabaseAddress)
			return
		}
		err = json.Unmarshal(mockData, &mock)
		if err != nil {
			fmt.Println("Failed to parse static configuration file:", config.DatabaseAddress)
			return
		}

		db, err = wayang.NewStaticStore(mock)
		if err != nil {
			fmt.Println(err)
			return
		}

		r.HandleFunc(config.HTTPPrefix+"/", optionsHandlerRoot).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/{endpoint:[a-zA-Z0-9/]+}", optionsHandler).Methods("OPTIONS")
		r.HandleFunc(config.HTTPPrefix+"/", mockRespondRoot).Methods("GET", "POST", "PUT", "PATCH", "DELETE")
		r.HandleFunc(config.HTTPPrefix+"/{endpoint:[a-zA-Z0-9/]+}", mockRespond).Methods("GET", "POST", "PUT", "PATCH", "DELETE")

	}
	defer db.Close()

	http.Handle("/", r)
	http.ListenAndServe(*listenAddr, nil)

}

func accessControlAllow(rw http.ResponseWriter, req *http.Request) {
	origin := req.Header.Get("Origin")
	if origin == "" {
		rw.Header().Set("Access-Control-Allow-Origin", "*")
		log.Println(req.Method, req.RequestURI)
	} else {
		rw.Header().Set("Access-Control-Allow-Origin", origin)
		log.Println(req.Method, req.RequestURI, "[Origin:", req.Header.Get("Origin")+"]")
	}
}

func rootOptions(rw http.ResponseWriter, req *http.Request) {

	defer req.Body.Close()

	accessControlAllow(rw, req)

	csv := "OPTIONS,GET,POST"
	rw.Header().Set("Allow", csv)
	rw.Header().Set("Access-Control-Allow-Methods", csv)

}

func rootGet(rw http.ResponseWriter, req *http.Request) {
	accessControlAllow(rw, req)
	helpMessage := `Hi!
To create an endpoint, do a POST request to the current URL.
The request should contain JSON of the following format:
{
	"/": {
		"GET": {
			"some_key": "some_val"
		},
		"POST": {
			"some_key": {
				"some_other_key": "some_val"
			}
		}
	},
	"/other_endpoint": {
		"DELETE": {
			"key": "value"
		}
	}
}`
	rw.Write([]byte(helpMessage))
}

type NewMockResponse struct {
	Status string `json:"status"`
	Detail string `json:"detail,omitempty"`
	URL    string `json:"url,omitempty"`
}

func rootPost(rw http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	accessControlAllow(rw, req)
	rw.Header().Set("Content-Type", "application/json")
	resp := NewMockResponse{}

	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		resp.URL = ""
		resp.Status = http.StatusText(http.StatusInternalServerError)
		resp.Detail = err.Error()
		rv, _ := json.Marshal(resp)
		http.Error(rw, string(rv), http.StatusInternalServerError)
		return
	}
	mock := wayang.Mock{}
	err = json.Unmarshal(body, &mock)
	if err != nil {
		resp.URL = ""
		resp.Status = http.StatusText(http.StatusBadRequest)
		resp.Detail = err.Error()
		rv, _ := json.Marshal(resp)
		http.Error(rw, string(rv), http.StatusBadRequest)
		return
	}

	// Ensure that all routes begin with a /
	for k, v := range mock {
		if k[:1] != "/" {
			mock["/"+k] = v
			delete(mock, k)
		}
	}

	// Ensure that all requests are in the right format
	for k, v := range mock {
		for kk, vv := range v {
			KK := strings.ToUpper(kk)
			if KK != "GET" &&
				KK != "POST" &&
				KK != "PUT" &&
				KK != "PATCH" &&
				KK != "DELETE" {
				resp.URL = ""
				resp.Status = http.StatusText(http.StatusBadRequest)
				resp.Detail = "Invalid key " + kk + " in JSON."
				rv, _ := json.Marshal(resp)
				http.Error(rw, string(rv), http.StatusBadRequest)
				return
			}
			if kk != KK {
				mock[k][KK] = vv
				delete(mock[k], kk)
			}
		}
	}

	fmt.Println(mock)

	id, err := db.NewMock(mock)
	if err != nil {
		resp.URL = ""
		resp.Status = http.StatusText(http.StatusInternalServerError)
		resp.Detail = err.Error()
		rv, _ := json.Marshal(resp)
		http.Error(rw, string(rv), http.StatusInternalServerError)
		return
	} else {
		resp.Status = "200 OK"
	}
	resp.URL = "http://" + req.Host + "/" + id

	rv, _ := json.Marshal(resp)

	rw.Write(rv)
}

func optionsHandlerRoot(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	endpoint := "/"

	processOptionsResponse(rw, req, id, endpoint)
}

func optionsHandler(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	endpoint := "/" + vars["endpoint"]

	processOptionsResponse(rw, req, id, endpoint)
}

func processOptionsResponse(rw http.ResponseWriter, req *http.Request, id string, endpoint string) {

	defer req.Body.Close()

	accessControlAllow(rw, req)
	ep, err := db.GetEndpoint(id, endpoint)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	var allowedMethods []string
	allowedMethods = append(allowedMethods, "OPTIONS")
	for k, _ := range ep {
		allowedMethods = append(allowedMethods, k)
	}

	csv := strings.Join(allowedMethods, ",")
	rw.Header().Set("Allow", csv)
	rw.Header().Set("Access-Control-Allow-Methods", csv)

}

func mockRespondRoot(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	endpoint := "/"

	processEndpointResponse(rw, req, id, endpoint)
}

func mockRespond(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	id := vars["id"]
	endpoint := "/" + vars["endpoint"]

	processEndpointResponse(rw, req, id, endpoint)
}

func processEndpointResponse(rw http.ResponseWriter, req *http.Request, id string, endpoint string) {

	defer req.Body.Close()

	accessControlAllow(rw, req)
	ep, err := db.GetEndpoint(id, endpoint)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	rw.Header().Set("Content-Type", "application/json")
	mockResponse := ep[req.Method]

	if len(mockResponse) == 0 {
		http.Error(rw, http.StatusText(http.StatusNotImplemented), http.StatusNotImplemented)
		return
	}

	rv, _ := json.Marshal(mockResponse)
	rw.Write([]byte(rv))

}
<|endoftext|>


<commit_before>package kubernetes

import (
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/extensions/v1beta1"

	"github.com/rusenask/keel/types"

	"testing"
)

type fakeImplementer struct {
	namespaces     *v1.NamespaceList
	deployment     *v1beta1.Deployment
	deploymentList *v1beta1.DeploymentList

	// stores value of an updated deployment
	updated *v1beta1.Deployment
}

func (i *fakeImplementer) Namespaces() (*v1.NamespaceList, error) {
	return i.namespaces, nil
}

func (i *fakeImplementer) Deployment(namespace, name string) (*v1beta1.Deployment, error) {
	return i.deployment, nil
}

func (i *fakeImplementer) Deployments(namespace string) (*v1beta1.DeploymentList, error) {
	return i.deploymentList, nil
}

func (i *fakeImplementer) Update(deployment *v1beta1.Deployment) error {
	i.updated = deployment
	return nil
}

func TestGetNamespaces(t *testing.T) {
	fi := &fakeImplementer{
		namespaces: &v1.NamespaceList{
			Items: []v1.Namespace{
				v1.Namespace{
					meta_v1.TypeMeta{},
					meta_v1.ObjectMeta{Name: "xxxx"},
					v1.NamespaceSpec{},
					v1.NamespaceStatus{},
				},
			},
		},
	}

	provider, err := NewProvider(fi)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	namespaces, err := provider.namespaces()
	if err != nil {
		t.Errorf("failed to get namespaces: %s", err)
	}

	if namespaces.Items[0].Name != "xxxx" {
		t.Errorf("expected xxxx but got %s", namespaces.Items[0].Name)
	}
}

func TestGetImageName(t *testing.T) {
	name := versionreg.ReplaceAllString("gcr.io/v2-namespace/hello-world:1.1", "")
	if name != "gcr.io/v2-namespace/hello-world" {
		t.Errorf("expected 'gcr.io/v2-namespace/hello-world' but got '%s'", name)
	}
}

func TestGetDeployments(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	deps, err := provider.deployments()
	if err != nil {
		t.Errorf("failed to get deployments: %s", err)
	}
	if len(deps) != 1 {
		t.Errorf("expected to find 1 deployment, got: %d", len(deps))
	}

	if deps[0].Items[0].GetName() != "dep-1" {
		t.Errorf("expected name %s, got %s", "dep-1", deps[0].Items[0].GetName())
	}
}

func TestGetImpacted(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-2",
					Namespace: "xxxx",
					Labels:    map[string]string{"whatever": "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	// creating "new version" event
	repo := &types.Repository{
		Name: "gcr.io/v2-namespace/hello-world",
		Tag:  "1.1.2",
	}

	deps, err := provider.impactedDeployments(repo)
	if err != nil {
		t.Errorf("failed to get deployments: %s", err)
	}

	if len(deps) != 1 {
		t.Errorf("expected to find 1 deployment but found %s", len(deps))
	}

	found := false
	for _, c := range deps[0].Spec.Template.Spec.Containers {

		containerImageName := versionreg.ReplaceAllString(c.Image, "")

		if containerImageName == repo.Name {
			found = true
		}
	}

	if !found {
		t.Errorf("couldn't find expected deployment in impacted deployment list")
	}

}

func TestProcessEvent(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "deployment-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "deployment-2",
					Namespace: "xxxx",
					Labels:    map[string]string{"whatever": "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/bye-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	repo := types.Repository{
		Name: "gcr.io/v2-namespace/hello-world",
		Tag:  "1.4.5",
	}

	event := &types.Event{Repository: repo}
	_, err = provider.processEvent(event)
	if err != nil {
		t.Errorf("got error while processing event: %s", err)
	}

	if fp.updated.Spec.Template.Spec.Containers[0].Image != repo.Name+":"+repo.Tag {
		t.Errorf("expected to find a deployment with updated image but found: %s", fp.updated.Spec.Template.Spec.Containers[0].Image)
	}
}

// Test to check how many deployments are "impacted" if we have sidecar container
func TestGetImpactedTwoContainersInSameDeployment(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
								v1.Container{
									Image: "gcr.io/v2-namespace/greetings-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-2",
					Namespace: "xxxx",
					Labels:    map[string]string{"whatever": "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	// creating "new version" event
	repo := &types.Repository{
		Name: "gcr.io/v2-namespace/hello-world",
		Tag:  "1.1.2",
	}

	deps, err := provider.impactedDeployments(repo)
	if err != nil {
		t.Errorf("failed to get deployments: %s", err)
	}

	if len(deps) != 1 {
		t.Errorf("expected to find 1 deployment but found %s", len(deps))
	}

	found := false
	for _, c := range deps[0].Spec.Template.Spec.Containers {

		containerImageName := versionreg.ReplaceAllString(c.Image, "")

		if containerImageName == repo.Name {
			found = true
		}
	}

	if !found {
		t.Errorf("couldn't find expected deployment in impacted deployment list")
	}

}
<commit_msg>created failing test<commit_after>package kubernetes

import (
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/extensions/v1beta1"

	"github.com/rusenask/keel/types"

	"testing"
)

type fakeImplementer struct {
	namespaces     *v1.NamespaceList
	deployment     *v1beta1.Deployment
	deploymentList *v1beta1.DeploymentList

	// stores value of an updated deployment
	updated *v1beta1.Deployment
}

func (i *fakeImplementer) Namespaces() (*v1.NamespaceList, error) {
	return i.namespaces, nil
}

func (i *fakeImplementer) Deployment(namespace, name string) (*v1beta1.Deployment, error) {
	return i.deployment, nil
}

func (i *fakeImplementer) Deployments(namespace string) (*v1beta1.DeploymentList, error) {
	return i.deploymentList, nil
}

func (i *fakeImplementer) Update(deployment *v1beta1.Deployment) error {
	i.updated = deployment
	return nil
}

func TestGetNamespaces(t *testing.T) {
	fi := &fakeImplementer{
		namespaces: &v1.NamespaceList{
			Items: []v1.Namespace{
				v1.Namespace{
					meta_v1.TypeMeta{},
					meta_v1.ObjectMeta{Name: "xxxx"},
					v1.NamespaceSpec{},
					v1.NamespaceStatus{},
				},
			},
		},
	}

	provider, err := NewProvider(fi)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	namespaces, err := provider.namespaces()
	if err != nil {
		t.Errorf("failed to get namespaces: %s", err)
	}

	if namespaces.Items[0].Name != "xxxx" {
		t.Errorf("expected xxxx but got %s", namespaces.Items[0].Name)
	}
}

func TestGetImageName(t *testing.T) {
	name := versionreg.ReplaceAllString("gcr.io/v2-namespace/hello-world:1.1", "")
	if name != "gcr.io/v2-namespace/hello-world" {
		t.Errorf("expected 'gcr.io/v2-namespace/hello-world' but got '%s'", name)
	}
}

func TestGetDeployments(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	deps, err := provider.deployments()
	if err != nil {
		t.Errorf("failed to get deployments: %s", err)
	}
	if len(deps) != 1 {
		t.Errorf("expected to find 1 deployment, got: %d", len(deps))
	}

	if deps[0].Items[0].GetName() != "dep-1" {
		t.Errorf("expected name %s, got %s", "dep-1", deps[0].Items[0].GetName())
	}
}

func TestGetImpacted(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-2",
					Namespace: "xxxx",
					Labels:    map[string]string{"whatever": "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	// creating "new version" event
	repo := &types.Repository{
		Name: "gcr.io/v2-namespace/hello-world",
		Tag:  "1.1.2",
	}

	deps, err := provider.impactedDeployments(repo)
	if err != nil {
		t.Errorf("failed to get deployments: %s", err)
	}

	if len(deps) != 1 {
		t.Errorf("expected to find 1 deployment but found %s", len(deps))
	}

	found := false
	for _, c := range deps[0].Spec.Template.Spec.Containers {

		containerImageName := versionreg.ReplaceAllString(c.Image, "")

		if containerImageName == repo.Name {
			found = true
		}
	}

	if !found {
		t.Errorf("couldn't find expected deployment in impacted deployment list")
	}

}

func TestProcessEvent(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "deployment-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "deployment-2",
					Namespace: "xxxx",
					Labels:    map[string]string{"whatever": "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/bye-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	repo := types.Repository{
		Name: "gcr.io/v2-namespace/hello-world",
		Tag:  "1.4.5",
	}

	event := &types.Event{Repository: repo}
	_, err = provider.processEvent(event)
	if err != nil {
		t.Errorf("got error while processing event: %s", err)
	}

	if fp.updated.Spec.Template.Spec.Containers[0].Image != repo.Name+":"+repo.Tag {
		t.Errorf("expected to find a deployment with updated image but found: %s", fp.updated.Spec.Template.Spec.Containers[0].Image)
	}
}

// Test to check how many deployments are "impacted" if we have sidecar container
func TestGetImpactedTwoContainersInSameDeployment(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
								v1.Container{
									Image: "gcr.io/v2-namespace/greetings-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-2",
					Namespace: "xxxx",
					Labels:    map[string]string{"whatever": "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	// creating "new version" event
	repo := &types.Repository{
		Name: "gcr.io/v2-namespace/hello-world",
		Tag:  "1.1.2",
	}

	deps, err := provider.impactedDeployments(repo)
	if err != nil {
		t.Errorf("failed to get deployments: %s", err)
	}

	if len(deps) != 1 {
		t.Errorf("expected to find 1 deployment but found %s", len(deps))
	}

	found := false
	for _, c := range deps[0].Spec.Template.Spec.Containers {

		containerImageName := versionreg.ReplaceAllString(c.Image, "")

		if containerImageName == repo.Name {
			found = true
		}
	}

	if !found {
		t.Errorf("couldn't find expected deployment in impacted deployment list")
	}

}

func TestGetImpactedTwoSameContainersInSameDeployment(t *testing.T) {
	fp := &fakeImplementer{}
	fp.namespaces = &v1.NamespaceList{
		Items: []v1.Namespace{
			v1.Namespace{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{Name: "xxxx"},
				v1.NamespaceSpec{},
				v1.NamespaceStatus{},
			},
		},
	}
	fp.deploymentList = &v1beta1.DeploymentList{
		Items: []v1beta1.Deployment{
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-1",
					Namespace: "xxxx",
					Labels:    map[string]string{types.KeelPolicyLabel: "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
			v1beta1.Deployment{
				meta_v1.TypeMeta{},
				meta_v1.ObjectMeta{
					Name:      "dep-2",
					Namespace: "xxxx",
					Labels:    map[string]string{"whatever": "all"},
				},
				v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Image: "gcr.io/v2-namespace/hello-world:1.1.1",
								},
							},
						},
					},
				},
				v1beta1.DeploymentStatus{},
			},
		},
	}

	provider, err := NewProvider(fp)
	if err != nil {
		t.Fatalf("failed to get provider: %s", err)
	}

	// creating "new version" event
	repo := &types.Repository{
		Name: "gcr.io/v2-namespace/hello-world",
		Tag:  "1.1.2",
	}

	deps, err := provider.impactedDeployments(repo)
	if err != nil {
		t.Errorf("failed to get deployments: %s", err)
	}

	if len(deps) != 1 {
		t.Errorf("expected to find 1 deployment but found %s", len(deps))
	}

	found := false
	for _, c := range deps[0].Spec.Template.Spec.Containers {

		containerImageName := versionreg.ReplaceAllString(c.Image, "")

		if containerImageName == repo.Name {
			found = true
		}
	}

	if !found {
		t.Errorf("couldn't find expected deployment in impacted deployment list")
	}

}
<|endoftext|>


<commit_before>/* Copyright (c) 2014-2015, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package storage

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log"
	"time"
)

const (
	// Length of the random hexadecimal ids assigned to pastes. At least 4.
	idSize = 8
	// Number of times to try getting an unused random paste id
	randTries = 10
	// How long to wait before retrying to delete a file
	deleteRetry = 2 * time.Minute
)

var (
	// ErrPasteNotFound means that we could not find the requested paste
	ErrPasteNotFound = errors.New("paste could not be found")
	// ErrNoUnusedIDFound means that we could not find an unused ID to
	// allocate to a new paste
	ErrNoUnusedIDFound = errors.New("gave up trying to find an unused random id")
)

// A Paste represents the paste's content and information
type Paste interface {
	io.Reader
	io.ReaderAt
	io.Seeker
	io.Closer
	ModTime() time.Time
	Size() int64
}

// ID is the binary representation of the identifier for a paste
type ID [idSize / 2]byte

// IDFromString parses a hexadecimal string into an ID. Returns the ID and an
// error, if any.
func IDFromString(hexID string) (id ID, err error) {
	if len(hexID) != idSize {
		return id, fmt.Errorf("invalid id at %s", hexID)
	}
	b, err := hex.DecodeString(hexID)
	if err != nil || len(b) != idSize/2 {
		return id, fmt.Errorf("invalid id at %s", hexID)
	}
	copy(id[:], b)
	return id, nil
}

func (id ID) String() string {
	return hex.EncodeToString(id[:])
}

// A Store represents a database holding multiple pastes identified by their
// ids
type Store interface {
	// Get the paste known by the given ID and an error, if any.
	Get(id ID) (Paste, error)

	// Put a new paste given its content. Will return the ID assigned to
	// the new paste and an error, if any.
	Put(content []byte) (ID, error)

	// Delete an existing paste by its ID. Will return an error, if any.
	Delete(id ID) error
}

func randomID(available func(ID) bool) (ID, error) {
	var id ID
	for try := 0; try < randTries; try++ {
		if _, err := rand.Read(id[:]); err != nil {
			continue
		}
		if available(id) {
			return id, nil
		}
	}
	return id, ErrNoUnusedIDFound
}

func SetupPasteDeletion(s Store, stats *Stats, id ID, size int64, after time.Duration) {
	if after == 0 {
		return
	}
	timer := time.NewTimer(after)
	go func() {
		for {
			<-timer.C
			if err := s.Delete(id); err == nil {
				stats.FreeSpace(size)
				break
			}
			log.Printf("Could not delete %s, will try again in %s", id, deleteRetry)
			timer.Reset(deleteRetry)
		}
	}()
}
<commit_msg>Use AfterFunc instead of a blocking goroutine for deletes<commit_after>/* Copyright (c) 2014-2015, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package storage

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log"
	"time"
)

const (
	// Length of the random hexadecimal ids assigned to pastes. At least 4.
	idSize = 8
	// Number of times to try getting an unused random paste id
	randTries = 10
	// How long to wait before retrying to delete a file
	deleteRetry = 2 * time.Minute
)

var (
	// ErrPasteNotFound means that we could not find the requested paste
	ErrPasteNotFound = errors.New("paste could not be found")
	// ErrNoUnusedIDFound means that we could not find an unused ID to
	// allocate to a new paste
	ErrNoUnusedIDFound = errors.New("gave up trying to find an unused random id")
)

// A Paste represents the paste's content and information
type Paste interface {
	io.Reader
	io.ReaderAt
	io.Seeker
	io.Closer
	ModTime() time.Time
	Size() int64
}

// ID is the binary representation of the identifier for a paste
type ID [idSize / 2]byte

// IDFromString parses a hexadecimal string into an ID. Returns the ID and an
// error, if any.
func IDFromString(hexID string) (id ID, err error) {
	if len(hexID) != idSize {
		return id, fmt.Errorf("invalid id at %s", hexID)
	}
	b, err := hex.DecodeString(hexID)
	if err != nil || len(b) != idSize/2 {
		return id, fmt.Errorf("invalid id at %s", hexID)
	}
	copy(id[:], b)
	return id, nil
}

func (id ID) String() string {
	return hex.EncodeToString(id[:])
}

// A Store represents a database holding multiple pastes identified by their
// ids
type Store interface {
	// Get the paste known by the given ID and an error, if any.
	Get(id ID) (Paste, error)

	// Put a new paste given its content. Will return the ID assigned to
	// the new paste and an error, if any.
	Put(content []byte) (ID, error)

	// Delete an existing paste by its ID. Will return an error, if any.
	Delete(id ID) error
}

func randomID(available func(ID) bool) (ID, error) {
	var id ID
	for try := 0; try < randTries; try++ {
		if _, err := rand.Read(id[:]); err != nil {
			continue
		}
		if available(id) {
			return id, nil
		}
	}
	return id, ErrNoUnusedIDFound
}

func SetupPasteDeletion(s Store, stats *Stats, id ID, size int64, after time.Duration) {
	if after == 0 {
		return
	}
	f := func() {
		del := func() error {
			if err := s.Delete(id); err != nil {
				return err
			}
			stats.FreeSpace(size)
			return nil
		}
		if err := del(); err == nil {
			return
		}
		timer := time.NewTimer(deleteRetry)
		for {
			log.Printf("Could not delete %s, will try again in %s", id, deleteRetry)
			<-timer.C
			if err := del(); err == nil {
				break
			}
			timer.Reset(deleteRetry)
		}
	}
	time.AfterFunc(after, f)
}
<|endoftext|>