<commit_before>package main

import (
	"encoding/binary"
	//"runtime/pprof"
	"bytes"
	"fmt"
	"github.com/gholt/brimstore"
	"github.com/gholt/brimutil"
	"os"
	"runtime"
	"sync"
	"time"
)

func main() {
	/*
	   cpuProf, err := os.Create("brimstore.cpu.pprof")
	   if err != nil {
	       panic(err)
	   }
	   pprof.StartCPUProfile(cpuProf)
	*/
	/*
	   blockPprof := pprof.Lookup("block")
	   runtime.SetBlockProfileRate(1)
	*/
	seed := int64(1)
	valueLength := 128
	fmt.Println(valueLength, "value length")
	targetBytes := 1 * 1024 * 1024
	fmt.Println(targetBytes, "target bytes")
	cores := runtime.GOMAXPROCS(0)
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
		cores = runtime.GOMAXPROCS(0)
	}
	fmt.Println(cores, "cores")
	clients := cores * cores
	fmt.Println(clients, "clients")
	keysPerClient := targetBytes / valueLength / clients
	fmt.Println(keysPerClient, "keys per client")
	totalKeys := clients * keysPerClient
	fmt.Println(totalKeys, "total keys")
	totalValueLength := totalKeys * valueLength
	fmt.Println(totalValueLength, "total value length")
	start := time.Now()
	wg := &sync.WaitGroup{}
	keys := make([][]byte, clients)
	wg.Add(clients)
	for i := 0; i < clients; i++ {
		keys[i] = make([]byte, keysPerClient*16)
		go func(i int) {
			brimutil.NewSeededScrambled(seed + int64(keysPerClient*i)).Read(keys[i])
			wg.Done()
		}(i)
	}
	wg.Wait()
	value := make([]byte, valueLength)
	brimutil.NewSeededScrambled(seed).Read(value)
	fmt.Println(time.Now().Sub(start), "to make keys and value")
	start = time.Now()
	speedStart := start
	s := brimstore.NewStore()
	s.Start()
	fmt.Println(time.Now().Sub(start), "to start store")
	start = time.Now()
	wg.Add(clients)
	for i := 0; i < clients; i++ {
		go func(keys []byte, seq uint64) {
			var err error
			w := &brimstore.WriteValue{
				Value:       value,
				WrittenChan: make(chan error, 1),
				Seq:         seq,
			}
			for o := 0; o < len(keys); o += 16 {
				w.KeyHashA = binary.LittleEndian.Uint64(keys[o:])
				w.KeyHashB = binary.LittleEndian.Uint64(keys[o+8:])
				w.Seq++
				s.Put(w)
				err = <-w.WrittenChan
				if err != nil {
					panic(err)
				}
			}
			wg.Done()
		}(keys[i], uint64(i*keysPerClient))
	}
	wg.Wait()
	fmt.Println(time.Now().Sub(start), "to add keys")
	start = time.Now()
	bytesWritten := s.Stop()
	speedStop := time.Now()
	fmt.Println(time.Now().Sub(start), "to stop store")
	fmt.Println(bytesWritten, "bytes written")
	seconds := float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	fmt.Printf("%.2fG/s based on total value length\n", float64(totalValueLength)/seconds/1024.0/1024.0/1024.0)
	fmt.Printf("%.2fG/s based on total bytes to disk\n", float64(bytesWritten)/seconds/1024.0/1024.0/1024.0)
	var st runtime.MemStats
	runtime.ReadMemStats(&st)
	fmt.Printf("%.2fG total alloc\n", float64(st.TotalAlloc)/1024/1024/1024)

	fmt.Println()
	start = time.Now()
	speedStart = start
	c := make([]chan int, clients)
	for i := 0; i < clients; i++ {
		c[i] = make(chan int)
		go func(keys []byte, c chan int) {
			m := 0
			for o := 0; o < len(keys); o += 16 {
				v, _ := s.Get(binary.LittleEndian.Uint64(keys[o:]), binary.LittleEndian.Uint64(keys[o+8:]))
				if v == nil {
					m++
				}
				if !bytes.Equal(v, value) {
					panic(fmt.Sprintf("%#v != %#v", string(v), string(value)))
				}
			}
			c <- m
		}(keys[i], c[i])
	}
	m := 0
	for i := 0; i < clients; i++ {
		m += <-c[i]
	}
	speedStop = time.Now()
	fmt.Println(time.Now().Sub(start), "to lookup keys")
	nanoseconds := speedStop.UnixNano() - speedStart.UnixNano()
	seconds = float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	fmt.Printf("%.2f key lookups per second\n", float64(totalKeys)/seconds)
	fmt.Printf("%.2fns per key lookup\n", float64(nanoseconds)/float64(totalKeys))
	fmt.Println(m, "keys missing")

	fmt.Println()
	start = time.Now()
	keys2 := make([][]byte, clients)
	wg.Add(clients)
	for i := 0; i < clients; i++ {
		keys2[i] = make([]byte, keysPerClient*16)
		go func(i int) {
			brimutil.NewSeededScrambled(seed + int64(totalKeys+keysPerClient*i)).Read(keys2[i])
			wg.Done()
		}(i)
	}
	wg.Wait()
	fmt.Println(time.Now().Sub(start), "to make second set of keys")
	start = time.Now()
	speedStart = start
	s.Start()
	fmt.Println(time.Now().Sub(start), "to restart store")
	start = time.Now()
	wg.Add(clients)
	for i := 0; i < clients; i++ {
		go func(keys []byte, seq uint64) {
			var err error
			w := &brimstore.WriteValue{
				Value:       value,
				WrittenChan: make(chan error, 1),
				Seq:         seq,
			}
			for o := 0; o < len(keys); o += 16 {
				w.KeyHashA = binary.LittleEndian.Uint64(keys[o:])
				w.KeyHashB = binary.LittleEndian.Uint64(keys[o+8:])
				w.Seq++
				s.Put(w)
				err = <-w.WrittenChan
				if err != nil {
					panic(err)
				}
			}
			wg.Done()
		}(keys2[i], uint64(i*keysPerClient))
		go func(keys []byte, c chan int) {
			m := 0
			for o := 0; o < len(keys); o += 16 {
				v, _ := s.Get(binary.LittleEndian.Uint64(keys[o:]), binary.LittleEndian.Uint64(keys[o+8:]))
				if v == nil {
					m++
				}
			}
			c <- m
		}(keys[i], c[i])
	}
	wg.Wait()
	m = 0
	for i := 0; i < clients; i++ {
		m += <-c[i]
	}
	fmt.Println(time.Now().Sub(start), "to add new keys while looking up old keys")
	start = time.Now()
	bytesWritten2 := s.Stop()
	speedStop = time.Now()
	fmt.Println(time.Now().Sub(start), "to stop store")
	fmt.Println(bytesWritten2-bytesWritten, "bytes written")
	nanoseconds = speedStop.UnixNano() - speedStart.UnixNano()
	seconds = float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	fmt.Printf("%.2fG/s based on total value length\n", float64(totalValueLength)/seconds/1024.0/1024.0/1024.0)
	fmt.Printf("%.2fG/s based on total bytes to disk\n", float64(bytesWritten2-bytesWritten)/seconds/1024.0/1024.0/1024.0)
	runtime.ReadMemStats(&st)
	fmt.Printf("%.2fG total alloc\n", float64(st.TotalAlloc)/1024/1024/1024)
	fmt.Printf("%.2f key lookups per second\n", float64(totalKeys)/seconds)
	fmt.Printf("%.2fns per key lookup\n", float64(nanoseconds)/float64(totalKeys))
	fmt.Println(m, "keys missing")

	fmt.Println()
	start = time.Now()
	speedStart = start
	for i := 0; i < clients; i++ {
		go func(keys []byte, c chan int) {
			m := 0
			for o := 0; o < len(keys); o += 16 {
				v, _ := s.Get(binary.LittleEndian.Uint64(keys[o:]), binary.LittleEndian.Uint64(keys[o+8:]))
				if v == nil {
					m++
				}
			}
			c <- m
		}(keys[i], c[i])
	}
	m = 0
	for i := 0; i < clients; i++ {
		m += <-c[i]
	}
	for i := 0; i < clients; i++ {
		go func(keys []byte, c chan int) {
			m := 0
			for o := 0; o < len(keys); o += 16 {
				v, _ := s.Get(binary.LittleEndian.Uint64(keys[o:]), binary.LittleEndian.Uint64(keys[o+8:]))
				if v == nil {
					m++
				}
			}
			c <- m
		}(keys2[i], c[i])
	}
	for i := 0; i < clients; i++ {
		m += <-c[i]
	}
	speedStop = time.Now()
	fmt.Println(time.Now().Sub(start), "to lookup both sets of keys")
	nanoseconds = speedStop.UnixNano() - speedStart.UnixNano()
	seconds = float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	fmt.Printf("%.2f key lookups per second\n", float64(totalKeys*2)/seconds)
	fmt.Printf("%.2fns per key lookup\n", float64(nanoseconds)/float64(totalKeys*2))
	fmt.Println(m, "keys missing")
	/*
	   f, err := os.Create("brimstore.blocking.pprof")
	   if err != nil {
	       panic(err)
	   }
	   blockPprof.WriteTo(f, 0)
	   f.Close()
	*/
	/*
	   pprof.StopCPUProfile()
	   cpuProf.Close()
	*/
}
<commit_msg>try to separate write timing from read timing<commit_after>package main

import (
	"encoding/binary"
	//"runtime/pprof"
	"bytes"
	"fmt"
	"github.com/gholt/brimstore"
	"github.com/gholt/brimutil"
	"os"
	"runtime"
	"sync"
	"time"
)

func main() {
	/*
	   cpuProf, err := os.Create("brimstore.cpu.pprof")
	   if err != nil {
	       panic(err)
	   }
	   pprof.StartCPUProfile(cpuProf)
	*/
	/*
	   blockPprof := pprof.Lookup("block")
	   runtime.SetBlockProfileRate(1)
	*/
	seed := int64(1)
	valueLength := 128
	fmt.Println(valueLength, "value length")
	targetBytes := 1 * 1024 * 1024
	fmt.Println(targetBytes, "target bytes")
	cores := runtime.GOMAXPROCS(0)
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
		cores = runtime.GOMAXPROCS(0)
	}
	fmt.Println(cores, "cores")
	clients := cores * cores
	fmt.Println(clients, "clients")
	keysPerClient := targetBytes / valueLength / clients
	fmt.Println(keysPerClient, "keys per client")
	totalKeys := clients * keysPerClient
	fmt.Println(totalKeys, "total keys")
	totalValueLength := totalKeys * valueLength
	fmt.Println(totalValueLength, "total value length")
	start := time.Now()
	wg := &sync.WaitGroup{}
	keys := make([][]byte, clients)
	wg.Add(clients)
	for i := 0; i < clients; i++ {
		keys[i] = make([]byte, keysPerClient*16)
		go func(i int) {
			brimutil.NewSeededScrambled(seed + int64(keysPerClient*i)).Read(keys[i])
			wg.Done()
		}(i)
	}
	wg.Wait()
	value := make([]byte, valueLength)
	brimutil.NewSeededScrambled(seed).Read(value)
	fmt.Println(time.Now().Sub(start), "to make keys and value")
	start = time.Now()
	speedStart := start
	s := brimstore.NewStore()
	s.Start()
	fmt.Println(time.Now().Sub(start), "to start store")
	start = time.Now()
	wg.Add(clients)
	for i := 0; i < clients; i++ {
		go func(keys []byte, seq uint64) {
			var err error
			w := &brimstore.WriteValue{
				Value:       value,
				WrittenChan: make(chan error, 1),
				Seq:         seq,
			}
			for o := 0; o < len(keys); o += 16 {
				w.KeyHashA = binary.LittleEndian.Uint64(keys[o:])
				w.KeyHashB = binary.LittleEndian.Uint64(keys[o+8:])
				w.Seq++
				s.Put(w)
				err = <-w.WrittenChan
				if err != nil {
					panic(err)
				}
			}
			wg.Done()
		}(keys[i], uint64(i*keysPerClient))
	}
	wg.Wait()
	fmt.Println(time.Now().Sub(start), "to add keys")
	start = time.Now()
	bytesWritten := s.Stop()
	speedStop := time.Now()
	fmt.Println(time.Now().Sub(start), "to stop store")
	fmt.Println(bytesWritten, "bytes written")
	seconds := float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	fmt.Printf("%.2fG/s based on total value length\n", float64(totalValueLength)/seconds/1024.0/1024.0/1024.0)
	fmt.Printf("%.2fG/s based on total bytes to disk\n", float64(bytesWritten)/seconds/1024.0/1024.0/1024.0)
	var st runtime.MemStats
	runtime.ReadMemStats(&st)
	fmt.Printf("%.2fG total alloc\n", float64(st.TotalAlloc)/1024/1024/1024)

	fmt.Println()
	start = time.Now()
	speedStart = start
	c := make([]chan int, clients)
	for i := 0; i < clients; i++ {
		c[i] = make(chan int)
		go func(keys []byte, c chan int) {
			m := 0
			for o := 0; o < len(keys); o += 16 {
				v, _ := s.Get(binary.LittleEndian.Uint64(keys[o:]), binary.LittleEndian.Uint64(keys[o+8:]))
				if v == nil {
					m++
				}
				if !bytes.Equal(v, value) {
					panic(fmt.Sprintf("%#v != %#v", string(v), string(value)))
				}
			}
			c <- m
		}(keys[i], c[i])
	}
	m := 0
	for i := 0; i < clients; i++ {
		m += <-c[i]
	}
	speedStop = time.Now()
	fmt.Println(time.Now().Sub(start), "to lookup keys")
	nanoseconds := speedStop.UnixNano() - speedStart.UnixNano()
	seconds = float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	fmt.Printf("%.2f key lookups per second\n", float64(totalKeys)/seconds)
	fmt.Printf("%.2fns per key lookup\n", float64(nanoseconds)/float64(totalKeys))
	fmt.Println(m, "keys missing")

	fmt.Println()
	start = time.Now()
	keys2 := make([][]byte, clients)
	wg.Add(clients)
	for i := 0; i < clients; i++ {
		keys2[i] = make([]byte, keysPerClient*16)
		go func(i int) {
			brimutil.NewSeededScrambled(seed + int64(totalKeys+keysPerClient*i)).Read(keys2[i])
			wg.Done()
		}(i)
	}
	wg.Wait()
	fmt.Println(time.Now().Sub(start), "to make second set of keys")
	start = time.Now()
	speedStart = start
	s.Start()
	fmt.Println(time.Now().Sub(start), "to restart store")
	start = time.Now()
	wg.Add(clients)
	for i := 0; i < clients; i++ {
		go func(keys []byte, seq uint64) {
			var err error
			w := &brimstore.WriteValue{
				Value:       value,
				WrittenChan: make(chan error, 1),
				Seq:         seq,
			}
			for o := 0; o < len(keys); o += 16 {
				w.KeyHashA = binary.LittleEndian.Uint64(keys[o:])
				w.KeyHashB = binary.LittleEndian.Uint64(keys[o+8:])
				w.Seq++
				s.Put(w)
				err = <-w.WrittenChan
				if err != nil {
					panic(err)
				}
			}
			wg.Done()
		}(keys2[i], uint64(i*keysPerClient))
		go func(keys []byte, c chan int) {
			m := 0
			for o := 0; o < len(keys); o += 16 {
				v, _ := s.Get(binary.LittleEndian.Uint64(keys[o:]), binary.LittleEndian.Uint64(keys[o+8:]))
				if v == nil {
					m++
				}
			}
			c <- m
		}(keys[i], c[i])
	}
	wg.Wait()
	fmt.Println(time.Now().Sub(start), "to add new keys while looking up old keys")
	start = time.Now()
	bytesWritten2 := s.Stop()
	speedStop = time.Now()
	fmt.Println(time.Now().Sub(start), "to stop store")
	fmt.Println(bytesWritten2-bytesWritten, "bytes written")
	nanoseconds = speedStop.UnixNano() - speedStart.UnixNano()
	seconds = float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	fmt.Printf("%.2fG/s based on total value length\n", float64(totalValueLength)/seconds/1024.0/1024.0/1024.0)
	fmt.Printf("%.2fG/s based on total bytes to disk\n", float64(bytesWritten2-bytesWritten)/seconds/1024.0/1024.0/1024.0)
	m = 0
	for i := 0; i < clients; i++ {
		m += <-c[i]
	}
	speedStop = time.Now()
	nanoseconds = speedStop.UnixNano() - speedStart.UnixNano()
	seconds = float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	runtime.ReadMemStats(&st)
	fmt.Printf("%.2fG total alloc\n", float64(st.TotalAlloc)/1024/1024/1024)
	fmt.Printf("%.2f key lookups per second\n", float64(totalKeys)/seconds)
	fmt.Printf("%.2fns per key lookup\n", float64(nanoseconds)/float64(totalKeys))
	fmt.Println(m, "keys missing")

	fmt.Println()
	start = time.Now()
	speedStart = start
	for i := 0; i < clients; i++ {
		go func(keys []byte, c chan int) {
			m := 0
			for o := 0; o < len(keys); o += 16 {
				v, _ := s.Get(binary.LittleEndian.Uint64(keys[o:]), binary.LittleEndian.Uint64(keys[o+8:]))
				if v == nil {
					m++
				}
			}
			c <- m
		}(keys[i], c[i])
	}
	m = 0
	for i := 0; i < clients; i++ {
		m += <-c[i]
	}
	for i := 0; i < clients; i++ {
		go func(keys []byte, c chan int) {
			m := 0
			for o := 0; o < len(keys); o += 16 {
				v, _ := s.Get(binary.LittleEndian.Uint64(keys[o:]), binary.LittleEndian.Uint64(keys[o+8:]))
				if v == nil {
					m++
				}
			}
			c <- m
		}(keys2[i], c[i])
	}
	for i := 0; i < clients; i++ {
		m += <-c[i]
	}
	speedStop = time.Now()
	fmt.Println(time.Now().Sub(start), "to lookup both sets of keys")
	nanoseconds = speedStop.UnixNano() - speedStart.UnixNano()
	seconds = float64(speedStop.UnixNano()-speedStart.UnixNano()) / 1000000000.0
	fmt.Printf("%.2f key lookups per second\n", float64(totalKeys*2)/seconds)
	fmt.Printf("%.2fns per key lookup\n", float64(nanoseconds)/float64(totalKeys*2))
	fmt.Println(m, "keys missing")
	/*
	   f, err := os.Create("brimstore.blocking.pprof")
	   if err != nil {
	       panic(err)
	   }
	   blockPprof.WriteTo(f, 0)
	   f.Close()
	*/
	/*
	   pprof.StopCPUProfile()
	   cpuProf.Close()
	*/
}
<|endoftext|>


<commit_before>package s2tools

import (
	"github.com/golang/geo/s2"
	"github.com/twpayne/go-geom"
	"github.com/twpayne/go-geom/encoding/geojson"
)

// CellUnionToGeoJSON helpers to display s2 cells on maps with GeoJSON
// exports cell union into its GeoJSON representation
func CellUnionToGeoJSON(cu s2.CellUnion) []byte {
	fc := geojson.FeatureCollection{}
	for _, cid := range cu {
		f := &geojson.Feature{}
		f.Properties = make(map[string]interface{})
		f.Properties["id"] = cid.ToToken()
		f.Properties["uid"] = uint64(cid)
		f.Properties["str"] = cid.String()
		f.Properties["level"] = cid.Level()

		c := s2.CellFromCellID(cid)
		coords := make([]float64, 5*2)
		for i := 0; i < 4; i++ {
			p := c.Vertex(i)
			ll := s2.LatLngFromPoint(p)
			coords[i*2] = ll.Lng.Degrees()
			coords[i*2+1] = ll.Lat.Degrees()
		}
		// last is first
		coords[8], coords[9] = coords[0], coords[1]
		ng := geom.NewPolygonFlat(geom.XY, coords, []int{10})
		f.Geometry = ng
		fc.Features = append(fc.Features, f)
	}
	b, _ := fc.MarshalJSON()
	return b
}

// CellUnionToTokens a cell union to a token string list
func CellUnionToTokens(cu s2.CellUnion) []string {
	res := make([]string, len(cu))

	for i, c := range cu {
		res[i] = c.ToToken()
	}
	return res
}
<commit_msg>convert uid to string<commit_after>package s2tools

import (
	"strconv"

	"github.com/golang/geo/s2"
	"github.com/twpayne/go-geom"
	"github.com/twpayne/go-geom/encoding/geojson"
)

// CellUnionToGeoJSON helpers to display s2 cells on maps with GeoJSON
// exports cell union into its GeoJSON representation
func CellUnionToGeoJSON(cu s2.CellUnion) []byte {
	fc := geojson.FeatureCollection{}
	for _, cid := range cu {
		f := &geojson.Feature{}
		f.Properties = make(map[string]interface{})
		f.Properties["id"] = cid.ToToken()
		f.Properties["uid"] = strconv.FormatUint(uint64(cid), 10)
		f.Properties["str"] = cid.String()
		f.Properties["level"] = cid.Level()

		c := s2.CellFromCellID(cid)
		coords := make([]float64, 5*2)
		for i := 0; i < 4; i++ {
			p := c.Vertex(i)
			ll := s2.LatLngFromPoint(p)
			coords[i*2] = ll.Lng.Degrees()
			coords[i*2+1] = ll.Lat.Degrees()
		}
		// last is first
		coords[8], coords[9] = coords[0], coords[1]
		ng := geom.NewPolygonFlat(geom.XY, coords, []int{10})
		f.Geometry = ng
		fc.Features = append(fc.Features, f)
	}
	b, _ := fc.MarshalJSON()
	return b
}

// CellUnionToTokens a cell union to a token string list
func CellUnionToTokens(cu s2.CellUnion) []string {
	res := make([]string, len(cu))

	for i, c := range cu {
		res[i] = c.ToToken()
	}
	return res
}
<|endoftext|>


<commit_before>//
// config.go -- samtun configuration
//
package samtun

import (
  "encoding/json"
  "io/ioutil"
)

type jsonConfig struct {
  Keyfile string
  Sam string
  Addr string
  Ifname string
  Session string
  Netmask string
  MTU int
  Map addrMap
}

func (conf *jsonConfig) Save(fname string) (err error) {
  var data []byte
  data, err = json.Marshal(conf)
  if err == nil {
    err = ioutil.WriteFile(fname, data, 0600)
  }
  return
}

// generate default config
func genConfig(fname string) (cfg jsonConfig) {
  cfg.Keyfile = "samtun.key"
  cfg.Sam = "127.0.0.1:7656"
  cfg.Ifname = "i2p0"
  cfg.MTU = 8192
  cfg.Addr = "10.9.0.1/24"
  cfg.Netmask = "255.255.0.0"
  cfg.Session = "samtun"
  cfg.Map = make(addrMap)
  return
}

// load samtun config
// does not check validity
func loadConfig(fname string) (conf jsonConfig, err error) {
  var data []byte
  data, err = ioutil.ReadFile(fname)
  if err == nil {
    err = json.Unmarshal(data, &conf)
  }
  return
}

<commit_msg>ammend json config<commit_after>//
// config.go -- samtun configuration
//
package samtun

import (
  "bytes"
  "encoding/json"
  "io/ioutil"
)

type jsonConfig struct {
  Keyfile string
  Sam string
  Addr string
  Ifname string
  Session string
  Netmask string
  MTU int
  Map addrMap
}

func (conf *jsonConfig) Save(fname string) (err error) {
  var data []byte
  data, err = json.Marshal(conf)
  if err == nil {
    var buff bytes.Buffer
    err = json.Indent(&buff, data, " ", "  ")
    if err == nil {
      err = ioutil.WriteFile(fname, buff.Bytes(), 0600)
    }
  }
  return
}

// generate default config
func genConfig(fname string) (cfg jsonConfig) {
  cfg.Keyfile = "samtun.key"
  cfg.Sam = "127.0.0.1:7656"
  cfg.Ifname = "i2p0"
  cfg.MTU = 8192
  cfg.Addr = "10.9.0.1/24"
  cfg.Netmask = "255.255.0.0"
  cfg.Session = "samtun"
  cfg.Map = make(addrMap)
  return
}

// load samtun config
// does not check validity
func loadConfig(fname string) (conf jsonConfig, err error) {
  var data []byte
  data, err = ioutil.ReadFile(fname)
  if err == nil {
    err = json.Unmarshal(data, &conf)
  }
  return
}

<|endoftext|>


<commit_before>package virt

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	log "github.com/sirupsen/logrus"

	virtapi "github.com/projecteru2/libyavirt/client"
	virttypes "github.com/projecteru2/libyavirt/types"

	"github.com/projecteru2/core/cluster"
	"github.com/projecteru2/core/engine"
	enginetypes "github.com/projecteru2/core/engine/types"
	coresource "github.com/projecteru2/core/source"
	coretypes "github.com/projecteru2/core/types"
)

const (
	// HTTPPrefixKey indicate http yavirtd
	HTTPPrefixKey = "virt://"
	// GRPCPrefixKey indicates grpc yavirtd
	GRPCPrefixKey = "virt-grpc://"
	// DmiUUIDKey indicates the key within deploy info.
	DmiUUIDKey = "DMIUUID"
)

// Virt implements the core engine.API interface.
type Virt struct {
	client virtapi.Client
	config coretypes.Config
}

// MakeClient makes a virt. client which wraps yavirt API client.
func MakeClient(ctx context.Context, config coretypes.Config, nodename, endpoint, ca, cert, key string) (engine.API, error) {
	var uri string
	if strings.HasPrefix(endpoint, HTTPPrefixKey) {
		uri = fmt.Sprintf("http://%s/%s", strings.TrimPrefix(endpoint, HTTPPrefixKey), config.Virt.APIVersion)
	} else if strings.HasPrefix(endpoint, GRPCPrefixKey) {
		uri = "grpc://" + strings.TrimPrefix(endpoint, GRPCPrefixKey)
	} else {
		return nil, fmt.Errorf("invalid endpoint: %s", endpoint)
	}

	cli, err := virtapi.New(uri)
	if err != nil {
		return nil, err
	}
	return &Virt{cli, config}, nil
}

// Info shows a connected node's information.
func (v *Virt) Info(ctx context.Context) (*enginetypes.Info, error) {
	resp, err := v.client.Info(ctx)
	if err != nil {
		return nil, err
	}

	return &enginetypes.Info{
		ID:           resp.ID,
		NCPU:         resp.Cpu,
		MemTotal:     resp.Mem,
		StorageTotal: resp.Storage,
	}, nil
}

// ExecCreate creates an execution.
func (v *Virt) ExecCreate(ctx context.Context, target string, config *enginetypes.ExecConfig) (id string, err error) {
	return "", fmt.Errorf("ExecCreate does not implement")
}

// ExecAttach executes an attachment.
func (v *Virt) ExecAttach(ctx context.Context, execID string, tty bool) (io.ReadCloser, io.WriteCloser, error) {
	return nil, nil, fmt.Errorf("ExecAttach does not implement")
}

// Execute executes a command in vm
func (v *Virt) Execute(ctx context.Context, target string, config *enginetypes.ExecConfig) (_ string, outputStream io.ReadCloser, inputStream io.WriteCloser, err error) {
	if config.Tty {
		flags := virttypes.AttachGuestFlags{Safe: true, Force: true}
		stream, err := v.client.AttachGuest(ctx, target, flags)
		if err != nil {
			return "", nil, nil, err
		}
		return "", ioutil.NopCloser(stream), stream, nil

	}

	msg, err := v.client.ExecuteGuest(ctx, target, config.Cmd)
	return "", ioutil.NopCloser(bytes.NewReader(msg.Data)), nil, err

}

// ExecExitCode gets return code of a specific execution.
func (v *Virt) ExecExitCode(ctx context.Context, execID string) (code int, err error) {
	return 0, nil
}

// ExecResize resize exec tty
func (v *Virt) ExecResize(ctx context.Context, execID string, height, width uint) (err error) {
	return nil
}

// NetworkConnect connects to a network.
func (v *Virt) NetworkConnect(ctx context.Context, network, target, ipv4, ipv6 string) (err error) {
	log.Warnf("NetworkConnect does not implement")
	return
}

// NetworkDisconnect disconnects from one network.
func (v *Virt) NetworkDisconnect(ctx context.Context, network, target string, force bool) (err error) {
	log.Warnf("NetworkDisconnect does not implement")
	return
}

// NetworkList lists all of networks.
func (v *Virt) NetworkList(ctx context.Context, drivers []string) (nets []*enginetypes.Network, err error) {
	log.Warnf("NetworkList does not implement")
	return
}

// BuildRefs builds references, it's not necessary for virt. presently.
func (v *Virt) BuildRefs(ctx context.Context, name string, tags []string) (refs []string) {
	log.Warnf("BuildRefs does not implement")
	return
}

// BuildContent builds content, the use of it is similar to BuildRefs.
func (v *Virt) BuildContent(ctx context.Context, scm coresource.Source, opts *enginetypes.BuildContentOptions) (string, io.Reader, error) {
	return "", nil, fmt.Errorf("BuildContent does not implement")
}

// VirtualizationCreate creates a guest.
func (v *Virt) VirtualizationCreate(ctx context.Context, opts *enginetypes.VirtualizationCreateOptions) (guest *enginetypes.VirtualizationCreated, err error) {
	vols, err := v.parseVolumes(opts.Volumes)
	if err != nil {
		return nil, err
	}

	stor := MinVirtStorage
	for _, cap := range vols {
		stor += cap
	}
	if opts.Storage < stor {
		return nil, coretypes.NewDetailedErr(coretypes.ErrInsufficientStorage,
			fmt.Sprintf("specify at least %d bytes for the storage", stor))
	}

	req := virttypes.CreateGuestReq{
		Cpu:       int(opts.Quota),
		Mem:       opts.Memory,
		ImageName: opts.Image,
		Volumes:   vols,
	}

	if dmiUUID, exists := opts.Labels[DmiUUIDKey]; exists {
		req.DmiUuid = dmiUUID
	}

	var resp virttypes.Guest
	if resp, err = v.client.CreateGuest(ctx, req); err != nil {
		return nil, err
	}

	return &enginetypes.VirtualizationCreated{ID: resp.ID, Name: opts.Name}, nil
}

// VirtualizationCopyTo copies one.
func (v *Virt) VirtualizationCopyTo(ctx context.Context, ID, target string, content io.Reader, AllowOverwriteDirWithFile, CopyUIDGID bool) (err error) {
	log.Warnf("VirtualizationCopyTo does not implement")
	return
}

// VirtualizationStart boots a guest.
func (v *Virt) VirtualizationStart(ctx context.Context, ID string) (err error) {
	_, err = v.client.StartGuest(ctx, ID)
	return
}

// VirtualizationStop stops it.
func (v *Virt) VirtualizationStop(ctx context.Context, ID string) (err error) {
	_, err = v.client.StopGuest(ctx, ID)
	return
}

// VirtualizationRemove removes a guest.
func (v *Virt) VirtualizationRemove(ctx context.Context, ID string, volumes, force bool) (err error) {
	_, err = v.client.DestroyGuest(ctx, ID, force)
	return
}

// VirtualizationInspect gets a guest.
func (v *Virt) VirtualizationInspect(ctx context.Context, ID string) (*enginetypes.VirtualizationInfo, error) {
	guest, err := v.client.GetGuest(ctx, ID)
	if err != nil {
		return nil, err
	}

	bytes, err := json.Marshal(coretypes.LabelMeta{Publish: []string{"PORT"}})
	if err != nil {
		return nil, err
	}

	return &enginetypes.VirtualizationInfo{
		ID:       guest.ID,
		Image:    guest.ImageName,
		Running:  guest.Status == "running",
		Networks: guest.Networks,
		Labels:   map[string]string{cluster.LabelMeta: string(bytes), cluster.ERUMark: "1"},
	}, nil
}

// VirtualizationLogs streams a specific guest's log.
func (v *Virt) VirtualizationLogs(ctx context.Context, opts *enginetypes.VirtualizationLogStreamOptions) (reader io.ReadCloser, err error) {
	return nil, fmt.Errorf("VirtualizationLogs does not implement")
}

// VirtualizationAttach attaches something to a guest.
func (v *Virt) VirtualizationAttach(ctx context.Context, ID string, stream, stdin bool) (io.ReadCloser, io.WriteCloser, error) {
	return nil, nil, fmt.Errorf("VirtualizationAttach does not implement")
}

// VirtualizationResize resized window size
func (v *Virt) VirtualizationResize(ctx context.Context, ID string, height, width uint) error {
	return fmt.Errorf("VirtualizationResize not implemented")
}

// VirtualizationWait is waiting for a shut-off
func (v *Virt) VirtualizationWait(ctx context.Context, ID, state string) (*enginetypes.VirtualizationWaitResult, error) {
	return nil, fmt.Errorf("VirtualizationWait does not implement")
}

// VirtualizationUpdateResource updates resource.
func (v *Virt) VirtualizationUpdateResource(ctx context.Context, ID string, opts *enginetypes.VirtualizationResource) error {
	vols, err := v.parseVolumes(opts.Volumes)
	if err != nil {
		return err
	}

	args := virttypes.ResizeGuestReq{
		Cpu:     int(opts.Quota),
		Mem:     opts.Memory,
		Volumes: vols,
	}
	args.ID = ID

	_, err = v.client.ResizeGuest(ctx, args)
	return err
}

// VirtualizationCopyFrom copies from another.
func (v *Virt) VirtualizationCopyFrom(ctx context.Context, ID, path string) (io.ReadCloser, string, error) {
	return nil, "", fmt.Errorf("VirtualizationCopyFrom does not implement")
}

// VirtualizationExecute executes commands in running virtual unit
func (v *Virt) VirtualizationExecute(ctx context.Context, ID string, commands, env []string, workdir string) (io.WriteCloser, io.ReadCloser, error) {
	return nil, nil, fmt.Errorf("VirtualizationExecute not implemented")
}

// ResourceValidate validate resource usage
func (v *Virt) ResourceValidate(ctx context.Context, cpu float64, cpumap map[string]int64, memory, storage int64) error {
	// TODO list all containers, calcuate resource
	return nil
}
<commit_msg>implement virt ExecResize (#217)<commit_after>package virt

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	log "github.com/sirupsen/logrus"

	virtapi "github.com/projecteru2/libyavirt/client"
	virttypes "github.com/projecteru2/libyavirt/types"

	"github.com/projecteru2/core/cluster"
	"github.com/projecteru2/core/engine"
	enginetypes "github.com/projecteru2/core/engine/types"
	coresource "github.com/projecteru2/core/source"
	coretypes "github.com/projecteru2/core/types"
)

const (
	// HTTPPrefixKey indicate http yavirtd
	HTTPPrefixKey = "virt://"
	// GRPCPrefixKey indicates grpc yavirtd
	GRPCPrefixKey = "virt-grpc://"
	// DmiUUIDKey indicates the key within deploy info.
	DmiUUIDKey = "DMIUUID"
)

// Virt implements the core engine.API interface.
type Virt struct {
	client virtapi.Client
	config coretypes.Config
}

// MakeClient makes a virt. client which wraps yavirt API client.
func MakeClient(ctx context.Context, config coretypes.Config, nodename, endpoint, ca, cert, key string) (engine.API, error) {
	var uri string
	if strings.HasPrefix(endpoint, HTTPPrefixKey) {
		uri = fmt.Sprintf("http://%s/%s", strings.TrimPrefix(endpoint, HTTPPrefixKey), config.Virt.APIVersion)
	} else if strings.HasPrefix(endpoint, GRPCPrefixKey) {
		uri = "grpc://" + strings.TrimPrefix(endpoint, GRPCPrefixKey)
	} else {
		return nil, fmt.Errorf("invalid endpoint: %s", endpoint)
	}

	cli, err := virtapi.New(uri)
	if err != nil {
		return nil, err
	}
	return &Virt{cli, config}, nil
}

// Info shows a connected node's information.
func (v *Virt) Info(ctx context.Context) (*enginetypes.Info, error) {
	resp, err := v.client.Info(ctx)
	if err != nil {
		return nil, err
	}

	return &enginetypes.Info{
		ID:           resp.ID,
		NCPU:         resp.Cpu,
		MemTotal:     resp.Mem,
		StorageTotal: resp.Storage,
	}, nil
}

// ExecCreate creates an execution.
func (v *Virt) ExecCreate(ctx context.Context, target string, config *enginetypes.ExecConfig) (id string, err error) {
	return "", fmt.Errorf("ExecCreate does not implement")
}

// ExecAttach executes an attachment.
func (v *Virt) ExecAttach(ctx context.Context, execID string, tty bool) (io.ReadCloser, io.WriteCloser, error) {
	return nil, nil, fmt.Errorf("ExecAttach does not implement")
}

// Execute executes a command in vm
func (v *Virt) Execute(ctx context.Context, target string, config *enginetypes.ExecConfig) (execID string, outputStream io.ReadCloser, inputStream io.WriteCloser, err error) {
	if config.Tty {
		flags := virttypes.AttachGuestFlags{Safe: true, Force: true}
		stream, err := v.client.AttachGuest(ctx, target, flags)
		if err != nil {
			return "", nil, nil, err
		}
		return target, ioutil.NopCloser(stream), stream, nil

	}

	msg, err := v.client.ExecuteGuest(ctx, target, config.Cmd)
	return target, ioutil.NopCloser(bytes.NewReader(msg.Data)), nil, err

}

// ExecExitCode gets return code of a specific execution.
func (v *Virt) ExecExitCode(ctx context.Context, execID string) (code int, err error) {
	return 0, nil
}

// ExecResize resize exec tty
func (v *Virt) ExecResize(ctx context.Context, execID string, height, width uint) (err error) {
	resizeCmd := fmt.Sprintf("/bin/stty -F /dev/ttyS0 rows %d cols %d", height, width)
	msg, err := v.client.ExecuteGuest(ctx, execID, strings.Split(resizeCmd, " "))
	log.Debugf("[ExecResize] resize got response: %v", msg)
	return err
}

// NetworkConnect connects to a network.
func (v *Virt) NetworkConnect(ctx context.Context, network, target, ipv4, ipv6 string) (err error) {
	log.Warnf("NetworkConnect does not implement")
	return
}

// NetworkDisconnect disconnects from one network.
func (v *Virt) NetworkDisconnect(ctx context.Context, network, target string, force bool) (err error) {
	log.Warnf("NetworkDisconnect does not implement")
	return
}

// NetworkList lists all of networks.
func (v *Virt) NetworkList(ctx context.Context, drivers []string) (nets []*enginetypes.Network, err error) {
	log.Warnf("NetworkList does not implement")
	return
}

// BuildRefs builds references, it's not necessary for virt. presently.
func (v *Virt) BuildRefs(ctx context.Context, name string, tags []string) (refs []string) {
	log.Warnf("BuildRefs does not implement")
	return
}

// BuildContent builds content, the use of it is similar to BuildRefs.
func (v *Virt) BuildContent(ctx context.Context, scm coresource.Source, opts *enginetypes.BuildContentOptions) (string, io.Reader, error) {
	return "", nil, fmt.Errorf("BuildContent does not implement")
}

// VirtualizationCreate creates a guest.
func (v *Virt) VirtualizationCreate(ctx context.Context, opts *enginetypes.VirtualizationCreateOptions) (guest *enginetypes.VirtualizationCreated, err error) {
	vols, err := v.parseVolumes(opts.Volumes)
	if err != nil {
		return nil, err
	}

	stor := MinVirtStorage
	for _, cap := range vols {
		stor += cap
	}
	if opts.Storage < stor {
		return nil, coretypes.NewDetailedErr(coretypes.ErrInsufficientStorage,
			fmt.Sprintf("specify at least %d bytes for the storage", stor))
	}

	req := virttypes.CreateGuestReq{
		Cpu:       int(opts.Quota),
		Mem:       opts.Memory,
		ImageName: opts.Image,
		Volumes:   vols,
	}

	if dmiUUID, exists := opts.Labels[DmiUUIDKey]; exists {
		req.DmiUuid = dmiUUID
	}

	var resp virttypes.Guest
	if resp, err = v.client.CreateGuest(ctx, req); err != nil {
		return nil, err
	}

	return &enginetypes.VirtualizationCreated{ID: resp.ID, Name: opts.Name}, nil
}

// VirtualizationCopyTo copies one.
func (v *Virt) VirtualizationCopyTo(ctx context.Context, ID, target string, content io.Reader, AllowOverwriteDirWithFile, CopyUIDGID bool) (err error) {
	log.Warnf("VirtualizationCopyTo does not implement")
	return
}

// VirtualizationStart boots a guest.
func (v *Virt) VirtualizationStart(ctx context.Context, ID string) (err error) {
	_, err = v.client.StartGuest(ctx, ID)
	return
}

// VirtualizationStop stops it.
func (v *Virt) VirtualizationStop(ctx context.Context, ID string) (err error) {
	_, err = v.client.StopGuest(ctx, ID)
	return
}

// VirtualizationRemove removes a guest.
func (v *Virt) VirtualizationRemove(ctx context.Context, ID string, volumes, force bool) (err error) {
	_, err = v.client.DestroyGuest(ctx, ID, force)
	return
}

// VirtualizationInspect gets a guest.
func (v *Virt) VirtualizationInspect(ctx context.Context, ID string) (*enginetypes.VirtualizationInfo, error) {
	guest, err := v.client.GetGuest(ctx, ID)
	if err != nil {
		return nil, err
	}

	bytes, err := json.Marshal(coretypes.LabelMeta{Publish: []string{"PORT"}})
	if err != nil {
		return nil, err
	}

	return &enginetypes.VirtualizationInfo{
		ID:       guest.ID,
		Image:    guest.ImageName,
		Running:  guest.Status == "running",
		Networks: guest.Networks,
		Labels:   map[string]string{cluster.LabelMeta: string(bytes), cluster.ERUMark: "1"},
	}, nil
}

// VirtualizationLogs streams a specific guest's log.
func (v *Virt) VirtualizationLogs(ctx context.Context, opts *enginetypes.VirtualizationLogStreamOptions) (reader io.ReadCloser, err error) {
	return nil, fmt.Errorf("VirtualizationLogs does not implement")
}

// VirtualizationAttach attaches something to a guest.
func (v *Virt) VirtualizationAttach(ctx context.Context, ID string, stream, stdin bool) (io.ReadCloser, io.WriteCloser, error) {
	return nil, nil, fmt.Errorf("VirtualizationAttach does not implement")
}

// VirtualizationResize resized window size
func (v *Virt) VirtualizationResize(ctx context.Context, ID string, height, width uint) error {
	return fmt.Errorf("VirtualizationResize not implemented")
}

// VirtualizationWait is waiting for a shut-off
func (v *Virt) VirtualizationWait(ctx context.Context, ID, state string) (*enginetypes.VirtualizationWaitResult, error) {
	return nil, fmt.Errorf("VirtualizationWait does not implement")
}

// VirtualizationUpdateResource updates resource.
func (v *Virt) VirtualizationUpdateResource(ctx context.Context, ID string, opts *enginetypes.VirtualizationResource) error {
	vols, err := v.parseVolumes(opts.Volumes)
	if err != nil {
		return err
	}

	args := virttypes.ResizeGuestReq{
		Cpu:     int(opts.Quota),
		Mem:     opts.Memory,
		Volumes: vols,
	}
	args.ID = ID

	_, err = v.client.ResizeGuest(ctx, args)
	return err
}

// VirtualizationCopyFrom copies from another.
func (v *Virt) VirtualizationCopyFrom(ctx context.Context, ID, path string) (io.ReadCloser, string, error) {
	return nil, "", fmt.Errorf("VirtualizationCopyFrom does not implement")
}

// VirtualizationExecute executes commands in running virtual unit
func (v *Virt) VirtualizationExecute(ctx context.Context, ID string, commands, env []string, workdir string) (io.WriteCloser, io.ReadCloser, error) {
	return nil, nil, fmt.Errorf("VirtualizationExecute not implemented")
}

// ResourceValidate validate resource usage
func (v *Virt) ResourceValidate(ctx context.Context, cpu float64, cpumap map[string]int64, memory, storage int64) error {
	// TODO list all containers, calcuate resource
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"syscall"

	"github.com/dickeyxxx/golock"
	"github.com/heroku/heroku-cli/gode"
)

// Plugin represents a javascript plugin
type Plugin struct {
	Name     string     `json:"name"`
	Version  string     `json:"version"`
	Topics   TopicSet   `json:"topics"`
	Topic    *Topic     `json:"topic"`
	Commands CommandSet `json:"commands"`
}

// SetupNode sets up node and npm in ~/.heroku
func SetupNode() {
	gode.SetRootPath(AppDir())
	setup, err := gode.IsSetup()
	WarnIfError(err)
	if !setup {
		if err := gode.Setup(); err != nil {
			panic(err)
		}
	}
}

// LoadPlugins loads the topics and commands from the JavaScript plugins into the CLI
func (cli *Cli) LoadPlugins(plugins map[string]*Plugin) {
	for _, plugin := range plugins {
		for _, topic := range plugin.Topics {
			cli.AddTopic(topic)
		}
		if plugin.Topic != nil {
			cli.AddTopic(plugin.Topic)
		}
		for _, command := range plugin.Commands {
			if !cli.AddCommand(command) {
				Errf("WARNING: command %s has already been defined\n", command)
			}
		}
	}
	sort.Sort(cli.Topics)
	sort.Sort(cli.Commands)
}

var pluginsTopic = &Topic{
	Name:        "plugins",
	Description: "manage plugins",
}

var pluginsInstallCmd = &Command{
	Topic:       "plugins",
	Command:     "install",
	Hidden:      true,
	Args:        []Arg{{Name: "name"}},
	Description: "Installs a plugin into the CLI",
	Help: `Install a Heroku plugin

  Example:
  $ heroku plugins:install dickeyxxx/heroku-production-status`,

	Run: func(ctx *Context) {
		name := ctx.Args.(map[string]string)["name"]
		if len(name) == 0 {
			Errln("Must specify a plugin name")
			return
		}
		action("Installing plugin "+name, "done", func() {
			ExitIfError(installPlugins(name))
		})
	},
}

var pluginsLinkCmd = &Command{
	Topic:       "plugins",
	Command:     "link",
	Description: "Links a local plugin into CLI",
	Args:        []Arg{{Name: "path", Optional: true}},
	Help: `Links a local plugin into CLI.
	This is useful when developing plugins locally.
	It simply symlinks the specified path into ~/.heroku/node_modules

  Example:
	$ heroku plugins:link .`,

	Run: func(ctx *Context) {
		path := ctx.Args.(map[string]string)["path"]
		if path == "" {
			path = "."
		}
		path, err := filepath.Abs(path)
		if err != nil {
			panic(err)
		}
		if _, err := os.Stat(path); err != nil {
			panic(err)
		}
		name := filepath.Base(path)
		newPath := pluginPath(name)
		os.Remove(newPath)
		os.RemoveAll(newPath)
		err = os.Symlink(path, newPath)
		if err != nil {
			panic(err)
		}
		plugin, err := ParsePlugin(name)
		ExitIfError(err)
		if name != plugin.Name {
			path = newPath
			newPath = pluginPath(plugin.Name)
			os.Remove(newPath)
			os.RemoveAll(newPath)
			os.Rename(path, newPath)
		}
		Println("Symlinked", plugin.Name)
		AddPluginsToCache(plugin)
	},
}

var pluginsUninstallCmd = &Command{
	Topic:       "plugins",
	Command:     "uninstall",
	Hidden:      true,
	Args:        []Arg{{Name: "name"}},
	Description: "Uninstalls a plugin from the CLI",
	Help: `Uninstalls a Heroku plugin

  Example:
  $ heroku plugins:uninstall heroku-production-status`,

	Run: func(ctx *Context) {
		name := ctx.Args.(map[string]string)["name"]
		if !contains(PluginNames(), name) {
			ExitIfError(errors.New(name + " is not installed"))
		}
		Errf("Uninstalling plugin %s...", name)
		ExitIfError(gode.RemovePackages(name))
		RemovePluginFromCache(name)
		Errln(" done")
	},
}

var pluginsListCmd = &Command{
	Topic:            "plugins",
	Hidden:           true,
	Description:      "Lists installed plugins",
	DisableAnalytics: true,
	Help: `
Example:
  $ heroku plugins`,

	Run: func(ctx *Context) {
		SetupBuiltinPlugins()
		var plugins []string
		for _, plugin := range GetPlugins() {
			symlinked := ""
			if isPluginSymlinked(plugin.Name) {
				symlinked = " (symlinked)"
			}
			plugins = append(plugins, fmt.Sprintf("%s %s %s", plugin.Name, plugin.Version, symlinked))
		}
		sort.Strings(plugins)
		for _, plugin := range plugins {
			Println(plugin)
		}
	},
}

func runFn(plugin *Plugin, topic, command string) func(ctx *Context) {
	return func(ctx *Context) {
		readLockPlugin(plugin.Name)
		ctx.Dev = isPluginSymlinked(plugin.Name)
		ctxJSON, err := json.Marshal(ctx)
		if err != nil {
			panic(err)
		}
		title, _ := json.Marshal(processTitle(ctx))
		script := fmt.Sprintf(`'use strict';
var moduleName = '%s';
var moduleVersion = '%s';
var topic = '%s';
var command = '%s';
process.title = %s;
var ctx = %s;
ctx.version = ctx.version + ' ' + moduleName + '/' + moduleVersion + ' node-' + process.version;
var logPath = %s;
process.chdir(ctx.cwd);
if (!ctx.dev) {
	process.on('uncaughtException', function (err) {
		// ignore EPIPE errors (usually from piping to head)
		if (err.code === "EPIPE") return;
		console.error(' !   Error in ' + moduleName + ':')
		console.error(' !   ' + err.message || err);
		if (err.stack) {
			var fs = require('fs');
			var log = function (line) {
				var d = new Date().toISOString()
				.replace(/T/, ' ')
				.replace(/-/g, '/')
				.replace(/\..+/, '');
				fs.appendFileSync(logPath, d + ' ' + line + '\n');
			}
			log('Error during ' + topic + ':' + command);
			log(err.stack);
			console.error(' !   See ' + logPath + ' for more info.');
		}
		process.exit(1);
	});
}
if (command === '') { command = null }
var module = require(moduleName);
var cmd = module.commands.filter(function (c) {
	return c.topic === topic && c.command == command;
})[0];
cmd.run(ctx);
`, plugin.Name, plugin.Version, topic, command, string(title), ctxJSON, strconv.Quote(ErrLogPath))

		// swallow sigint since the plugin will handle it
		swallowSigint = true

		currentAnalyticsCommand.Plugin = plugin.Name
		currentAnalyticsCommand.Version = plugin.Version
		currentAnalyticsCommand.Language = fmt.Sprintf("node/%s", gode.NodeVersion)

		cmd, done := gode.RunScript(script)
		defer done()
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if ctx.Flags["debugger"] == true {
			cmd = gode.DebugScript(script)
		}
		err = cmd.Run()
		Exit(getExitCode(err))
	}
}

func getExitCode(err error) int {
	switch e := err.(type) {
	case nil:
		return 0
	case *exec.ExitError:
		status, ok := e.Sys().(syscall.WaitStatus)
		if !ok {
			panic(err)
		}
		return status.ExitStatus()
	default:
		panic(err)
	}
}

var pluginInstallRetrying = false

// ParsePlugin requires the plugin's node module
// to get the commands and metadata
func ParsePlugin(name string) (*Plugin, error) {
	script := `
	var plugin = require('` + name + `');
	if (!plugin.commands) throw new Error('Contains no commands. Is this a real plugin?');
	var pjson  = require('` + name + `/package.json');

	plugin.name    = pjson.name;
	plugin.version = pjson.version;

	console.log(JSON.stringify(plugin))`
	cmd, done := gode.RunScript(script)
	cmd.Stderr = Stderr
	output, err := cmd.Output()
	done()
	if err != nil {
		// try again but this time grab stdout and stderr
		cmd, done := gode.RunScript(script)
		output, _ := cmd.CombinedOutput()
		done()
		if !pluginInstallRetrying && strings.Contains(string(output), "Error: Cannot find module") {
			pluginInstallRetrying = true
			Warn("Failed to install " + name + ". Retrying...")
			WarnIfError(gode.RemovePackages(name))
			WarnIfError(gode.ClearCache())
			WarnIfError(gode.InstallPackages(name))
			return ParsePlugin(name)
		}
		return nil, fmt.Errorf("Error reading plugin: %s\n%s\n%s", name, err, output)
	}
	var plugin Plugin
	err = json.Unmarshal([]byte(output), &plugin)
	if err != nil {
		return nil, fmt.Errorf("Error parsing plugin: %s\n%s\n%s", name, err, string(output))
	}
	for _, command := range plugin.Commands {
		if command == nil {
			continue
		}
		command.Plugin = plugin.Name
		command.Help = strings.TrimSpace(command.Help)
	}
	return &plugin, nil
}

// GetPlugins goes through all the node plugins and returns them in Go stucts
func GetPlugins() map[string]*Plugin {
	plugins := FetchPluginCache()
	for name, plugin := range plugins {
		if plugin == nil || !pluginExists(name) || plugin.Commands.Len() == 0 {
			delete(plugins, name)
		} else {
			for _, command := range plugin.Commands {
				command.Run = runFn(plugin, command.Topic, command.Command)
			}
		}
	}
	return plugins
}

// PluginNames lists all the plugin names
func PluginNames() []string {
	plugins := FetchPluginCache()
	names := make([]string, 0, len(plugins))
	for _, plugin := range plugins {
		if plugin != nil && pluginExists(plugin.Name) && len(plugin.Commands) > 0 {
			names = append(names, plugin.Name)
		}
	}
	return names
}

// PluginNamesNotSymlinked returns all the plugins that are not symlinked
func PluginNamesNotSymlinked() []string {
	a := PluginNames()
	b := make([]string, 0, len(a))
	for _, plugin := range a {
		if !isPluginSymlinked(plugin) {
			b = append(b, plugin)
		}
	}
	return b
}

func isPluginSymlinked(plugin string) bool {
	path := filepath.Join(AppDir(), "node_modules", plugin)
	fi, err := os.Lstat(path)
	if err != nil {
		return false
	}
	return fi.Mode()&os.ModeSymlink != 0
}

// SetupBuiltinPlugins ensures all the builtinPlugins are installed
func SetupBuiltinPlugins() {
	pluginNames := difference(BuiltinPlugins, PluginNames())
	if len(pluginNames) == 0 {
		return
	}
	action("heroku-cli: Installing core plugins", "done", func() {
		if err := installPlugins(pluginNames...); err != nil {
			// retry once
			WarnIfError(gode.RemovePackages(pluginNames...))
			WarnIfError(gode.ClearCache())
			Err("\rheroku-cli: Installing core plugins (retrying)...")
			ExitIfError(installPlugins(pluginNames...))
		}
	})
}

func difference(a, b []string) []string {
	res := make([]string, 0, len(a))
	for _, aa := range a {
		if !contains(b, aa) {
			res = append(res, aa)
		}
	}
	return res
}

func contains(arr []string, s string) bool {
	for _, a := range arr {
		if a == s {
			return true
		}
	}
	return false
}

func installPlugins(names ...string) error {
	for _, name := range names {
		lockPlugin(name)
	}
	defer func() {
		for _, name := range names {
			unlockPlugin(name)
		}
	}()
	err := gode.InstallPackages(names...)
	if err != nil {
		return err
	}
	plugins := make([]*Plugin, len(names))
	for i, name := range names {
		plugin, err := ParsePlugin(name)
		if err != nil {
			return err
		}
		plugins[i] = plugin
	}
	AddPluginsToCache(plugins...)
	return nil
}

func pluginExists(plugin string) bool {
	exists, _ := fileExists(pluginPath(plugin))
	return exists
}

// directory location of plugin
func pluginPath(plugin string) string {
	return filepath.Join(AppDir(), "node_modules", plugin)
}

// lock a plugin for reading
func readLockPlugin(name string) {
	lockfile := updateLockPath + "." + name
	locked, err := golock.IsLocked(lockfile)
	LogIfError(err)
	if locked {
		lockPlugin(name)
		unlockPlugin(name)
	}
}

// lock a plugin for writing
func lockPlugin(name string) {
	LogIfError(golock.Lock(updateLockPath + "." + name))
}

// unlock a plugin
func unlockPlugin(name string) {
	LogIfError(golock.Unlock(updateLockPath + "." + name))
}
<commit_msg>better plugin retry logic<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"syscall"

	"github.com/dickeyxxx/golock"
	"github.com/heroku/heroku-cli/gode"
)

// Plugin represents a javascript plugin
type Plugin struct {
	Name     string     `json:"name"`
	Version  string     `json:"version"`
	Topics   TopicSet   `json:"topics"`
	Topic    *Topic     `json:"topic"`
	Commands CommandSet `json:"commands"`
}

// SetupNode sets up node and npm in ~/.heroku
func SetupNode() {
	gode.SetRootPath(AppDir())
	setup, err := gode.IsSetup()
	WarnIfError(err)
	if !setup {
		if err := gode.Setup(); err != nil {
			panic(err)
		}
	}
}

// LoadPlugins loads the topics and commands from the JavaScript plugins into the CLI
func (cli *Cli) LoadPlugins(plugins map[string]*Plugin) {
	for _, plugin := range plugins {
		for _, topic := range plugin.Topics {
			cli.AddTopic(topic)
		}
		if plugin.Topic != nil {
			cli.AddTopic(plugin.Topic)
		}
		for _, command := range plugin.Commands {
			if !cli.AddCommand(command) {
				Errf("WARNING: command %s has already been defined\n", command)
			}
		}
	}
	sort.Sort(cli.Topics)
	sort.Sort(cli.Commands)
}

var pluginsTopic = &Topic{
	Name:        "plugins",
	Description: "manage plugins",
}

var pluginsInstallCmd = &Command{
	Topic:       "plugins",
	Command:     "install",
	Hidden:      true,
	Args:        []Arg{{Name: "name"}},
	Description: "Installs a plugin into the CLI",
	Help: `Install a Heroku plugin

  Example:
  $ heroku plugins:install dickeyxxx/heroku-production-status`,

	Run: func(ctx *Context) {
		name := ctx.Args.(map[string]string)["name"]
		if len(name) == 0 {
			Errln("Must specify a plugin name")
			return
		}
		action("Installing plugin "+name, "done", func() {
			ExitIfError(installPlugins(name))
		})
	},
}

var pluginsLinkCmd = &Command{
	Topic:       "plugins",
	Command:     "link",
	Description: "Links a local plugin into CLI",
	Args:        []Arg{{Name: "path", Optional: true}},
	Help: `Links a local plugin into CLI.
	This is useful when developing plugins locally.
	It simply symlinks the specified path into ~/.heroku/node_modules

  Example:
	$ heroku plugins:link .`,

	Run: func(ctx *Context) {
		path := ctx.Args.(map[string]string)["path"]
		if path == "" {
			path = "."
		}
		path, err := filepath.Abs(path)
		if err != nil {
			panic(err)
		}
		if _, err := os.Stat(path); err != nil {
			panic(err)
		}
		name := filepath.Base(path)
		newPath := pluginPath(name)
		os.Remove(newPath)
		os.RemoveAll(newPath)
		err = os.Symlink(path, newPath)
		if err != nil {
			panic(err)
		}
		plugin, err := ParsePlugin(name)
		ExitIfError(err)
		if name != plugin.Name {
			path = newPath
			newPath = pluginPath(plugin.Name)
			os.Remove(newPath)
			os.RemoveAll(newPath)
			os.Rename(path, newPath)
		}
		Println("Symlinked", plugin.Name)
		AddPluginsToCache(plugin)
	},
}

var pluginsUninstallCmd = &Command{
	Topic:       "plugins",
	Command:     "uninstall",
	Hidden:      true,
	Args:        []Arg{{Name: "name"}},
	Description: "Uninstalls a plugin from the CLI",
	Help: `Uninstalls a Heroku plugin

  Example:
  $ heroku plugins:uninstall heroku-production-status`,

	Run: func(ctx *Context) {
		name := ctx.Args.(map[string]string)["name"]
		if !contains(PluginNames(), name) {
			ExitIfError(errors.New(name + " is not installed"))
		}
		Errf("Uninstalling plugin %s...", name)
		ExitIfError(gode.RemovePackages(name))
		RemovePluginFromCache(name)
		Errln(" done")
	},
}

var pluginsListCmd = &Command{
	Topic:            "plugins",
	Hidden:           true,
	Description:      "Lists installed plugins",
	DisableAnalytics: true,
	Help: `
Example:
  $ heroku plugins`,

	Run: func(ctx *Context) {
		SetupBuiltinPlugins()
		var plugins []string
		for _, plugin := range GetPlugins() {
			symlinked := ""
			if isPluginSymlinked(plugin.Name) {
				symlinked = " (symlinked)"
			}
			plugins = append(plugins, fmt.Sprintf("%s %s %s", plugin.Name, plugin.Version, symlinked))
		}
		sort.Strings(plugins)
		for _, plugin := range plugins {
			Println(plugin)
		}
	},
}

func runFn(plugin *Plugin, topic, command string) func(ctx *Context) {
	return func(ctx *Context) {
		readLockPlugin(plugin.Name)
		ctx.Dev = isPluginSymlinked(plugin.Name)
		ctxJSON, err := json.Marshal(ctx)
		if err != nil {
			panic(err)
		}
		title, _ := json.Marshal(processTitle(ctx))
		script := fmt.Sprintf(`'use strict';
var moduleName = '%s';
var moduleVersion = '%s';
var topic = '%s';
var command = '%s';
process.title = %s;
var ctx = %s;
ctx.version = ctx.version + ' ' + moduleName + '/' + moduleVersion + ' node-' + process.version;
var logPath = %s;
process.chdir(ctx.cwd);
if (!ctx.dev) {
	process.on('uncaughtException', function (err) {
		// ignore EPIPE errors (usually from piping to head)
		if (err.code === "EPIPE") return;
		console.error(' !   Error in ' + moduleName + ':')
		console.error(' !   ' + err.message || err);
		if (err.stack) {
			var fs = require('fs');
			var log = function (line) {
				var d = new Date().toISOString()
				.replace(/T/, ' ')
				.replace(/-/g, '/')
				.replace(/\..+/, '');
				fs.appendFileSync(logPath, d + ' ' + line + '\n');
			}
			log('Error during ' + topic + ':' + command);
			log(err.stack);
			console.error(' !   See ' + logPath + ' for more info.');
		}
		process.exit(1);
	});
}
if (command === '') { command = null }
var module = require(moduleName);
var cmd = module.commands.filter(function (c) {
	return c.topic === topic && c.command == command;
})[0];
cmd.run(ctx);
`, plugin.Name, plugin.Version, topic, command, string(title), ctxJSON, strconv.Quote(ErrLogPath))

		// swallow sigint since the plugin will handle it
		swallowSigint = true

		currentAnalyticsCommand.Plugin = plugin.Name
		currentAnalyticsCommand.Version = plugin.Version
		currentAnalyticsCommand.Language = fmt.Sprintf("node/%s", gode.NodeVersion)

		cmd, done := gode.RunScript(script)
		defer done()
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if ctx.Flags["debugger"] == true {
			cmd = gode.DebugScript(script)
		}
		err = cmd.Run()
		Exit(getExitCode(err))
	}
}

func getExitCode(err error) int {
	switch e := err.(type) {
	case nil:
		return 0
	case *exec.ExitError:
		status, ok := e.Sys().(syscall.WaitStatus)
		if !ok {
			panic(err)
		}
		return status.ExitStatus()
	default:
		panic(err)
	}
}

var pluginInstallRetrying = false

// ParsePlugin requires the plugin's node module
// to get the commands and metadata
func ParsePlugin(name string) (*Plugin, error) {
	script := `
	var plugin = require('` + name + `');
	if (!plugin.commands) throw new Error('Contains no commands. Is this a real plugin?');
	var pjson  = require('` + name + `/package.json');

	plugin.name    = pjson.name;
	plugin.version = pjson.version;

	console.log(JSON.stringify(plugin))`
	cmd, done := gode.RunScript(script)
	cmd.Stderr = Stderr
	output, err := cmd.Output()
	done()

	if err != nil {
		// try again but this time grab stdout and stderr
		cmd, done := gode.RunScript(script)
		output, err = cmd.CombinedOutput() // sometimes this actually works the second time
		if err != nil {
			done()
			if !pluginInstallRetrying && strings.Contains(string(output), "Error: Cannot find module") {
				pluginInstallRetrying = true
				Warn("Failed to install " + name + ". Retrying...")
				WarnIfError(gode.RemovePackages(name))
				WarnIfError(gode.ClearCache())
				WarnIfError(gode.InstallPackages(name))
				return ParsePlugin(name)
			}
			return nil, fmt.Errorf("Error reading plugin: %s\n%s\n%s", name, err, output)
		}
	}
	var plugin Plugin
	err = json.Unmarshal([]byte(output), &plugin)
	if err != nil {
		return nil, fmt.Errorf("Error parsing plugin: %s\n%s\n%s", name, err, string(output))
	}
	for _, command := range plugin.Commands {
		if command == nil {
			continue
		}
		command.Plugin = plugin.Name
		command.Help = strings.TrimSpace(command.Help)
	}
	return &plugin, nil
}

// GetPlugins goes through all the node plugins and returns them in Go stucts
func GetPlugins() map[string]*Plugin {
	plugins := FetchPluginCache()
	for name, plugin := range plugins {
		if plugin == nil || !pluginExists(name) || plugin.Commands.Len() == 0 {
			delete(plugins, name)
		} else {
			for _, command := range plugin.Commands {
				command.Run = runFn(plugin, command.Topic, command.Command)
			}
		}
	}
	return plugins
}

// PluginNames lists all the plugin names
func PluginNames() []string {
	plugins := FetchPluginCache()
	names := make([]string, 0, len(plugins))
	for _, plugin := range plugins {
		if plugin != nil && pluginExists(plugin.Name) && len(plugin.Commands) > 0 {
			names = append(names, plugin.Name)
		}
	}
	return names
}

// PluginNamesNotSymlinked returns all the plugins that are not symlinked
func PluginNamesNotSymlinked() []string {
	a := PluginNames()
	b := make([]string, 0, len(a))
	for _, plugin := range a {
		if !isPluginSymlinked(plugin) {
			b = append(b, plugin)
		}
	}
	return b
}

func isPluginSymlinked(plugin string) bool {
	path := filepath.Join(AppDir(), "node_modules", plugin)
	fi, err := os.Lstat(path)
	if err != nil {
		return false
	}
	return fi.Mode()&os.ModeSymlink != 0
}

// SetupBuiltinPlugins ensures all the builtinPlugins are installed
func SetupBuiltinPlugins() {
	pluginNames := difference(BuiltinPlugins, PluginNames())
	if len(pluginNames) == 0 {
		return
	}
	action("heroku-cli: Installing core plugins", "done", func() {
		if err := installPlugins(pluginNames...); err != nil {
			// retry once
			WarnIfError(gode.RemovePackages(pluginNames...))
			WarnIfError(gode.ClearCache())
			Err("\rheroku-cli: Installing core plugins (retrying)...")
			ExitIfError(installPlugins(pluginNames...))
		}
	})
}

func difference(a, b []string) []string {
	res := make([]string, 0, len(a))
	for _, aa := range a {
		if !contains(b, aa) {
			res = append(res, aa)
		}
	}
	return res
}

func contains(arr []string, s string) bool {
	for _, a := range arr {
		if a == s {
			return true
		}
	}
	return false
}

func installPlugins(names ...string) error {
	for _, name := range names {
		lockPlugin(name)
	}
	defer func() {
		for _, name := range names {
			unlockPlugin(name)
		}
	}()
	err := gode.InstallPackages(names...)
	if err != nil {
		return err
	}
	plugins := make([]*Plugin, len(names))
	for i, name := range names {
		plugin, err := ParsePlugin(name)
		if err != nil {
			return err
		}
		plugins[i] = plugin
	}
	AddPluginsToCache(plugins...)
	return nil
}

func pluginExists(plugin string) bool {
	exists, _ := fileExists(pluginPath(plugin))
	return exists
}

// directory location of plugin
func pluginPath(plugin string) string {
	return filepath.Join(AppDir(), "node_modules", plugin)
}

// lock a plugin for reading
func readLockPlugin(name string) {
	lockfile := updateLockPath + "." + name
	locked, err := golock.IsLocked(lockfile)
	LogIfError(err)
	if locked {
		lockPlugin(name)
		unlockPlugin(name)
	}
}

// lock a plugin for writing
func lockPlugin(name string) {
	LogIfError(golock.Lock(updateLockPath + "." + name))
}

// unlock a plugin
func unlockPlugin(name string) {
	LogIfError(golock.Unlock(updateLockPath + "." + name))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"syscall"

	"github.com/ansel1/merry"
	"github.com/dickeyxxx/golock"
)

func init() {
	Topics = append(Topics, &Topic{
		Name:        "plugins",
		Description: "manage plugins",
		Commands: CommandSet{
			{
				Topic:            "plugins",
				Hidden:           true,
				Description:      "Lists installed plugins",
				DisableAnalytics: true,
				Flags: []Flag{
					{Name: "core", Description: "show core plugins"},
				},
				Help: `
Example:
  $ heroku plugins`,

				Run: pluginsList,
			},
			{
				Topic:        "plugins",
				Command:      "install",
				Hidden:       true,
				VariableArgs: true,
				Description:  "Installs a plugin into the CLI",
				Help: `Install a Heroku plugin

  Example:
  $ heroku plugins:install heroku-production-status`,

				Run: pluginsInstall,
			},
			{
				Topic:       "plugins",
				Command:     "link",
				Description: "Links a local plugin into CLI",
				Args:        []Arg{{Name: "path", Optional: true}},
				Help: `Links a local plugin into CLI.
	This is useful when developing plugins locally.
	It simply symlinks the specified path into the plugins directory
	and parses the plugin.

	You will need to run it again if you change any of the plugin metadata.

  Example:
	$ heroku plugins:link .`,

				Run: pluginsLink,
			},
			{
				Topic:       "plugins",
				Command:     "uninstall",
				Hidden:      true,
				Args:        []Arg{{Name: "name"}},
				Description: "Uninstalls a plugin from the CLI",
				Help: `Uninstalls a Heroku plugin

  Example:
  $ heroku plugins:uninstall heroku-production-status`,

				Run: pluginsUninstall,
			},
		},
	})
}

func pluginsList(ctx *Context) {
	var names []string
	for _, plugin := range UserPlugins.Plugins() {
		symlinked := ""
		if UserPlugins.isPluginSymlinked(plugin.Name) {
			symlinked = " (symlinked)"
		}
		names = append(names, fmt.Sprintf("%s %s%s", plugin.Name, plugin.Version, symlinked))
	}
	if ctx.Flags["core"] != nil {
		UserPluginNames := UserPlugins.PluginNames()
		for _, plugin := range CorePlugins.Plugins() {
			if contains(UserPluginNames, plugin.Name) {
				continue
			}
			names = append(names, fmt.Sprintf("%s %s (core)", plugin.Name, plugin.Version))
		}
	}
	sort.Strings(names)
	for _, plugin := range names {
		Println(plugin)
	}
}
func pluginsInstall(ctx *Context) {
	plugins := ctx.Args.([]string)
	if len(plugins) == 0 {
		ExitWithMessage("Must specify a plugin name.\nUSAGE: heroku plugins:install heroku-debug")
	}
	toinstall := make([]string, 0, len(plugins))
	core := CorePlugins.PluginNames()
	for _, plugin := range plugins {
		if contains(core, strings.Split(plugin, "@")[0]) {
			Warn("Not installing " + plugin + " because it is already installed as a core plugin.")
			continue
		}
		toinstall = append(toinstall, plugin)
	}
	if len(toinstall) == 0 {
		Exit(0)
	}
	action("Installing "+plural("plugin", len(toinstall))+" "+strings.Join(toinstall, " "), "done", func() {
		err := UserPlugins.InstallPlugins(toinstall...)
		if err != nil {
			if strings.Contains(err.Error(), "no such package available") {
				ExitWithMessage("Plugin not found")
			}
			must(err)
		}
	})
}

func pluginsLink(ctx *Context) {
	path := ctx.Args.(map[string]string)["path"]
	if path == "" {
		path = "."
	}
	path, err := filepath.Abs(path)
	must(err)
	_, err = os.Stat(path)
	must(err)
	name := filepath.Base(path)
	action("Symlinking "+name, "done", func() {
		newPath := UserPlugins.pluginPath(name)
		os.Remove(newPath)
		os.RemoveAll(newPath)
		os.MkdirAll(filepath.Dir(newPath), 0755)
		err = os.Symlink(path, newPath)
		must(err)
		plugin, err := UserPlugins.ParsePlugin(name)
		must(err)
		if name != plugin.Name {
			path = newPath
			newPath = UserPlugins.pluginPath(plugin.Name)
			os.Remove(newPath)
			os.RemoveAll(newPath)
			os.Rename(path, newPath)
		}
		UserPlugins.addToCache(plugin)
	})
}

func pluginsUninstall(ctx *Context) {
	name := ctx.Args.(map[string]string)["name"]
	if !contains(UserPlugins.PluginNames(), name) {
		ExitWithMessage("%s is not installed", name)
	}
	Errf("Uninstalling plugin %s...", name)
	must(UserPlugins.RemovePackages(name))
	UserPlugins.removeFromCache(name)
	Errln(" done")
}

// Plugins represents either core or user plugins
type Plugins struct {
	Path    string
	plugins []*Plugin
}

// CorePlugins are built in plugins
var CorePlugins = &Plugins{Path: filepath.Join(AppDir, "lib")}

// UserPlugins are user-installable plugins
var UserPlugins = &Plugins{Path: filepath.Join(DataHome, "plugins")}

// Plugin represents a javascript plugin
type Plugin struct {
	Name     string     `json:"name"`
	Version  string     `json:"version"`
	Topics   TopicSet   `json:"topics"`
	Topic    *Topic     `json:"topic"`
	Commands CommandSet `json:"commands"`
}

// Commands lists all the commands of the plugins
func (p *Plugins) Commands() (commands CommandSet) {
	for _, plugin := range p.Plugins() {
		for _, command := range plugin.Commands {
			command.Run = p.runFn(plugin, command.Topic, command.Command)
			commands = append(commands, command)
		}
	}
	return
}

// Topics gets all the plugin's topics
func (p *Plugins) Topics() (topics TopicSet) {
	for _, plugin := range p.Plugins() {
		if plugin.Topic != nil {
			topics = append(topics, plugin.Topic)
		}
		topics = append(topics, plugin.Topics...)
	}
	return
}

func (p *Plugins) runFn(plugin *Plugin, topic, command string) func(ctx *Context) {
	return func(ctx *Context) {
		p.readLockPlugin(plugin.Name)
		ctx.Dev = p.isPluginSymlinked(plugin.Name)
		ctxJSON, err := json.Marshal(ctx)
		must(err)
		args, err := json.Marshal(Args)
		must(err)
		title, _ := json.Marshal("heroku " + strings.Join(Args[1:], " "))

		script := fmt.Sprintf(`'use strict'
process.argv = %s
let pluginName = '%s'
let pluginVersion = '%s'
let topic = '%s'
let command = '%s'
process.title = %s
let ctx = %s
ctx.version = ctx.version + ' ' + pluginName + '/' + pluginVersion + ' node-' + process.version
process.chdir(ctx.cwd)
if (command === '') { command = null }
let plugin = require(pluginName)
let cmd = plugin.commands.filter((c) => c.topic === topic && c.command == command)[0]
cmd.run(ctx)
`, args, plugin.Name, plugin.Version, topic, command, string(title), ctxJSON)

		// swallow sigint since the plugin will handle it
		swallowSigint = true

		currentAnalyticsCommand.Plugin = plugin.Name
		currentAnalyticsCommand.Version = plugin.Version
		currentAnalyticsCommand.Language = fmt.Sprintf("node/" + NodeVersion)

		cmd, done := p.RunScript(script)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		err = cmd.Run()
		done()
		Exit(getExitCode(err))
	}
}

func getExitCode(err error) int {
	switch e := err.(type) {
	case nil:
		return 0
	case *exec.ExitError:
		status, ok := e.Sys().(syscall.WaitStatus)
		if !ok {
			must(err)
		}
		return status.ExitStatus()
	}
	must(err)
	return -1
}

// ParsePlugin requires the plugin's node module
// to get the commands and metadata
func (p *Plugins) ParsePlugin(name string) (*Plugin, error) {
	script := `
	var plugin = require('` + name + `');
	var pjson  = require('` + name + `/package.json');

	plugin.name    = pjson.name;
	plugin.version = pjson.version;

	console.log(JSON.stringify(plugin))`
	cmd, done := p.RunScript(script)
	cmd.Stderr = Stderr
	output, err := cmd.Output()
	done()

	if err != nil {
		return nil, merry.Errorf("Error installing plugin %s", name)
	}
	var plugin Plugin
	err = json.Unmarshal(output, &plugin)
	if err != nil {
		return nil, fmt.Errorf("Error parsing plugin: %s\n%s\n%s", name, err, string(output))
	}
	if len(plugin.Commands) == 0 {
		return nil, fmt.Errorf("Invalid plugin. No commands found.")
	}
	for _, command := range plugin.Commands {
		if command == nil {
			continue
		}
		command.Plugin = plugin.Name
		command.Help = strings.TrimSpace(command.Help)
	}
	return &plugin, nil
}

// PluginNames lists all the plugin names
func (p *Plugins) PluginNames() []string {
	plugins := p.Plugins()
	names := make([]string, 0, len(plugins))
	for _, plugin := range plugins {
		names = append(names, plugin.Name)
	}
	return names
}

// PluginNamesNotSymlinked lists all the plugin names that are not symlinked
func (p *Plugins) PluginNamesNotSymlinked() []string {
	plugins := p.PluginNames()
	names := make([]string, 0, len(plugins))
	for _, plugin := range plugins {
		if !p.isPluginSymlinked(plugin) {
			names = append(names, plugin)
		}
	}
	return names
}

func (p *Plugins) isPluginSymlinked(plugin string) bool {
	path := filepath.Join(p.modulesPath(), plugin)
	fi, err := os.Lstat(path)
	if err != nil {
		return false
	}
	return fi.Mode()&os.ModeSymlink != 0
}

func contains(arr []string, s string) bool {
	for _, a := range arr {
		if a == s {
			return true
		}
	}
	return false
}

// InstallPlugins installs plugins
func (p *Plugins) InstallPlugins(names ...string) error {
	for _, name := range names {
		p.lockPlugin(name)
	}
	defer func() {
		for _, name := range names {
			p.unlockPlugin(name)
		}
	}()
	err := p.installPackages(names...)
	if err != nil {
		return err
	}
	plugins := make([]*Plugin, len(names))
	for i, name := range names {
		plugin, err := p.ParsePlugin(name)
		if err != nil {
			return err
		}
		plugins[i] = plugin
	}
	p.addToCache(plugins...)
	return nil
}

// directory location of plugin
func (p *Plugins) pluginPath(plugin string) string {
	return filepath.Join(p.Path, "node_modules", plugin)
}

// name of lockfile
func (p *Plugins) lockfile(name string) string {
	return filepath.Join(p.Path, name+".updating")
}

// lock a plugin for reading
func (p *Plugins) readLockPlugin(name string) {
	locked, err := golock.IsLocked(p.lockfile(name))
	LogIfError(err)
	if locked {
		p.lockPlugin(name)
		p.unlockPlugin(name)
	}
}

// lock a plugin for writing
func (p *Plugins) lockPlugin(name string) {
	os.MkdirAll(filepath.Dir(p.lockfile(name)), 0755)
	LogIfError(golock.Lock(p.lockfile(name)))
}

// unlock a plugin
func (p *Plugins) unlockPlugin(name string) {
	LogIfError(golock.Unlock(p.lockfile(name)))
}

// Update updates the plugins
func (p *Plugins) Update() {
	plugins := p.PluginNamesNotSymlinked()
	if len(plugins) == 0 {
		return
	}
	packages, err := p.OutdatedPackages(plugins...)
	WarnIfError(err)
	if len(packages) > 0 {
		action("heroku-cli: Updating plugins", "", func() {
			for name, version := range packages {
				p.lockPlugin(name)
				WarnIfError(p.installPackages(name + "@" + version))
				plugin, err := p.ParsePlugin(name)
				WarnIfError(err)
				p.addToCache(plugin)
				p.unlockPlugin(name)
			}
		})
		Errf(" done. Updated %d %s.\n", len(packages), plural("package", len(packages)))
	}
}

func (p *Plugins) addToCache(plugins ...*Plugin) {
	contains := func(name string) int {
		for i, plugin := range p.plugins {
			if plugin.Name == name {
				return i
			}
		}
		return -1
	}
	for _, plugin := range plugins {
		// find or replace
		i := contains(plugin.Name)
		if i == -1 {
			p.plugins = append(p.plugins, plugin)
		} else {
			p.plugins[i] = plugin
		}
	}
	p.saveCache()
}

func (p *Plugins) removeFromCache(name string) {
	for i, plugin := range p.plugins {
		if plugin.Name == name {
			p.plugins = append(p.plugins[:i], p.plugins[i+1:]...)
		}
	}
	p.saveCache()
}

func (p *Plugins) saveCache() {
	if err := saveJSON(p.plugins, p.cachePath()); err != nil {
		must(err)
	}
}

// Plugins reads the cache file into the struct
func (p *Plugins) Plugins() []*Plugin {
	if p.plugins == nil {
		p.plugins = []*Plugin{}
		if exists, _ := FileExists(p.cachePath()); !exists {
			return p.plugins
		}
		f, err := os.Open(p.cachePath())
		if err != nil {
			LogIfError(err)
			return p.plugins
		}
		err = json.NewDecoder(f).Decode(&p.plugins)
		WarnIfError(err)
	}
	return p.plugins
}

func (p *Plugins) cachePath() string {
	return filepath.Join(p.Path, "plugins.json")
}
<commit_msg>show better error on invalid plugin parsing<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"syscall"

	"github.com/ansel1/merry"
	"github.com/dickeyxxx/golock"
)

func init() {
	Topics = append(Topics, &Topic{
		Name:        "plugins",
		Description: "manage plugins",
		Commands: CommandSet{
			{
				Topic:            "plugins",
				Hidden:           true,
				Description:      "Lists installed plugins",
				DisableAnalytics: true,
				Flags: []Flag{
					{Name: "core", Description: "show core plugins"},
				},
				Help: `
Example:
  $ heroku plugins`,

				Run: pluginsList,
			},
			{
				Topic:        "plugins",
				Command:      "install",
				Hidden:       true,
				VariableArgs: true,
				Description:  "Installs a plugin into the CLI",
				Help: `Install a Heroku plugin

  Example:
  $ heroku plugins:install heroku-production-status`,

				Run: pluginsInstall,
			},
			{
				Topic:       "plugins",
				Command:     "link",
				Description: "Links a local plugin into CLI",
				Args:        []Arg{{Name: "path", Optional: true}},
				Help: `Links a local plugin into CLI.
	This is useful when developing plugins locally.
	It simply symlinks the specified path into the plugins directory
	and parses the plugin.

	You will need to run it again if you change any of the plugin metadata.

  Example:
	$ heroku plugins:link .`,

				Run: pluginsLink,
			},
			{
				Topic:       "plugins",
				Command:     "uninstall",
				Hidden:      true,
				Args:        []Arg{{Name: "name"}},
				Description: "Uninstalls a plugin from the CLI",
				Help: `Uninstalls a Heroku plugin

  Example:
  $ heroku plugins:uninstall heroku-production-status`,

				Run: pluginsUninstall,
			},
		},
	})
}

func pluginsList(ctx *Context) {
	var names []string
	for _, plugin := range UserPlugins.Plugins() {
		symlinked := ""
		if UserPlugins.isPluginSymlinked(plugin.Name) {
			symlinked = " (symlinked)"
		}
		names = append(names, fmt.Sprintf("%s %s%s", plugin.Name, plugin.Version, symlinked))
	}
	if ctx.Flags["core"] != nil {
		UserPluginNames := UserPlugins.PluginNames()
		for _, plugin := range CorePlugins.Plugins() {
			if contains(UserPluginNames, plugin.Name) {
				continue
			}
			names = append(names, fmt.Sprintf("%s %s (core)", plugin.Name, plugin.Version))
		}
	}
	sort.Strings(names)
	for _, plugin := range names {
		Println(plugin)
	}
}
func pluginsInstall(ctx *Context) {
	plugins := ctx.Args.([]string)
	if len(plugins) == 0 {
		ExitWithMessage("Must specify a plugin name.\nUSAGE: heroku plugins:install heroku-debug")
	}
	toinstall := make([]string, 0, len(plugins))
	core := CorePlugins.PluginNames()
	for _, plugin := range plugins {
		if contains(core, strings.Split(plugin, "@")[0]) {
			Warn("Not installing " + plugin + " because it is already installed as a core plugin.")
			continue
		}
		toinstall = append(toinstall, plugin)
	}
	if len(toinstall) == 0 {
		Exit(0)
	}
	action("Installing "+plural("plugin", len(toinstall))+" "+strings.Join(toinstall, " "), "done", func() {
		err := UserPlugins.InstallPlugins(toinstall...)
		if err != nil {
			if strings.Contains(err.Error(), "no such package available") {
				ExitWithMessage("Plugin not found")
			}
			must(err)
		}
	})
}

func pluginsLink(ctx *Context) {
	path := ctx.Args.(map[string]string)["path"]
	if path == "" {
		path = "."
	}
	path, err := filepath.Abs(path)
	must(err)
	_, err = os.Stat(path)
	must(err)
	name := filepath.Base(path)
	action("Symlinking "+name, "done", func() {
		newPath := UserPlugins.pluginPath(name)
		os.Remove(newPath)
		os.RemoveAll(newPath)
		os.MkdirAll(filepath.Dir(newPath), 0755)
		err = os.Symlink(path, newPath)
		must(err)
		plugin, err := UserPlugins.ParsePlugin(name)
		must(err)
		if name != plugin.Name {
			path = newPath
			newPath = UserPlugins.pluginPath(plugin.Name)
			os.Remove(newPath)
			os.RemoveAll(newPath)
			os.Rename(path, newPath)
		}
		UserPlugins.addToCache(plugin)
	})
}

func pluginsUninstall(ctx *Context) {
	name := ctx.Args.(map[string]string)["name"]
	if !contains(UserPlugins.PluginNames(), name) {
		ExitWithMessage("%s is not installed", name)
	}
	Errf("Uninstalling plugin %s...", name)
	must(UserPlugins.RemovePackages(name))
	UserPlugins.removeFromCache(name)
	Errln(" done")
}

// Plugins represents either core or user plugins
type Plugins struct {
	Path    string
	plugins []*Plugin
}

// CorePlugins are built in plugins
var CorePlugins = &Plugins{Path: filepath.Join(AppDir, "lib")}

// UserPlugins are user-installable plugins
var UserPlugins = &Plugins{Path: filepath.Join(DataHome, "plugins")}

// Plugin represents a javascript plugin
type Plugin struct {
	Name     string     `json:"name"`
	Version  string     `json:"version"`
	Topics   TopicSet   `json:"topics"`
	Topic    *Topic     `json:"topic"`
	Commands CommandSet `json:"commands"`
}

// Commands lists all the commands of the plugins
func (p *Plugins) Commands() (commands CommandSet) {
	for _, plugin := range p.Plugins() {
		for _, command := range plugin.Commands {
			command.Run = p.runFn(plugin, command.Topic, command.Command)
			commands = append(commands, command)
		}
	}
	return
}

// Topics gets all the plugin's topics
func (p *Plugins) Topics() (topics TopicSet) {
	for _, plugin := range p.Plugins() {
		if plugin.Topic != nil {
			topics = append(topics, plugin.Topic)
		}
		topics = append(topics, plugin.Topics...)
	}
	return
}

func (p *Plugins) runFn(plugin *Plugin, topic, command string) func(ctx *Context) {
	return func(ctx *Context) {
		p.readLockPlugin(plugin.Name)
		ctx.Dev = p.isPluginSymlinked(plugin.Name)
		ctxJSON, err := json.Marshal(ctx)
		must(err)
		args, err := json.Marshal(Args)
		must(err)
		title, _ := json.Marshal("heroku " + strings.Join(Args[1:], " "))

		script := fmt.Sprintf(`'use strict'
process.argv = %s
let pluginName = '%s'
let pluginVersion = '%s'
let topic = '%s'
let command = '%s'
process.title = %s
let ctx = %s
ctx.version = ctx.version + ' ' + pluginName + '/' + pluginVersion + ' node-' + process.version
process.chdir(ctx.cwd)
if (command === '') { command = null }
let plugin = require(pluginName)
let cmd = plugin.commands.filter((c) => c.topic === topic && c.command == command)[0]
cmd.run(ctx)
`, args, plugin.Name, plugin.Version, topic, command, string(title), ctxJSON)

		// swallow sigint since the plugin will handle it
		swallowSigint = true

		currentAnalyticsCommand.Plugin = plugin.Name
		currentAnalyticsCommand.Version = plugin.Version
		currentAnalyticsCommand.Language = fmt.Sprintf("node/" + NodeVersion)

		cmd, done := p.RunScript(script)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		err = cmd.Run()
		done()
		Exit(getExitCode(err))
	}
}

func getExitCode(err error) int {
	switch e := err.(type) {
	case nil:
		return 0
	case *exec.ExitError:
		status, ok := e.Sys().(syscall.WaitStatus)
		if !ok {
			must(err)
		}
		return status.ExitStatus()
	}
	must(err)
	return -1
}

// ParsePlugin requires the plugin's node module
// to get the commands and metadata
func (p *Plugins) ParsePlugin(name string) (*Plugin, error) {
	script := `
	var plugin = require('` + name + `')
	var pjson  = require('` + name + `/package.json')

	plugin.name    = pjson.name
	plugin.version = pjson.version

	console.log(JSON.stringify(plugin))`
	cmd, done := p.RunScript(script)
	cmd.Stderr = Stderr
	output, err := cmd.Output()
	done()

	if err != nil {
		return nil, merry.Errorf("Error installing plugin %s", name)
	}
	var plugin Plugin
	err = json.Unmarshal(output, &plugin)
	if err != nil {
		return nil, fmt.Errorf("Error parsing plugin: %s\n%s\n%s\nIs this a real CLI plugin?", name, err, string(output))
	}
	if len(plugin.Commands) == 0 {
		return nil, fmt.Errorf("Invalid plugin. No commands found.")
	}
	for _, command := range plugin.Commands {
		if command == nil {
			continue
		}
		command.Plugin = plugin.Name
		command.Help = strings.TrimSpace(command.Help)
	}
	return &plugin, nil
}

// PluginNames lists all the plugin names
func (p *Plugins) PluginNames() []string {
	plugins := p.Plugins()
	names := make([]string, 0, len(plugins))
	for _, plugin := range plugins {
		names = append(names, plugin.Name)
	}
	return names
}

// PluginNamesNotSymlinked lists all the plugin names that are not symlinked
func (p *Plugins) PluginNamesNotSymlinked() []string {
	plugins := p.PluginNames()
	names := make([]string, 0, len(plugins))
	for _, plugin := range plugins {
		if !p.isPluginSymlinked(plugin) {
			names = append(names, plugin)
		}
	}
	return names
}

func (p *Plugins) isPluginSymlinked(plugin string) bool {
	path := filepath.Join(p.modulesPath(), plugin)
	fi, err := os.Lstat(path)
	if err != nil {
		return false
	}
	return fi.Mode()&os.ModeSymlink != 0
}

func contains(arr []string, s string) bool {
	for _, a := range arr {
		if a == s {
			return true
		}
	}
	return false
}

// InstallPlugins installs plugins
func (p *Plugins) InstallPlugins(names ...string) error {
	for _, name := range names {
		p.lockPlugin(name)
	}
	defer func() {
		for _, name := range names {
			p.unlockPlugin(name)
		}
	}()
	err := p.installPackages(names...)
	if err != nil {
		return err
	}
	plugins := make([]*Plugin, len(names))
	for i, name := range names {
		plugin, err := p.ParsePlugin(name)
		if err != nil {
			return err
		}
		plugins[i] = plugin
	}
	p.addToCache(plugins...)
	return nil
}

// directory location of plugin
func (p *Plugins) pluginPath(plugin string) string {
	return filepath.Join(p.Path, "node_modules", plugin)
}

// name of lockfile
func (p *Plugins) lockfile(name string) string {
	return filepath.Join(p.Path, name+".updating")
}

// lock a plugin for reading
func (p *Plugins) readLockPlugin(name string) {
	locked, err := golock.IsLocked(p.lockfile(name))
	LogIfError(err)
	if locked {
		p.lockPlugin(name)
		p.unlockPlugin(name)
	}
}

// lock a plugin for writing
func (p *Plugins) lockPlugin(name string) {
	os.MkdirAll(filepath.Dir(p.lockfile(name)), 0755)
	LogIfError(golock.Lock(p.lockfile(name)))
}

// unlock a plugin
func (p *Plugins) unlockPlugin(name string) {
	LogIfError(golock.Unlock(p.lockfile(name)))
}

// Update updates the plugins
func (p *Plugins) Update() {
	plugins := p.PluginNamesNotSymlinked()
	if len(plugins) == 0 {
		return
	}
	packages, err := p.OutdatedPackages(plugins...)
	WarnIfError(err)
	if len(packages) > 0 {
		action("heroku-cli: Updating plugins", "", func() {
			for name, version := range packages {
				p.lockPlugin(name)
				WarnIfError(p.installPackages(name + "@" + version))
				plugin, err := p.ParsePlugin(name)
				WarnIfError(err)
				p.addToCache(plugin)
				p.unlockPlugin(name)
			}
		})
		Errf(" done. Updated %d %s.\n", len(packages), plural("package", len(packages)))
	}
}

func (p *Plugins) addToCache(plugins ...*Plugin) {
	contains := func(name string) int {
		for i, plugin := range p.plugins {
			if plugin.Name == name {
				return i
			}
		}
		return -1
	}
	for _, plugin := range plugins {
		// find or replace
		i := contains(plugin.Name)
		if i == -1 {
			p.plugins = append(p.plugins, plugin)
		} else {
			p.plugins[i] = plugin
		}
	}
	p.saveCache()
}

func (p *Plugins) removeFromCache(name string) {
	for i, plugin := range p.plugins {
		if plugin.Name == name {
			p.plugins = append(p.plugins[:i], p.plugins[i+1:]...)
		}
	}
	p.saveCache()
}

func (p *Plugins) saveCache() {
	if err := saveJSON(p.plugins, p.cachePath()); err != nil {
		must(err)
	}
}

// Plugins reads the cache file into the struct
func (p *Plugins) Plugins() []*Plugin {
	if p.plugins == nil {
		p.plugins = []*Plugin{}
		if exists, _ := FileExists(p.cachePath()); !exists {
			return p.plugins
		}
		f, err := os.Open(p.cachePath())
		if err != nil {
			LogIfError(err)
			return p.plugins
		}
		err = json.NewDecoder(f).Decode(&p.plugins)
		WarnIfError(err)
	}
	return p.plugins
}

func (p *Plugins) cachePath() string {
	return filepath.Join(p.Path, "plugins.json")
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/coopernurse/gorp"
	. "gopkg.in/check.v1"
)

func createTestDatabase() *gorp.DbMap {
	db, err := sql.Open("sqlite3", ":memory:")
	checkErr(err, "sql.Open failed")

	// construct a gorp DbMap
	dbmap := &gorp.DbMap{Db: db, Dialect: gorp.SqliteDialect{}}

	return dbmap
}

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

type PRSuite struct {
	dbmap     *gorp.DbMap
	app       Application
	createURL string
	tableName string
}

func (s *PRSuite) SetUpSuite(c *C) {
	s.dbmap = createTestDatabase()
	s.app = CreateApplication(s.dbmap)
}

func (s *PRSuite) SetUpTest(c *C) {
	err := s.dbmap.TruncateTables()
	if err != nil {
		c.Error(err)
	}
}

func (s *PRSuite) TearDownSuite(c *C) {
	s.dbmap.Db.Close()
}

func (s *PRSuite) PerformRequest(method string, relativePath string, body string) *httptest.ResponseRecorder {
	path := fmt.Sprintf("http://test.example.com%s", relativePath)
	w := httptest.NewRecorder()
	r, err := http.NewRequest(method, path, strings.NewReader(body))
	checkErr(err, "Request creation failed")

	s.app.handler.ServeHTTP(w, r)
	return w
}

func (s *PRSuite) TestAddReturns201(c *C) {
	recorder := s.PerformRequest("POST", s.createURL, `{"name": "Test Name"}`)

	c.Check(recorder.Code, Equals, 201)
}

func (s *PRSuite) TestAddCreatesOneEntity(c *C) {
	s.PerformRequest("POST", s.createURL, `{"name": "Test Name"}`)

	query := fmt.Sprintf("SELECT count(*) FROM %s", s.tableName)
	count, err := s.dbmap.SelectInt(query)
	checkErr(err, "Getting count failed")
	c.Assert(count, Equals, int64(1))
}

type ElectionSuite struct {
	PRSuite
}

func (s *ElectionSuite) SetUpSuite(c *C) {
	s.PRSuite.SetUpSuite(c)
	s.createURL = "/elections"
	s.tableName = "elections"
}

var _ = Suite(&ElectionSuite{})

func (s *ElectionSuite) TestAddElectionCreatesElectionWithCorrectName(c *C) {
	s.PerformRequest("POST", "/elections", `{"name": "Test Election"}`)

	var createdElection Election
	err := s.dbmap.SelectOne(&createdElection, "select * from elections")
	if err != nil {
		c.Error(err)
	}
	c.Assert(createdElection.Name, Matches, "Test Election")
}

func (s *ElectionSuite) TestAddElectionRejectsZeroLengthName(c *C) {
	recorder := s.PerformRequest("POST", "/elections", `{"name": ""}`)

	c.Check(recorder.Code, Equals, 400)
	c.Check(recorder.Body.String(), Matches, "Empty name forbidden.\n?")

	count, err := s.dbmap.SelectInt("select count(*) from elections")
	checkErr(err, "Getting count failed")
	c.Check(count, Equals, int64(0))
}

func (s *ElectionSuite) TestAddElectionRejectsDuplicateNames(c *C) {
	s.PerformRequest("POST", "/elections", `{"name": "Duplicate"}`)
	recorder := s.PerformRequest("POST", "/elections", `{"name": "Duplicate"}`)

	c.Check(recorder.Code, Equals, 400)
	c.Check(recorder.Body.String(), Matches, "Name taken.\n?")

	count, err := s.dbmap.SelectInt("select count(*) from elections")
	checkErr(err, "Getting count failed")
	c.Check(count, Equals, int64(1))
}

func (s *ElectionSuite) TestGetElectionReturns200(c *C) {
	election := Election{Name: "my test name"}
	s.dbmap.Insert(&election)

	recorder := s.PerformRequest("GET", fmt.Sprintf("/elections/%d", election.Id), "")

	c.Check(recorder.Code, Equals, 200)
}

func (s *ElectionSuite) TestGetElectionReturnsElectionName(c *C) {
	election := Election{Name: "my test name"}
	s.dbmap.Insert(&election)

	recorder := s.PerformRequest("GET", fmt.Sprintf("/elections/%d", election.Id), "")
	returnedElection := Election{}
	json.Unmarshal(recorder.Body.Bytes(), &returnedElection)
	c.Assert(returnedElection.Name, Matches, "my test name")
}

func (s *ElectionSuite) TestGetElection404sForUnknownElection(c *C) {
	recorder := s.PerformRequest("GET", "/elections/1", "")

	c.Check(recorder.Code, Equals, 404)
}

func (s *ElectionSuite) TestListElectionsReturnsEmptyList(c *C) {
	recorder := s.PerformRequest("GET", "/elections", "")

	c.Check(recorder.Code, Equals, 200)
	c.Check(recorder.Body.String(), Equals, "[]")
}

func (s *ElectionSuite) TestListElectionsReturnsListOfCorrectLength(c *C) {
	election := Election{Name: "my test name"}
	other_election := Election{Name: "my other name"}
	third_election := Election{Name: "my third name"}
	s.dbmap.Insert(&election, &other_election, &third_election)

	recorder := s.PerformRequest("GET", "/elections", "")

	var electionList []Election
	json.Unmarshal(recorder.Body.Bytes(), &electionList)
	c.Check(len(electionList), Equals, 3)
}

func (s *ElectionSuite) TestListElectionReturnsExistingElections(c *C) {
	election := Election{Name: "my test name"}
	other_election := Election{Name: "my other name"}
	s.dbmap.Insert(&election, &other_election)

	recorder := s.PerformRequest("GET", "/elections", "")

	expectedElectionNames := map[string]int{
		"my test name":  0,
		"my other name": 0,
	}
	var electionList []Election
	json.Unmarshal(recorder.Body.Bytes(), &electionList)
	actualElectionNames := make(map[string]int)
	for _, election := range electionList {
		actualElectionNames[election.Name] = 0
	}
	c.Check(actualElectionNames, DeepEquals, expectedElectionNames)
}

type CandidatesSuite struct {
	PRSuite
}

func (s *CandidatesSuite) SetUpTest(c *C) {
	s.PRSuite.SetUpTest(c)

	// Set up test election
	election := Election{Name: "my test name"}
	err := s.dbmap.Insert(&election)
	if err != nil {
		c.Error(err)
	}

	s.createURL = fmt.Sprintf("/elections/%d/candidates", election.Id)
	s.tableName = "candidates"
}

var _ = Suite(&CandidatesSuite{})

func (s *CandidatesSuite) TestAddCandidateReturns404ForMissingElection(c *C) {
	recorder := s.PerformRequest("POST", "/elections/1234/candidates", `{"name": "Test Candidate"}`)

	c.Check(recorder.Code, Equals, 404)
}
<commit_msg>Add another test for candidate creation.<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/coopernurse/gorp"
	. "gopkg.in/check.v1"
)

func createTestDatabase() *gorp.DbMap {
	db, err := sql.Open("sqlite3", ":memory:")
	checkErr(err, "sql.Open failed")

	// construct a gorp DbMap
	dbmap := &gorp.DbMap{Db: db, Dialect: gorp.SqliteDialect{}}

	return dbmap
}

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

type PRSuite struct {
	dbmap     *gorp.DbMap
	app       Application
	createURL string
	tableName string
}

func (s *PRSuite) SetUpSuite(c *C) {
	s.dbmap = createTestDatabase()
	s.app = CreateApplication(s.dbmap)
}

func (s *PRSuite) SetUpTest(c *C) {
	err := s.dbmap.TruncateTables()
	if err != nil {
		c.Error(err)
	}
}

func (s *PRSuite) TearDownSuite(c *C) {
	s.dbmap.Db.Close()
}

func (s *PRSuite) PerformRequest(method string, relativePath string, body string) *httptest.ResponseRecorder {
	path := fmt.Sprintf("http://test.example.com%s", relativePath)
	w := httptest.NewRecorder()
	r, err := http.NewRequest(method, path, strings.NewReader(body))
	checkErr(err, "Request creation failed")

	s.app.handler.ServeHTTP(w, r)
	return w
}

func (s *PRSuite) TestAddReturns201(c *C) {
	recorder := s.PerformRequest("POST", s.createURL, `{"name": "Test Name"}`)

	c.Check(recorder.Code, Equals, 201)
}

func (s *PRSuite) TestAddCreatesOneEntity(c *C) {
	s.PerformRequest("POST", s.createURL, `{"name": "Test Name"}`)

	query := fmt.Sprintf("SELECT count(*) FROM %s", s.tableName)
	count, err := s.dbmap.SelectInt(query)
	checkErr(err, "Getting count failed")
	c.Assert(count, Equals, int64(1))
}

type ElectionSuite struct {
	PRSuite
}

func (s *ElectionSuite) SetUpSuite(c *C) {
	s.PRSuite.SetUpSuite(c)
	s.createURL = "/elections"
	s.tableName = "elections"
}

var _ = Suite(&ElectionSuite{})

func (s *ElectionSuite) TestAddElectionCreatesElectionWithCorrectName(c *C) {
	s.PerformRequest("POST", "/elections", `{"name": "Test Election"}`)

	var createdElection Election
	err := s.dbmap.SelectOne(&createdElection, "select * from elections")
	if err != nil {
		c.Error(err)
	}
	c.Assert(createdElection.Name, Matches, "Test Election")
}

func (s *ElectionSuite) TestAddElectionRejectsZeroLengthName(c *C) {
	recorder := s.PerformRequest("POST", "/elections", `{"name": ""}`)

	c.Check(recorder.Code, Equals, 400)
	c.Check(recorder.Body.String(), Matches, "Empty name forbidden.\n?")

	count, err := s.dbmap.SelectInt("select count(*) from elections")
	checkErr(err, "Getting count failed")
	c.Check(count, Equals, int64(0))
}

func (s *ElectionSuite) TestAddElectionRejectsDuplicateNames(c *C) {
	s.PerformRequest("POST", "/elections", `{"name": "Duplicate"}`)
	recorder := s.PerformRequest("POST", "/elections", `{"name": "Duplicate"}`)

	c.Check(recorder.Code, Equals, 400)
	c.Check(recorder.Body.String(), Matches, "Name taken.\n?")

	count, err := s.dbmap.SelectInt("select count(*) from elections")
	checkErr(err, "Getting count failed")
	c.Check(count, Equals, int64(1))
}

func (s *ElectionSuite) TestGetElectionReturns200(c *C) {
	election := Election{Name: "my test name"}
	s.dbmap.Insert(&election)

	recorder := s.PerformRequest("GET", fmt.Sprintf("/elections/%d", election.Id), "")

	c.Check(recorder.Code, Equals, 200)
}

func (s *ElectionSuite) TestGetElectionReturnsElectionName(c *C) {
	election := Election{Name: "my test name"}
	s.dbmap.Insert(&election)

	recorder := s.PerformRequest("GET", fmt.Sprintf("/elections/%d", election.Id), "")
	returnedElection := Election{}
	json.Unmarshal(recorder.Body.Bytes(), &returnedElection)
	c.Assert(returnedElection.Name, Matches, "my test name")
}

func (s *ElectionSuite) TestGetElection404sForUnknownElection(c *C) {
	recorder := s.PerformRequest("GET", "/elections/1", "")

	c.Check(recorder.Code, Equals, 404)
}

func (s *ElectionSuite) TestListElectionsReturnsEmptyList(c *C) {
	recorder := s.PerformRequest("GET", "/elections", "")

	c.Check(recorder.Code, Equals, 200)
	c.Check(recorder.Body.String(), Equals, "[]")
}

func (s *ElectionSuite) TestListElectionsReturnsListOfCorrectLength(c *C) {
	election := Election{Name: "my test name"}
	other_election := Election{Name: "my other name"}
	third_election := Election{Name: "my third name"}
	s.dbmap.Insert(&election, &other_election, &third_election)

	recorder := s.PerformRequest("GET", "/elections", "")

	var electionList []Election
	json.Unmarshal(recorder.Body.Bytes(), &electionList)
	c.Check(len(electionList), Equals, 3)
}

func (s *ElectionSuite) TestListElectionReturnsExistingElections(c *C) {
	election := Election{Name: "my test name"}
	other_election := Election{Name: "my other name"}
	s.dbmap.Insert(&election, &other_election)

	recorder := s.PerformRequest("GET", "/elections", "")

	expectedElectionNames := map[string]int{
		"my test name":  0,
		"my other name": 0,
	}
	var electionList []Election
	json.Unmarshal(recorder.Body.Bytes(), &electionList)
	actualElectionNames := make(map[string]int)
	for _, election := range electionList {
		actualElectionNames[election.Name] = 0
	}
	c.Check(actualElectionNames, DeepEquals, expectedElectionNames)
}

type CandidatesSuite struct {
	PRSuite
}

func (s *CandidatesSuite) SetUpTest(c *C) {
	s.PRSuite.SetUpTest(c)

	// Set up test election
	election := Election{Name: "my test name"}
	err := s.dbmap.Insert(&election)
	if err != nil {
		c.Error(err)
	}

	s.createURL = fmt.Sprintf("/elections/%d/candidates", election.Id)
	s.tableName = "candidates"
}

var _ = Suite(&CandidatesSuite{})

func (s *CandidatesSuite) TestAddCandidateReturns404ForMissingElection(c *C) {
	recorder := s.PerformRequest("POST", "/elections/1234/candidates", `{"name": "Test Candidate"}`)

	c.Check(recorder.Code, Equals, 404)
}

func (s *CandidatesSuite) TestAddCandidateCreatesCandidateWithCorrectName(c *C) {
	s.PerformRequest("POST", s.createURL, `{"name": "Test Candidate"}`)

	var createdCandidate Candidate
	err := s.dbmap.SelectOne(&createdCandidate, "select * from candidates")
	if err != nil {
		c.Error(err)
	}
	c.Assert(createdCandidate.Name, Matches, "Test Candidate")
}
<|endoftext|>