<commit_before>package godbg

import (
	"fmt"
	"os"
	"runtime"
	"testing"
	. "github.com/smartystreets/goconvey/convey"
)

func TestProject(t *testing.T) {
	SkipConvey("Test buffers", t, func() {

		Convey("By Default, equals to std", func() {
			So(Out(), ShouldEqual, os.Stdout)
			So(Err(), ShouldEqual, os.Stderr)
		})
		Convey("When set to buffer, no longer equals to std", func() {
			SetBuffers(nil)
			So(Out(), ShouldNotEqual, os.Stdout)
			So(Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance buffer equals to std", func() {
			apdbg := NewPdbg()
			So(apdbg.Out(), ShouldEqual, os.Stdout)
			So(apdbg.Err(), ShouldEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance set to buffer writes no longer equals to std", func() {
			apdbg := NewPdbg(SetBuffers)
			So(apdbg.Out(), ShouldNotEqual, os.Stdout)
			So(apdbg.Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("Test custom buffer on global pdbg", func() {
			pdbg.bout = nil
			pdbg.sout = nil
			pdbg.berr = nil
			pdbg.serr = nil
			fmt.Fprintln(Out(), "test0 content0")
			So(OutString(), ShouldEqual, ``)
			fmt.Fprintln(Err(), "err0 content0")
			So(ErrString(), ShouldEqual, ``)
			SetBuffers(nil)
			fmt.Fprintln(Out(), "test content")
			fmt.Fprintln(Err(), "err1 cerr")
			fmt.Fprintln(Err(), "err2 cerr2")
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test content
test2 content2`)
			So(ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})

		Convey("Test custom buffer reset on global pdbg", func() {
			SetBuffers(nil)
			fmt.Fprint(Out(), "test content")
			So(OutString(), ShouldEqual, `test content`)
			fmt.Fprint(Err(), "err1 cerr")
			So(ErrString(), ShouldEqual, `err1 cerr`)
			ResetIOs()
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(Err(), "err2 cerr2")
			So(ErrString(), ShouldEqual, `err2 cerr2`)
		})

		Convey("Test custom buffer on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprintln(apdbg.Out(), "test content")
			fmt.Fprintln(apdbg.Err(), "err1 cerr")
			fmt.Fprintln(apdbg.Err(), "err2 cerr2")
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test content
test2 content2`)
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})
		Convey("Test custom buffer reset on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprint(apdbg.Out(), "test content")
			So(apdbg.OutString(), ShouldEqual, `test content`)
			fmt.Fprint(apdbg.Err(), "err1 cerr")
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr`)
			apdbg.ResetIOs()
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(apdbg.Err(), "err2 cerr2")
			So(apdbg.ErrString(), ShouldEqual, `err2 cerr2`)
		})
	})

	Convey("Test pdbg print functions", t, func() {
		Convey("Test pdbg print with global instance", func() {
			SetBuffers(nil)
			Pdbgf("test")
			So(ErrString(), ShouldEqual,
				`[func.012:96]
  test
`)
			ResetIOs()
			prbgtest()
			So(ErrString(), ShouldEqual,
				`  [prbgtest:4] (func.012:102)
    prbgtest content
`)
		})

		Convey("Test pdbg print with custom instance", func() {
			apdbg := NewPdbg(SetBuffers)
			apdbg.Pdbgf("test2")
			So(apdbg.ErrString(), ShouldEqual,
				`[func.013:111]
  test2
`)
			apdbg.ResetIOs()
			prbgtestCustom(apdbg)
			So(apdbg.ErrString(), ShouldEqual,
				`  [prbgtestCustom:8] (func.013:117)
    prbgtest content2
`)
			apdbg.ResetIOs()
			apdbg.pdbgTestInstance()
			So(apdbg.ErrString(), ShouldEqual,
				`  [*Pdbg.pdbgTestInstance:12] (func.013:123)
    pdbgTestInstance content3
`)
		})
		Convey("Test pdbg prints nothing if runtime.Caller fails", func() {
			mycaller = failCaller
			apdbg := NewPdbg(SetBuffers)
			apdbg.Pdbgf("test fail")
			So(apdbg.ErrString(), ShouldEqual, `  test fail
`)
			mycaller = runtime.Caller
		})
	})

	SkipConvey("Test pdbg excludes functions", t, func() {
		Convey("Test pdbg exclude with global instance", func() {
			SetBuffers(nil)
			pdbg.SetExcludes([]string{"globalNo"})
			globalPdbgExcludeTest()
			So(ErrString(), ShouldEqual,
				`  [globalPdbgExcludeTest:16] (func.016:143)
    calling no
      [globalCNo:26] (globalPdbgExcludeTest:17) (func.016:143)
        gcalled2
`)
		})
		Convey("Test pdbg exclude with custom instance", func() {
			apdbg := NewPdbg(SetBuffers, OptExcludes([]string{"customNo"}))
			customPdbgExcludeTest(apdbg)
			So(apdbg.ErrString(), ShouldEqual,
				`  [customPdbgExcludeTest:30] (func.017:153)
    calling cno
      [customCNo:40] (customPdbgExcludeTest:31) (func.017:153)
        ccalled2
`)
		})
	})

	SkipConvey("Test pdbg skips functions", t, func() {
		Convey("Test pdbg skip with global instance", func() {
			SetBuffers(nil)
			pdbg.SetSkips([]string{"globalNo"})
			globalPdbgExcludeTest()
			So(ErrString(), ShouldEqual,
				`  [globalPdbgExcludeTest:16] (func.019:167)
    calling no
    [globalPdbgExcludeTest:17] (func.019:167)
      gcalled1
      [globalCNo:26] (globalPdbgExcludeTest:17) (func.019:167)
        gcalled2
`)
		})
	})
}

func failCaller(skip int) (pc uintptr, file string, line int, ok bool) {
	return 0, "fail", skip, false
}
<commit_msg>Fix exclude test<commit_after>package godbg

import (
	"fmt"
	"os"
	"runtime"
	"testing"
	. "github.com/smartystreets/goconvey/convey"
)

func TestProject(t *testing.T) {
	SkipConvey("Test buffers", t, func() {

		Convey("By Default, equals to std", func() {
			So(Out(), ShouldEqual, os.Stdout)
			So(Err(), ShouldEqual, os.Stderr)
		})
		Convey("When set to buffer, no longer equals to std", func() {
			SetBuffers(nil)
			So(Out(), ShouldNotEqual, os.Stdout)
			So(Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance buffer equals to std", func() {
			apdbg := NewPdbg()
			So(apdbg.Out(), ShouldEqual, os.Stdout)
			So(apdbg.Err(), ShouldEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance set to buffer writes no longer equals to std", func() {
			apdbg := NewPdbg(SetBuffers)
			So(apdbg.Out(), ShouldNotEqual, os.Stdout)
			So(apdbg.Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("Test custom buffer on global pdbg", func() {
			pdbg.bout = nil
			pdbg.sout = nil
			pdbg.berr = nil
			pdbg.serr = nil
			fmt.Fprintln(Out(), "test0 content0")
			So(OutString(), ShouldEqual, ``)
			fmt.Fprintln(Err(), "err0 content0")
			So(ErrString(), ShouldEqual, ``)
			SetBuffers(nil)
			fmt.Fprintln(Out(), "test content")
			fmt.Fprintln(Err(), "err1 cerr")
			fmt.Fprintln(Err(), "err2 cerr2")
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test content
test2 content2`)
			So(ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})

		Convey("Test custom buffer reset on global pdbg", func() {
			SetBuffers(nil)
			fmt.Fprint(Out(), "test content")
			So(OutString(), ShouldEqual, `test content`)
			fmt.Fprint(Err(), "err1 cerr")
			So(ErrString(), ShouldEqual, `err1 cerr`)
			ResetIOs()
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(Err(), "err2 cerr2")
			So(ErrString(), ShouldEqual, `err2 cerr2`)
		})

		Convey("Test custom buffer on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprintln(apdbg.Out(), "test content")
			fmt.Fprintln(apdbg.Err(), "err1 cerr")
			fmt.Fprintln(apdbg.Err(), "err2 cerr2")
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test content
test2 content2`)
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})
		Convey("Test custom buffer reset on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprint(apdbg.Out(), "test content")
			So(apdbg.OutString(), ShouldEqual, `test content`)
			fmt.Fprint(apdbg.Err(), "err1 cerr")
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr`)
			apdbg.ResetIOs()
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(apdbg.Err(), "err2 cerr2")
			So(apdbg.ErrString(), ShouldEqual, `err2 cerr2`)
		})
	})

	Convey("Test pdbg print functions", t, func() {
		Convey("Test pdbg print with global instance", func() {
			SetBuffers(nil)
			Pdbgf("test")
			So(ErrString(), ShouldEqual,
				`[func.012:96]
  test
`)
			ResetIOs()
			prbgtest()
			So(ErrString(), ShouldEqual,
				`  [prbgtest:4] (func.012:102)
    prbgtest content
`)
		})

		Convey("Test pdbg print with custom instance", func() {
			apdbg := NewPdbg(SetBuffers)
			apdbg.Pdbgf("test2")
			So(apdbg.ErrString(), ShouldEqual,
				`[func.013:111]
  test2
`)
			apdbg.ResetIOs()
			prbgtestCustom(apdbg)
			So(apdbg.ErrString(), ShouldEqual,
				`  [prbgtestCustom:8] (func.013:117)
    prbgtest content2
`)
			apdbg.ResetIOs()
			apdbg.pdbgTestInstance()
			So(apdbg.ErrString(), ShouldEqual,
				`  [*Pdbg.pdbgTestInstance:12] (func.013:123)
    pdbgTestInstance content3
`)
		})
		Convey("Test pdbg prints nothing if runtime.Caller fails", func() {
			mycaller = failCaller
			apdbg := NewPdbg(SetBuffers)
			apdbg.Pdbgf("test fail")
			So(apdbg.ErrString(), ShouldEqual, `  test fail
`)
			mycaller = runtime.Caller
		})
	})

	SkipConvey("Test pdbg excludes functions", t, func() {
		Convey("Test pdbg exclude with global instance", func() {
			SetBuffers(nil)
			pdbg.SetExcludes([]string{"globalNo"})
			globalPdbgExcludeTest()
			So(ErrString(), ShouldEqual,
				`  [globalPdbgExcludeTest:16] (func.016:143)
    calling no
      [globalCNo:26] (globalPdbgExcludeTest:17) (func.016:143)
        gcalled2
`)
		})
		Convey("Test pdbg exclude with custom instance", func() {
			apdbg := NewPdbg(SetBuffers, OptExcludes([]string{"customNo"}))
			customPdbgExcludeTest(apdbg)
			So(apdbg.ErrString(), ShouldEqual,
				`  [customPdbgExcludeTest:30] (func.017:153)
    calling cno
      [customCNo:40] (customPdbgExcludeTest:31) (func.017:153)
        ccalled2
`)
		})
	})

	SkipConvey("Test pdbg skips functions", t, func() {
		Convey("Test pdbg skip with global instance", func() {
			SetBuffers(nil)
			pdbg.SetSkips([]string{"globalNo"})
			globalPdbgExcludeTest()
			So(ErrString(), ShouldEqual,
				`  [globalPdbgExcludeTest:16] (func.019:167)
    calling no
      [globalCNo:26] (globalPdbgExcludeTest:17) (func.019:167)
        gcalled2
`)
		})
	})
}

func failCaller(skip int) (pc uintptr, file string, line int, ok bool) {
	return 0, "fail", skip, false
}
<|endoftext|>


<commit_before>package godbg

import (
	"fmt"
	"os"
	"runtime"
	"testing"
	. "github.com/smartystreets/goconvey/convey"
)

func TestProject(t *testing.T) {
	SkipConvey("Test buffers", t, func() {

		Convey("By Default, equals to std", func() {
			So(Out(), ShouldEqual, os.Stdout)
			So(Err(), ShouldEqual, os.Stderr)
		})
		Convey("When set to buffer, no longer equals to std", func() {
			SetBuffers(nil)
			So(Out(), ShouldNotEqual, os.Stdout)
			So(Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance buffer equals to std", func() {
			apdbg := NewPdbg()
			So(apdbg.Out(), ShouldEqual, os.Stdout)
			So(apdbg.Err(), ShouldEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance set to buffer writes no longer equals to std", func() {
			apdbg := NewPdbg(SetBuffers)
			So(apdbg.Out(), ShouldNotEqual, os.Stdout)
			So(apdbg.Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("Test custom buffer on global pdbg", func() {
			pdbg.bout = nil
			pdbg.sout = nil
			pdbg.berr = nil
			pdbg.serr = nil
			fmt.Fprintln(Out(), "test0 content0")
			So(OutString(), ShouldEqual, ``)
			fmt.Fprintln(Err(), "err0 content0")
			So(ErrString(), ShouldEqual, ``)
			SetBuffers(nil)
			fmt.Fprintln(Out(), "test content")
			fmt.Fprintln(Err(), "err1 cerr")
			fmt.Fprintln(Err(), "err2 cerr2")
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test content
test2 content2`)
			So(ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})

		Convey("Test custom buffer reset on global pdbg", func() {
			SetBuffers(nil)
			fmt.Fprint(Out(), "test content")
			So(OutString(), ShouldEqual, `test content`)
			fmt.Fprint(Err(), "err1 cerr")
			So(ErrString(), ShouldEqual, `err1 cerr`)
			ResetIOs()
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(Err(), "err2 cerr2")
			So(ErrString(), ShouldEqual, `err2 cerr2`)
		})

		Convey("Test custom buffer on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprintln(apdbg.Out(), "test content")
			fmt.Fprintln(apdbg.Err(), "err1 cerr")
			fmt.Fprintln(apdbg.Err(), "err2 cerr2")
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test content
test2 content2`)
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})
		Convey("Test custom buffer reset on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprint(apdbg.Out(), "test content")
			So(apdbg.OutString(), ShouldEqual, `test content`)
			fmt.Fprint(apdbg.Err(), "err1 cerr")
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr`)
			apdbg.ResetIOs()
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(apdbg.Err(), "err2 cerr2")
			So(apdbg.ErrString(), ShouldEqual, `err2 cerr2`)
		})
	})

	SkipConvey("Test pdbg print functions", t, func() {
		Convey("Test pdbg print with global instance", func() {
			SetBuffers(nil)
			Pdbgf("test")
			So(ErrString(), ShouldEqual,
				`[func.012:96]
  test
`)
			ResetIOs()
			prbgtest()
			So(ErrString(), ShouldEqual,
				` [prbgtest:4] (func.012:102)
   prbgtest content
`)
		})

		Convey("Test pdbg print with custom instance", func() {
			apdbg := NewPdbg(SetBuffers)
			apdbg.Pdbgf("test2")
			So(apdbg.ErrString(), ShouldEqual,
				`[func.013:111]
  test2
`)
			apdbg.ResetIOs()
			prbgtestCustom(apdbg)
			So(apdbg.ErrString(), ShouldEqual,
				` [prbgtestCustom:8] (func.013:117)
   prbgtest content2
`)
			apdbg.ResetIOs()
			apdbg.pdbgTestInstance()
			So(apdbg.ErrString(), ShouldEqual,
				` [*Pdbg.pdbgTestInstance:12] (func.013:123)
   pdbgTestInstance content3
`)
		})
		Convey("Test pdbg prints nothing if runtime.Caller fails", func() {
			mycaller = failCaller
			apdbg := NewPdbg(SetBuffers)
			apdbg.Pdbgf("test fail")
			So(apdbg.ErrString(), ShouldEqual, `  test fail
`)
			mycaller = runtime.Caller
		})
	})

	Convey("Test pdbg excludes functions", t, func() {
		Convey("Test pdbg exclude with global instance", func() {
			SetBuffers(nil)
			pdbg.SetExcludes([]string{"globalNo"})
			globalPdbgExcludeTest()
			So(ErrString(), ShouldEqual,
				`  [globalPdbgExcludeTest:16] (func.016:143)
    calling no
      [globalCNo:26] (globalPdbgExcludeTest:17) (func.016:143)
        gcalled2
`)
		})
		SkipConvey("Test pdbg exclude with custom instance", func() {
			apdbg := NewPdbg(SetBuffers, OptExcludes([]string{"customNo"}))
			customPdbgExcludeTest(apdbg)
			So(apdbg.ErrString(), ShouldEqual,
				`  [customPdbgExcludeTest:30] (func.017:153)
    calling cno
      [customCNo:40] (customPdbgExcludeTest:31) (func.017:153)
        ccalled2
`)
		})
	})

	SkipConvey("Test pdbg skips functions", t, func() {
		Convey("Test pdbg skip with global instance", func() {
			SetBuffers(nil)
			pdbg.SetSkips([]string{"globalNo"})
			globalPdbgExcludeTest()
			So(ErrString(), ShouldEqual,
				`  [globalPdbgExcludeTest:16] (func.019:167)
    calling no
    [globalPdbgExcludeTest:17] (func.019:167)
      gcalled1
      [globalCNo:26] (globalPdbgExcludeTest:17) (func.019:167)
        gcalled2
`)
		})
	})
}

func failCaller(skip int) (pc uintptr, file string, line int, ok bool) {
	return 0, "fail", skip, false
}
<commit_msg>Focus on second global exclude test<commit_after>package godbg

import (
	"fmt"
	"os"
	"runtime"
	"testing"
	. "github.com/smartystreets/goconvey/convey"
)

func TestProject(t *testing.T) {
	SkipConvey("Test buffers", t, func() {

		Convey("By Default, equals to std", func() {
			So(Out(), ShouldEqual, os.Stdout)
			So(Err(), ShouldEqual, os.Stderr)
		})
		Convey("When set to buffer, no longer equals to std", func() {
			SetBuffers(nil)
			So(Out(), ShouldNotEqual, os.Stdout)
			So(Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance buffer equals to std", func() {
			apdbg := NewPdbg()
			So(apdbg.Out(), ShouldEqual, os.Stdout)
			So(apdbg.Err(), ShouldEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance set to buffer writes no longer equals to std", func() {
			apdbg := NewPdbg(SetBuffers)
			So(apdbg.Out(), ShouldNotEqual, os.Stdout)
			So(apdbg.Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("Test custom buffer on global pdbg", func() {
			pdbg.bout = nil
			pdbg.sout = nil
			pdbg.berr = nil
			pdbg.serr = nil
			fmt.Fprintln(Out(), "test0 content0")
			So(OutString(), ShouldEqual, ``)
			fmt.Fprintln(Err(), "err0 content0")
			So(ErrString(), ShouldEqual, ``)
			SetBuffers(nil)
			fmt.Fprintln(Out(), "test content")
			fmt.Fprintln(Err(), "err1 cerr")
			fmt.Fprintln(Err(), "err2 cerr2")
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test content
test2 content2`)
			So(ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})

		Convey("Test custom buffer reset on global pdbg", func() {
			SetBuffers(nil)
			fmt.Fprint(Out(), "test content")
			So(OutString(), ShouldEqual, `test content`)
			fmt.Fprint(Err(), "err1 cerr")
			So(ErrString(), ShouldEqual, `err1 cerr`)
			ResetIOs()
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(Err(), "err2 cerr2")
			So(ErrString(), ShouldEqual, `err2 cerr2`)
		})

		Convey("Test custom buffer on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprintln(apdbg.Out(), "test content")
			fmt.Fprintln(apdbg.Err(), "err1 cerr")
			fmt.Fprintln(apdbg.Err(), "err2 cerr2")
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test content
test2 content2`)
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})
		Convey("Test custom buffer reset on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprint(apdbg.Out(), "test content")
			So(apdbg.OutString(), ShouldEqual, `test content`)
			fmt.Fprint(apdbg.Err(), "err1 cerr")
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr`)
			apdbg.ResetIOs()
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(apdbg.Err(), "err2 cerr2")
			So(apdbg.ErrString(), ShouldEqual, `err2 cerr2`)
		})
	})

	SkipConvey("Test pdbg print functions", t, func() {
		Convey("Test pdbg print with global instance", func() {
			SetBuffers(nil)
			Pdbgf("test")
			So(ErrString(), ShouldEqual,
				`[func.012:96]
  test
`)
			ResetIOs()
			prbgtest()
			So(ErrString(), ShouldEqual,
				` [prbgtest:4] (func.012:102)
   prbgtest content
`)
		})

		Convey("Test pdbg print with custom instance", func() {
			apdbg := NewPdbg(SetBuffers)
			apdbg.Pdbgf("test2")
			So(apdbg.ErrString(), ShouldEqual,
				`[func.013:111]
  test2
`)
			apdbg.ResetIOs()
			prbgtestCustom(apdbg)
			So(apdbg.ErrString(), ShouldEqual,
				` [prbgtestCustom:8] (func.013:117)
   prbgtest content2
`)
			apdbg.ResetIOs()
			apdbg.pdbgTestInstance()
			So(apdbg.ErrString(), ShouldEqual,
				` [*Pdbg.pdbgTestInstance:12] (func.013:123)
   pdbgTestInstance content3
`)
		})
		Convey("Test pdbg prints nothing if runtime.Caller fails", func() {
			mycaller = failCaller
			apdbg := NewPdbg(SetBuffers)
			apdbg.Pdbgf("test fail")
			So(apdbg.ErrString(), ShouldEqual, `  test fail
`)
			mycaller = runtime.Caller
		})
	})

	Convey("Test pdbg excludes functions", t, func() {
		SkipConvey("Test pdbg exclude with global instance", func() {
			SetBuffers(nil)
			pdbg.SetExcludes([]string{"globalNo"})
			globalPdbgExcludeTest()
			So(ErrString(), ShouldEqual,
				`  [globalPdbgExcludeTest:16] (func.016:143)
    calling no
      [globalCNo:26] (globalPdbgExcludeTest:17) (func.016:143)
        gcalled2
`)
		})
		Convey("Test pdbg exclude with custom instance", func() {
			apdbg := NewPdbg(SetBuffers, OptExcludes([]string{"customNo"}))
			customPdbgExcludeTest(apdbg)
			So(apdbg.ErrString(), ShouldEqual,
				`  [customPdbgExcludeTest:30] (func.017:153)
    calling cno
      [customCNo:40] (customPdbgExcludeTest:31) (func.017:153)
        ccalled2
`)
		})
	})

	SkipConvey("Test pdbg skips functions", t, func() {
		Convey("Test pdbg skip with global instance", func() {
			SetBuffers(nil)
			pdbg.SetSkips([]string{"globalNo"})
			globalPdbgExcludeTest()
			So(ErrString(), ShouldEqual,
				`  [globalPdbgExcludeTest:16] (func.019:167)
    calling no
    [globalPdbgExcludeTest:17] (func.019:167)
      gcalled1
      [globalCNo:26] (globalPdbgExcludeTest:17) (func.019:167)
        gcalled2
`)
		})
	})
}

func failCaller(skip int) (pc uintptr, file string, line int, ok bool) {
	return 0, "fail", skip, false
}
<|endoftext|>


<commit_before>package gode

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/dickeyxxx/golock"
	"github.com/franela/goreq"
	"github.com/mitchellh/ioprogress"
	"github.com/ulikunitz/xz"
)

var errInvalidSha = errors.New("Invalid SHA")

// IsSetup returns true if node is setup in RootPath
func IsSetup() (bool, error) {
	exists, err := fileExists(nodeBinPath)
	if !exists {
		return exists, err
	}
	return fileExists(npmBinPath)
}

// Setup downloads and sets up node in the RootPath directory
func Setup() error {
	golock.Lock(lockPath)
	defer golock.Unlock(lockPath)
	if setup, _ := IsSetup(); setup {
		return nil
	}
	if t == nil {
		return errors.New(`node does not offer a prebuilt binary for your OS.
You'll need to compile the tarball from nodejs.org and place the binary at ` + nodeBinPath)
	}
	if err := downloadFile(nodeBinPath, t.URL, t.Sha); err != nil {
		return err
	}
	if err := os.Chmod(nodeBinPath, 0755); err != nil {
		return err
	}
	if err := downloadNpm(); err != nil {
		return err
	}
	if err := os.MkdirAll(modulesDir, 0755); err != nil {
		return err
	}
	return clearOldNodeInstalls()
}

func downloadNpm() error {
	reader, getSha, err := downloadXZ(npmURL)
	if err != nil {
		return err
	}
	tmpDir := tmpDir()

	if err := extractTar(reader, tmpDir); err != nil {
		return err
	}
	if getSha() != npmSha {
		return errInvalidSha
	}
	os.RemoveAll(filepath.Join(npmBasePath))
	os.Rename(filepath.Join(tmpDir, "npm-"+NpmVersion), npmBasePath)
	return os.RemoveAll(tmpDir)
}

func downloadXZ(url string) (io.Reader, func() string, error) {
	req := goreq.Request{Uri: url}
	resp, err := req.Do()
	if err != nil {
		return nil, nil, err
	}
	size, _ := strconv.Atoi(resp.Header.Get("Content-Length"))
	progress := &ioprogress.Reader{
		Reader:   resp.Body,
		Size:     int64(size),
		DrawFunc: ioprogress.DrawTerminalf(os.Stderr, progressDrawFn),
	}
	getSha, reader := computeSha(progress)
	uncompressed, err := xz.NewReader(reader)
	return uncompressed, getSha, nil
}

func downloadFile(path, url, sha string) error {
	reader, getSha, err := downloadXZ(url)
	if err != nil {
		return err
	}
	tmp := filepath.Join(tmpDir(), "file")
	file, err := os.Create(tmp)
	if err != nil {
		return err
	}
	if _, err = io.Copy(file, reader); err != nil {
		return err
	}
	file.Close()
	if getSha() != sha {
		return errInvalidSha
	}
	if err = os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return err
	}
	os.Remove(path)
	err = os.Rename(tmp, path)
	if err != nil {
		return err
	}
	return os.RemoveAll(filepath.Dir(tmp))
}

func progressDrawFn(progress, total int64) string {
	return fmt.Sprintf("heroku-cli: Adding dependencies... %15s", ioprogress.DrawTextFormatBytes(progress, total))
}

func clearOldNodeInstalls() error {
	for _, name := range getDirsWithPrefix(rootPath, "node-") {
		if !strings.HasPrefix(name, "node-"+NodeVersion) {
			if err := os.RemoveAll(filepath.Join(rootPath, name)); err != nil {
				return err
			}
		}
	}
	for _, name := range getDirsWithPrefix(rootPath, "npm-") {
		if name != "npm-"+NpmVersion {
			if err := os.RemoveAll(filepath.Join(rootPath, name)); err != nil {
				return err
			}
		}
	}
	return nil
}
<commit_msg>catch http errors in gode<commit_after>package gode

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/dickeyxxx/golock"
	"github.com/franela/goreq"
	"github.com/mitchellh/ioprogress"
	"github.com/ulikunitz/xz"
)

var errInvalidSha = errors.New("Invalid SHA")

// IsSetup returns true if node is setup in RootPath
func IsSetup() (bool, error) {
	exists, err := fileExists(nodeBinPath)
	if !exists {
		return exists, err
	}
	return fileExists(npmBinPath)
}

// Setup downloads and sets up node in the RootPath directory
func Setup() error {
	golock.Lock(lockPath)
	defer golock.Unlock(lockPath)
	if setup, _ := IsSetup(); setup {
		return nil
	}
	if t == nil {
		return errors.New(`node does not offer a prebuilt binary for your OS.
You'll need to compile the tarball from nodejs.org and place the binary at ` + nodeBinPath)
	}
	if err := downloadFile(nodeBinPath, t.URL, t.Sha); err != nil {
		return err
	}
	if err := os.Chmod(nodeBinPath, 0755); err != nil {
		return err
	}
	if err := downloadNpm(); err != nil {
		return err
	}
	if err := os.MkdirAll(modulesDir, 0755); err != nil {
		return err
	}
	return clearOldNodeInstalls()
}

func downloadNpm() error {
	reader, getSha, err := downloadXZ(npmURL)
	if err != nil {
		return err
	}
	tmpDir := tmpDir()

	if err := extractTar(reader, tmpDir); err != nil {
		return err
	}
	if getSha() != npmSha {
		return errInvalidSha
	}
	os.RemoveAll(filepath.Join(npmBasePath))
	os.Rename(filepath.Join(tmpDir, "npm-"+NpmVersion), npmBasePath)
	return os.RemoveAll(tmpDir)
}

func downloadXZ(url string) (io.Reader, func() string, error) {
	req := goreq.Request{Uri: url}
	resp, err := req.Do()
	if err != nil {
		return nil, nil, err
	}
	if err := getHTTPError(resp); err != nil {
		return nil, nil, err
	}
	size, _ := strconv.Atoi(resp.Header.Get("Content-Length"))
	progress := &ioprogress.Reader{
		Reader:   resp.Body,
		Size:     int64(size),
		DrawFunc: ioprogress.DrawTerminalf(os.Stderr, progressDrawFn),
	}
	getSha, reader := computeSha(progress)
	uncompressed, err := xz.NewReader(reader)
	return uncompressed, getSha, nil
}

func downloadFile(path, url, sha string) error {
	reader, getSha, err := downloadXZ(url)
	if err != nil {
		return err
	}
	tmp := filepath.Join(tmpDir(), "file")
	file, err := os.Create(tmp)
	if err != nil {
		return err
	}
	if _, err = io.Copy(file, reader); err != nil {
		return err
	}
	file.Close()
	if getSha() != sha {
		return errInvalidSha
	}
	if err = os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return err
	}
	os.Remove(path)
	err = os.Rename(tmp, path)
	if err != nil {
		return err
	}
	return os.RemoveAll(filepath.Dir(tmp))
}

func progressDrawFn(progress, total int64) string {
	return fmt.Sprintf("heroku-cli: Adding dependencies... %15s", ioprogress.DrawTextFormatBytes(progress, total))
}

func clearOldNodeInstalls() error {
	for _, name := range getDirsWithPrefix(rootPath, "node-") {
		if !strings.HasPrefix(name, "node-"+NodeVersion) {
			if err := os.RemoveAll(filepath.Join(rootPath, name)); err != nil {
				return err
			}
		}
	}
	for _, name := range getDirsWithPrefix(rootPath, "npm-") {
		if name != "npm-"+NpmVersion {
			if err := os.RemoveAll(filepath.Join(rootPath, name)); err != nil {
				return err
			}
		}
	}
	return nil
}

func getHTTPError(resp *goreq.Response) error {
	if resp.StatusCode < 400 {
		return nil
	}
	var body string
	body = resp.Header.Get("Content-Type")
	return fmt.Errorf("%s: %s", resp.Status, body)
}
<|endoftext|>


<commit_before>package gostc

import (
	"net"
	"testing"
	"time"

	"github.com/cespare/a"
)

type TestServer struct {
	Addr     string
	Conn     *net.UDPConn
	Messages chan []byte
}

func NewTestServer() *TestServer {
	s := &TestServer{}
	u, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	conn, err := net.ListenUDP("udp", u)
	if err != nil {
		panic(err)
	}
	s.Conn = conn
	s.Addr = conn.LocalAddr().String()
	s.Messages = make(chan []byte)
	go func() {
		for {
			buf := make([]byte, 1000)
			n, _, err := s.Conn.ReadFromUDP(buf)
			if err != nil {
				return
			}
			s.Messages <- buf[:n]
		}
	}()
	return s
}

func (s *TestServer) Close() {
	s.Conn.Close()
}

func (s *TestServer) NextMessage() string {
	return string(<-s.Messages)
}

func MakeServerAndClient() (*TestServer, *Client) {
	s := NewTestServer()
	c, err := NewClient(s.Addr)
	if err != nil {
		panic(err)
	}
	return s, c
}

func MakeNonRandom(seq []float64) func() float64 {
	i := 0
	return func() float64 {
		v := seq[i]
		i++
		if i >= len(seq) {
			i = 0
		}
		return v
	}
}

func TestCount(t *testing.T) {
	server, client := MakeServerAndClient()
	defer server.Close()

	client.Count("foo", 3, 1)
	a.Assert(t, server.NextMessage(), a.Equals, "foo:3|c")

	client.Count("foo", 3, 0.5)
	a.Assert(t, server.NextMessage(), a.Equals, "foo:3|c@0.5")

	client.Count("blah", -123.456, 1)
	a.Assert(t, server.NextMessage(), a.Equals, "blah:-123.456|c")

	client.Inc("incme")
	a.Assert(t, server.NextMessage(), a.Equals, "incme:1|c")

	randFloat = MakeNonRandom([]float64{0.6, 0.4})
	client.CountProb("foo", 3, 0.5) // nothin
	client.CountProb("bar", 3, 0.5)
	a.Assert(t, server.NextMessage(), a.Equals, "bar:3|c@0.5")

	client.IncProb("foo", 0.5)
	client.IncProb("bar", 0.5)
	a.Assert(t, server.NextMessage(), a.Equals, "bar:1|c@0.5")
}

func TestTime(t *testing.T) {
	server, client := MakeServerAndClient()
	defer server.Close()

	client.Time("foo", 3*time.Second)
	a.Assert(t, server.NextMessage(), a.Equals, "foo:3000|ms")
}

func TestGauge(t *testing.T) {
	server, client := MakeServerAndClient()
	defer server.Close()

	client.Gauge("foo", 123.456)
	a.Assert(t, server.NextMessage(), a.Equals, "foo:123.456|g")
}

func TestSet(t *testing.T) {
	server, client := MakeServerAndClient()
	defer server.Close()

	client.Set("foo", []byte("hello"))
	a.Assert(t, server.NextMessage(), a.Equals, "foo:hello|s")
}

func TestBufferedMaxSize(t *testing.T) {
	s := NewTestServer()
	// 5 ms is hopefully enough time to be processed. Kind of a fragile test, but simple.
	c, err := NewBufferedClient(s.Addr, 100, 12, 5*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	for i := byte(0); i < 4; i++ {
		c.Set("a", []byte{'a' + i})
	}
	c.Close()
	a.Assert(t, s.NextMessage(), a.Equals, "a:a|s\na:b|s")
	a.Assert(t, s.NextMessage(), a.Equals, "a:c|s\na:d|s")
}

func TestBufferedMinFlush(t *testing.T) {
	s := NewTestServer()
	c, err := NewBufferedClient(s.Addr, 100, 100, 5*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	for i := byte(0); i < 4; i++ {
		c.Set("a", []byte{'a' + i})
		if i == 1 {
			time.Sleep(10 * time.Millisecond)
		}
	}
	c.Close()
	a.Assert(t, s.NextMessage(), a.Equals, "a:a|s\na:b|s")
	a.Assert(t, s.NextMessage(), a.Equals, "a:c|s\na:d|s")
}
<commit_msg>Switch assertion lib from a to asrt<commit_after>package gostc

import (
	"net"
	"testing"
	"time"

	"github.com/cespare/asrt"
)

type TestServer struct {
	Addr     string
	Conn     *net.UDPConn
	Messages chan []byte
}

func NewTestServer() *TestServer {
	s := &TestServer{}
	u, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	conn, err := net.ListenUDP("udp", u)
	if err != nil {
		panic(err)
	}
	s.Conn = conn
	s.Addr = conn.LocalAddr().String()
	s.Messages = make(chan []byte)
	go func() {
		for {
			buf := make([]byte, 1000)
			n, _, err := s.Conn.ReadFromUDP(buf)
			if err != nil {
				return
			}
			s.Messages <- buf[:n]
		}
	}()
	return s
}

func (s *TestServer) Close() {
	s.Conn.Close()
}

func (s *TestServer) NextMessage() string {
	return string(<-s.Messages)
}

func MakeServerAndClient() (*TestServer, *Client) {
	s := NewTestServer()
	c, err := NewClient(s.Addr)
	if err != nil {
		panic(err)
	}
	return s, c
}

func MakeNonRandom(seq []float64) func() float64 {
	i := 0
	return func() float64 {
		v := seq[i]
		i++
		if i >= len(seq) {
			i = 0
		}
		return v
	}
}

func TestCount(t *testing.T) {
	server, client := MakeServerAndClient()
	defer server.Close()

	client.Count("foo", 3, 1)
	asrt.Equal(t, server.NextMessage(), "foo:3|c")

	client.Count("foo", 3, 0.5)
	asrt.Equal(t, server.NextMessage(), "foo:3|c@0.5")

	client.Count("blah", -123.456, 1)
	asrt.Equal(t, server.NextMessage(), "blah:-123.456|c")

	client.Inc("incme")
	asrt.Equal(t, server.NextMessage(), "incme:1|c")

	randFloat = MakeNonRandom([]float64{0.6, 0.4})
	client.CountProb("foo", 3, 0.5) // nothin
	client.CountProb("bar", 3, 0.5)
	asrt.Equal(t, server.NextMessage(), "bar:3|c@0.5")

	client.IncProb("foo", 0.5)
	client.IncProb("bar", 0.5)
	asrt.Equal(t, server.NextMessage(), "bar:1|c@0.5")
}

func TestTime(t *testing.T) {
	server, client := MakeServerAndClient()
	defer server.Close()

	client.Time("foo", 3*time.Second)
	asrt.Equal(t, server.NextMessage(), "foo:3000|ms")
}

func TestGauge(t *testing.T) {
	server, client := MakeServerAndClient()
	defer server.Close()

	client.Gauge("foo", 123.456)
	asrt.Equal(t, server.NextMessage(), "foo:123.456|g")
}

func TestSet(t *testing.T) {
	server, client := MakeServerAndClient()
	defer server.Close()

	client.Set("foo", []byte("hello"))
	asrt.Equal(t, server.NextMessage(), "foo:hello|s")
}

func TestBufferedMaxSize(t *testing.T) {
	s := NewTestServer()
	// 5 ms is hopefully enough time to be processed. Kind of a fragile test, but simple.
	c, err := NewBufferedClient(s.Addr, 100, 12, 5*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	for i := byte(0); i < 4; i++ {
		c.Set("a", []byte{'a' + i})
	}
	c.Close()
	asrt.Equal(t, s.NextMessage(), "a:a|s\na:b|s")
	asrt.Equal(t, s.NextMessage(), "a:c|s\na:d|s")
}

func TestBufferedMinFlush(t *testing.T) {
	s := NewTestServer()
	c, err := NewBufferedClient(s.Addr, 100, 100, 5*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	for i := byte(0); i < 4; i++ {
		c.Set("a", []byte{'a' + i})
		if i == 1 {
			time.Sleep(10 * time.Millisecond)
		}
	}
	c.Close()
	asrt.Equal(t, s.NextMessage(), "a:a|s\na:b|s")
	asrt.Equal(t, s.NextMessage(), "a:c|s\na:d|s")
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	stdContext "context"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/kataras/iris"
)

func logger(app *iris.Application) *bytes.Buffer {
	buf := &bytes.Buffer{}

	app.Logger().Out = buf

	// disable the "Now running at...." in order to have a clean log of the error.
	// we could attach that on `Run` but better to keep things simple here.
	app.Configure(iris.WithoutStartupLog)
	return buf
}

func TestListenAddr(t *testing.T) {
	app := iris.New()
	// we keep the logger running as well but in a controlled way.
	log := logger(app)

	// close the server at 3-6 seconds
	go func() {
		time.Sleep(3 * time.Second)
		ctx, cancel := stdContext.WithTimeout(stdContext.TODO(), 3*time.Second)
		defer cancel()
		app.Shutdown(ctx)
	}()

	err := app.Run(iris.Addr(":9829"))
	// in this case the error should be logged and return as well.
	if err != iris.ErrServerClosed {
		t.Fatalf("expecting err to be `iris.ErrServerClosed` but got: %v", err)
	}

	// println(log.Bytes())
	// println(len(log.Bytes()))

	expected := fmt.Sprintln("\"" + iris.ErrServerClosed.Error() + "\" ")
	// println([]byte(expected))
	// println(len([]byte(expected)))

	got := log.String()
	got = strings.Split(got, "msg=")[1]
	if expected != got {
		t.Fatalf("expecting to log the:\n'%s'\ninstead of:\n'%s'", expected, got)
	}
}

func TestListenAddrWithoutServerErr(t *testing.T) {
	app := iris.New()
	// we keep the logger running as well but in a controlled way.
	log := logger(app)

	// close the server at 3-6 seconds
	go func() {
		time.Sleep(3 * time.Second)
		ctx, cancel := stdContext.WithTimeout(stdContext.TODO(), 3*time.Second)
		defer cancel()
		app.Shutdown(ctx)
	}()

	// we disable the ErrServerClosed, so the error should be nil when server is closed by `app.Shutdown`.

	// so in this case the iris/http.ErrServerClosed should be NOT logged and NOT return.
	err := app.Run(iris.Addr(":9827"), iris.WithoutServerError(iris.ErrServerClosed))
	if err != nil {
		t.Fatalf("expecting err to be nil but got: %v", err)
	}

	if got := log.String(); got != "" {
		t.Fatalf("expecting to log nothing but logged: '%s'", got)
	}
}
<commit_msg>fix the listen-addr test fails on unix :one:<commit_after>package main

import (
	"bytes"
	stdContext "context"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/kataras/iris"
)

func logger(app *iris.Application) *bytes.Buffer {
	buf := &bytes.Buffer{}

	app.Logger().Out = buf

	// disable the "Now running at...." in order to have a clean log of the error.
	// we could attach that on `Run` but better to keep things simple here.
	app.Configure(iris.WithoutStartupLog)
	return buf
}

func TestListenAddr(t *testing.T) {
	app := iris.New()
	// we keep the logger running as well but in a controlled way.
	log := logger(app)

	// close the server at 3-6 seconds
	go func() {
		time.Sleep(3 * time.Second)
		ctx, cancel := stdContext.WithTimeout(stdContext.TODO(), 3*time.Second)
		defer cancel()
		app.Shutdown(ctx)
	}()

	err := app.Run(iris.Addr(":9829"))
	// in this case the error should be logged and return as well.
	if err != iris.ErrServerClosed {
		t.Fatalf("expecting err to be `iris.ErrServerClosed` but got: %v", err)
	}

	// println(log.Bytes())
	// println(len(log.Bytes()))

	expected := fmt.Sprintln("\"" + iris.ErrServerClosed.Error() + "\" ")
	expected = strings.TrimSpace(expected)
	// println([]byte(expected))
	// println(len([]byte(expected)))

	got := log.String()
	got = strings.Split(got, "msg=")[1]
	got = strings.TrimSpace(got)
	if expected != got {
		t.Fatalf("expecting to log the:\n'%s'\ninstead of:\n'%s'", expected, got)
	}
}

func TestListenAddrWithoutServerErr(t *testing.T) {
	app := iris.New()
	// we keep the logger running as well but in a controlled way.
	log := logger(app)

	// close the server at 3-6 seconds
	go func() {
		time.Sleep(3 * time.Second)
		ctx, cancel := stdContext.WithTimeout(stdContext.TODO(), 3*time.Second)
		defer cancel()
		app.Shutdown(ctx)
	}()

	// we disable the ErrServerClosed, so the error should be nil when server is closed by `app.Shutdown`.

	// so in this case the iris/http.ErrServerClosed should be NOT logged and NOT return.
	err := app.Run(iris.Addr(":9827"), iris.WithoutServerError(iris.ErrServerClosed))
	if err != nil {
		t.Fatalf("expecting err to be nil but got: %v", err)
	}

	if got := log.String(); got != "" {
		t.Fatalf("expecting to log nothing but logged: '%s'", got)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Claudemiro Alves Feitosa Neto. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package ipe

import (
	log "github.com/golang/glog"
	"github.com/gorilla/websocket"
)

// An User Connection
type connection struct {
	SocketID string
	Socket   *websocket.Conn
}

// Create a new Subscriber
func newConnection(socketID string, s *websocket.Conn) *connection {
	log.Infof("Creating a new Subscriber %+v", socketID)

	return &connection{SocketID: socketID, Socket: s}
}

// Publish the message to websocket atached to this client
func (conn *connection) Publish(m interface{}) {
	go func() {
		if err := conn.Socket.WriteJSON(m); err != nil {
			log.Errorf("Error publishing message to connection %+v, %s", conn, err)
		}
	}()
}
<commit_msg>Added Created at to connection.<commit_after>// Copyright 2014 Claudemiro Alves Feitosa Neto. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package ipe

import (
	"time"

	log "github.com/golang/glog"
	"github.com/gorilla/websocket"
)

// An User Connection
type connection struct {
	SocketID  string
	Socket    *websocket.Conn
	CreatedAt time.Time
}

// Create a new Subscriber
func newConnection(socketID string, s *websocket.Conn) *connection {
	log.Infof("Creating a new Subscriber %+v", socketID)

	return &connection{SocketID: socketID, Socket: s, CreatedAt: time.Now()}
}

// Publish the message to websocket atached to this client
func (conn *connection) Publish(m interface{}) {
	go func() {
		if err := conn.Socket.WriteJSON(m); err != nil {
			log.Errorf("Error publishing message to connection %+v, %s", conn, err)
		}
	}()
}
<|endoftext|>


<commit_before>package nodepool

import (
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"time"

	"reflect"

	"github.com/rancher/rancher/pkg/ref"
	"github.com/rancher/rke/services"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
)

var (
	nameRegexp = regexp.MustCompile("^(.*?)([0-9]+)$")
)

type Controller struct {
	NodePoolController v3.NodePoolController
	NodePoolLister     v3.NodePoolLister
	NodePools          v3.NodePoolInterface
	NodeLister         v3.NodeLister
	Nodes              v3.NodeInterface
}

func Register(management *config.ManagementContext) {
	p := &Controller{
		NodePoolController: management.Management.NodePools("").Controller(),
		NodePoolLister:     management.Management.NodePools("").Controller().Lister(),
		NodePools:          management.Management.NodePools(""),
		NodeLister:         management.Management.Nodes("").Controller().Lister(),
		Nodes:              management.Management.Nodes(""),
	}

	// Add handlers
	p.NodePools.AddLifecycle("nodepool-provisioner", p)
	management.Management.Nodes("").AddHandler("nodepool-provisioner", p.machineChanged)
}

func (c *Controller) Create(nodePool *v3.NodePool) (*v3.NodePool, error) {
	return nodePool, nil
}

func (c *Controller) Updated(nodePool *v3.NodePool) (*v3.NodePool, error) {
	obj, err := v3.NodePoolConditionUpdated.Do(nodePool, func() (runtime.Object, error) {
		return nodePool, c.createNodes(nodePool)
	})
	return obj.(*v3.NodePool), err
}

func (c *Controller) Remove(nodePool *v3.NodePool) (*v3.NodePool, error) {
	logrus.Infof("Deleting nodePool [%s]", nodePool.Name)

	allNodes, err := c.nodes(nodePool, false)
	if err != nil {
		return nodePool, err
	}

	for _, node := range allNodes {
		_, nodePoolName := ref.Parse(node.Spec.NodePoolName)
		if nodePoolName != nodePool.Name {
			continue
		}

		err := c.deleteNode(node, time.Duration(0))
		if err != nil {
			return nil, err
		}
	}

	return nodePool, nil
}

func (c *Controller) machineChanged(key string, machine *v3.Node) error {
	if machine == nil {
		nps, err := c.NodePoolLister.List("", labels.Everything())
		if err != nil {
			return err
		}
		for _, np := range nps {
			c.NodePoolController.Enqueue(np.Namespace, np.Name)
		}
	} else if machine.Spec.NodePoolName != "" {
		ns, name := ref.Parse(machine.Spec.NodePoolName)
		c.NodePoolController.Enqueue(ns, name)
	}

	return nil
}

func (c *Controller) createNode(name string, nodePool *v3.NodePool, simulate bool) (*v3.Node, error) {
	newNode := &v3.Node{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "m-",
			Namespace:    nodePool.Namespace,
			Labels:       nodePool.Labels,
			Annotations:  nodePool.Annotations,
		},
		Spec: v3.NodeSpec{
			Etcd:              nodePool.Spec.Etcd,
			ControlPlane:      nodePool.Spec.ControlPlane,
			Worker:            nodePool.Spec.Worker,
			NodeTemplateName:  nodePool.Spec.NodeTemplateName,
			NodePoolName:      ref.Ref(nodePool),
			RequestedHostname: name,
		},
	}

	if simulate {
		return newNode, nil
	}

	return c.Nodes.Create(newNode)
}

func (c *Controller) deleteNode(node *v3.Node, duration time.Duration) error {
	f := metav1.DeletePropagationBackground

	if duration > time.Duration(0) {
		go func() {
			time.Sleep(duration)
			c.Nodes.DeleteNamespaced(node.Namespace, node.Name, &metav1.DeleteOptions{
				PropagationPolicy: &f,
			})
		}()
		return nil
	}

	return c.Nodes.DeleteNamespaced(node.Namespace, node.Name, &metav1.DeleteOptions{
		PropagationPolicy: &f,
	})
}

func (c *Controller) createNodes(nodePool *v3.NodePool) error {
	changed, err := c.createOrCheckNodes(nodePool, true)
	if err != nil {
		return err
	}

	if changed {
		_, err = c.createOrCheckNodes(nodePool, false)
	}

	return err
}

func parsePrefix(fullPrefix string) (prefix string, minLength, start int) {
	m := nameRegexp.FindStringSubmatch(fullPrefix)
	if len(m) == 0 {
		return fullPrefix, 1, 1
	}
	prefix = m[1]
	start, _ = strconv.Atoi(m[2])
	return prefix, len(m[2]), start
}

func (c *Controller) nodes(nodePool *v3.NodePool, simulate bool) ([]*v3.Node, error) {
	if simulate {
		return c.NodeLister.List(nodePool.Namespace, labels.Everything())
	}

	nodeList, err := c.Nodes.List(metav1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var nodes []*v3.Node
	for i := range nodeList.Items {
		if nodeList.Items[i].Namespace == nodePool.Namespace {
			nodes = append(nodes, &nodeList.Items[i])
		}
	}

	return nodes, nil
}

func (c *Controller) createOrCheckNodes(nodePool *v3.NodePool, simulate bool) (bool, error) {
	var (
		err     error
		byName  = map[string]*v3.Node{}
		changed = false
		nodes   []*v3.Node
	)

	allNodes, err := c.nodes(nodePool, simulate)
	if err != nil {
		return false, err
	}

	for _, node := range allNodes {
		byName[node.Spec.RequestedHostname] = node

		_, nodePoolName := ref.Parse(node.Spec.NodePoolName)
		if nodePoolName != nodePool.Name {
			continue
		}

		if v3.NodeConditionProvisioned.IsFalse(node) || v3.NodeConditionInitialized.IsFalse(node) || v3.NodeConditionConfigSaved.IsFalse(node) {
			changed = true
			if !simulate {
				c.deleteNode(node, 2*time.Minute)
			}
		}

		nodes = append(nodes, node)
	}

	quantity := nodePool.Spec.Quantity
	if quantity < 0 {
		quantity = 0
	}

	prefix, minLength, start := parsePrefix(nodePool.Spec.HostnamePrefix)

	for i := start; len(nodes) < quantity; i++ {
		ia := strconv.Itoa(i)
		name := prefix + ia
		if len(ia) < minLength {
			name = fmt.Sprintf("%s%0"+strconv.Itoa(minLength)+"d", prefix, i)
		}

		if byName[name] != nil {
			continue
		}

		changed = true
		newNode, err := c.createNode(name, nodePool, simulate)
		if err != nil {
			return false, err
		}

		byName[newNode.Spec.RequestedHostname] = newNode
		nodes = append(nodes, newNode)
	}

	for len(nodes) > quantity {
		sort.Slice(nodes, func(i, j int) bool {
			return nodes[i].Spec.RequestedHostname < nodes[j].Spec.RequestedHostname
		})

		toDelete := nodes[len(nodes)-1]

		changed = true
		if !simulate {
			c.deleteNode(toDelete, 0)
		}

		nodes = nodes[:len(nodes)-1]
		delete(byName, toDelete.Spec.RequestedHostname)
	}

	for _, n := range nodes {
		if needRoleUpdate(n, nodePool) {
			changed = true
			_, err := c.updateNodeRoles(n, nodePool, simulate)
			if err != nil {
				return false, err
			}
		}
	}

	return changed, nil
}

func needRoleUpdate(node *v3.Node, nodePool *v3.NodePool) bool {
	if node.Status.NodeConfig == nil {
		return false
	}
	if len(node.Status.NodeConfig.Role) == 0 && !nodePool.Spec.Worker {
		return true
	}

	nodeRolesMap := map[string]bool{}
	for _, role := range node.Status.NodeConfig.Role {
		switch r := role; r {
		case services.ETCDRole:
			nodeRolesMap[services.ETCDRole] = true
		case services.ControlRole:
			nodeRolesMap[services.ControlRole] = true
		case services.WorkerRole:
			nodeRolesMap[services.WorkerRole] = true
		}
	}

	poolRolesMap := map[string]bool{}
	nodeRolesMap[services.ETCDRole] = nodePool.Spec.Etcd
	nodeRolesMap[services.ControlRole] = nodePool.Spec.ControlPlane
	nodeRolesMap[services.WorkerRole] = nodePool.Spec.Worker
	return !reflect.DeepEqual(nodeRolesMap, poolRolesMap)
}

func (c *Controller) updateNodeRoles(existing *v3.Node, nodePool *v3.NodePool, simulate bool) (*v3.Node, error) {
	toUpdate := existing.DeepCopy()
	var newRoles []string

	if nodePool.Spec.ControlPlane {
		newRoles = append(newRoles, "controlplane")
	}
	if nodePool.Spec.Etcd {
		newRoles = append(newRoles, "etcd")
	}
	if nodePool.Spec.Worker {
		newRoles = append(newRoles, "worker")
	}

	toUpdate.Status.NodeConfig.Role = newRoles
	if simulate {
		return toUpdate, nil
	}

	return c.Nodes.Update(toUpdate)
}
<commit_msg>Compare node roles correctly<commit_after>package nodepool

import (
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"time"

	"reflect"

	"github.com/rancher/rancher/pkg/ref"
	"github.com/rancher/rke/services"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
)

var (
	nameRegexp = regexp.MustCompile("^(.*?)([0-9]+)$")
)

type Controller struct {
	NodePoolController v3.NodePoolController
	NodePoolLister     v3.NodePoolLister
	NodePools          v3.NodePoolInterface
	NodeLister         v3.NodeLister
	Nodes              v3.NodeInterface
}

func Register(management *config.ManagementContext) {
	p := &Controller{
		NodePoolController: management.Management.NodePools("").Controller(),
		NodePoolLister:     management.Management.NodePools("").Controller().Lister(),
		NodePools:          management.Management.NodePools(""),
		NodeLister:         management.Management.Nodes("").Controller().Lister(),
		Nodes:              management.Management.Nodes(""),
	}

	// Add handlers
	p.NodePools.AddLifecycle("nodepool-provisioner", p)
	management.Management.Nodes("").AddHandler("nodepool-provisioner", p.machineChanged)
}

func (c *Controller) Create(nodePool *v3.NodePool) (*v3.NodePool, error) {
	return nodePool, nil
}

func (c *Controller) Updated(nodePool *v3.NodePool) (*v3.NodePool, error) {
	obj, err := v3.NodePoolConditionUpdated.Do(nodePool, func() (runtime.Object, error) {
		return nodePool, c.createNodes(nodePool)
	})
	return obj.(*v3.NodePool), err
}

func (c *Controller) Remove(nodePool *v3.NodePool) (*v3.NodePool, error) {
	logrus.Infof("Deleting nodePool [%s]", nodePool.Name)

	allNodes, err := c.nodes(nodePool, false)
	if err != nil {
		return nodePool, err
	}

	for _, node := range allNodes {
		_, nodePoolName := ref.Parse(node.Spec.NodePoolName)
		if nodePoolName != nodePool.Name {
			continue
		}

		err := c.deleteNode(node, time.Duration(0))
		if err != nil {
			return nil, err
		}
	}

	return nodePool, nil
}

func (c *Controller) machineChanged(key string, machine *v3.Node) error {
	if machine == nil {
		nps, err := c.NodePoolLister.List("", labels.Everything())
		if err != nil {
			return err
		}
		for _, np := range nps {
			c.NodePoolController.Enqueue(np.Namespace, np.Name)
		}
	} else if machine.Spec.NodePoolName != "" {
		ns, name := ref.Parse(machine.Spec.NodePoolName)
		c.NodePoolController.Enqueue(ns, name)
	}

	return nil
}

func (c *Controller) createNode(name string, nodePool *v3.NodePool, simulate bool) (*v3.Node, error) {
	newNode := &v3.Node{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "m-",
			Namespace:    nodePool.Namespace,
			Labels:       nodePool.Labels,
			Annotations:  nodePool.Annotations,
		},
		Spec: v3.NodeSpec{
			Etcd:              nodePool.Spec.Etcd,
			ControlPlane:      nodePool.Spec.ControlPlane,
			Worker:            nodePool.Spec.Worker,
			NodeTemplateName:  nodePool.Spec.NodeTemplateName,
			NodePoolName:      ref.Ref(nodePool),
			RequestedHostname: name,
		},
	}

	if simulate {
		return newNode, nil
	}

	return c.Nodes.Create(newNode)
}

func (c *Controller) deleteNode(node *v3.Node, duration time.Duration) error {
	f := metav1.DeletePropagationBackground

	if duration > time.Duration(0) {
		go func() {
			time.Sleep(duration)
			c.Nodes.DeleteNamespaced(node.Namespace, node.Name, &metav1.DeleteOptions{
				PropagationPolicy: &f,
			})
		}()
		return nil
	}

	return c.Nodes.DeleteNamespaced(node.Namespace, node.Name, &metav1.DeleteOptions{
		PropagationPolicy: &f,
	})
}

func (c *Controller) createNodes(nodePool *v3.NodePool) error {
	changed, err := c.createOrCheckNodes(nodePool, true)
	if err != nil {
		return err
	}

	if changed {
		_, err = c.createOrCheckNodes(nodePool, false)
	}

	return err
}

func parsePrefix(fullPrefix string) (prefix string, minLength, start int) {
	m := nameRegexp.FindStringSubmatch(fullPrefix)
	if len(m) == 0 {
		return fullPrefix, 1, 1
	}
	prefix = m[1]
	start, _ = strconv.Atoi(m[2])
	return prefix, len(m[2]), start
}

func (c *Controller) nodes(nodePool *v3.NodePool, simulate bool) ([]*v3.Node, error) {
	if simulate {
		return c.NodeLister.List(nodePool.Namespace, labels.Everything())
	}

	nodeList, err := c.Nodes.List(metav1.ListOptions{})
	if err != nil {
		return nil, err
	}

	var nodes []*v3.Node
	for i := range nodeList.Items {
		if nodeList.Items[i].Namespace == nodePool.Namespace {
			nodes = append(nodes, &nodeList.Items[i])
		}
	}

	return nodes, nil
}

func (c *Controller) createOrCheckNodes(nodePool *v3.NodePool, simulate bool) (bool, error) {
	var (
		err     error
		byName  = map[string]*v3.Node{}
		changed = false
		nodes   []*v3.Node
	)

	allNodes, err := c.nodes(nodePool, simulate)
	if err != nil {
		return false, err
	}

	for _, node := range allNodes {
		byName[node.Spec.RequestedHostname] = node

		_, nodePoolName := ref.Parse(node.Spec.NodePoolName)
		if nodePoolName != nodePool.Name {
			continue
		}

		if v3.NodeConditionProvisioned.IsFalse(node) || v3.NodeConditionInitialized.IsFalse(node) || v3.NodeConditionConfigSaved.IsFalse(node) {
			changed = true
			if !simulate {
				c.deleteNode(node, 2*time.Minute)
			}
		}

		nodes = append(nodes, node)
	}

	quantity := nodePool.Spec.Quantity
	if quantity < 0 {
		quantity = 0
	}

	prefix, minLength, start := parsePrefix(nodePool.Spec.HostnamePrefix)

	for i := start; len(nodes) < quantity; i++ {
		ia := strconv.Itoa(i)
		name := prefix + ia
		if len(ia) < minLength {
			name = fmt.Sprintf("%s%0"+strconv.Itoa(minLength)+"d", prefix, i)
		}

		if byName[name] != nil {
			continue
		}

		changed = true
		newNode, err := c.createNode(name, nodePool, simulate)
		if err != nil {
			return false, err
		}

		byName[newNode.Spec.RequestedHostname] = newNode
		nodes = append(nodes, newNode)
	}

	for len(nodes) > quantity {
		sort.Slice(nodes, func(i, j int) bool {
			return nodes[i].Spec.RequestedHostname < nodes[j].Spec.RequestedHostname
		})

		toDelete := nodes[len(nodes)-1]

		changed = true
		if !simulate {
			c.deleteNode(toDelete, 0)
		}

		nodes = nodes[:len(nodes)-1]
		delete(byName, toDelete.Spec.RequestedHostname)
	}

	for _, n := range nodes {
		if needRoleUpdate(n, nodePool) {
			changed = true
			_, err := c.updateNodeRoles(n, nodePool, simulate)
			if err != nil {
				return false, err
			}
		}
	}

	return changed, nil
}

func needRoleUpdate(node *v3.Node, nodePool *v3.NodePool) bool {
	if node.Status.NodeConfig == nil {
		return false
	}
	if len(node.Status.NodeConfig.Role) == 0 && !nodePool.Spec.Worker {
		return true
	}

	nodeRolesMap := map[string]bool{}
	for _, role := range node.Status.NodeConfig.Role {
		switch r := role; r {
		case services.ETCDRole:
			nodeRolesMap[services.ETCDRole] = true
		case services.ControlRole:
			nodeRolesMap[services.ControlRole] = true
		case services.WorkerRole:
			nodeRolesMap[services.WorkerRole] = true
		}
	}

	poolRolesMap := map[string]bool{}
	poolRolesMap[services.ETCDRole] = nodePool.Spec.Etcd
	poolRolesMap[services.ControlRole] = nodePool.Spec.ControlPlane
	poolRolesMap[services.WorkerRole] = nodePool.Spec.Worker
	return !reflect.DeepEqual(nodeRolesMap, poolRolesMap)
}

func (c *Controller) updateNodeRoles(existing *v3.Node, nodePool *v3.NodePool, simulate bool) (*v3.Node, error) {
	toUpdate := existing.DeepCopy()
	var newRoles []string

	if nodePool.Spec.ControlPlane {
		newRoles = append(newRoles, "controlplane")
	}
	if nodePool.Spec.Etcd {
		newRoles = append(newRoles, "etcd")
	}
	if nodePool.Spec.Worker {
		newRoles = append(newRoles, "worker")
	}

	toUpdate.Status.NodeConfig.Role = newRoles
	if simulate {
		return toUpdate, nil
	}

	return c.Nodes.Update(toUpdate)
}
<|endoftext|>


<commit_before>/*-
 * Copyright 2014 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/codegangsta/cli"
	"github.com/square/go-jose"
)

func main() {
	app := cli.NewApp()
	app.Name = "jose-util"
	app.Usage = "command-line utility to deal with JOSE objects"
	app.Version = "0.0.2"
	app.Author = ""
	app.Email = ""

	app.Commands = []cli.Command{
		{
			Name:  "encrypt",
			Usage: "encrypt a plaintext",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "key, k",
					Usage: "Path to key file (PEM/DER)",
				},
				cli.StringFlag{
					Name:  "input, in",
					Usage: "Path to input file (stdin if missing)",
				},
				cli.StringFlag{
					Name:  "output, out",
					Usage: "Path to output file (stdout if missing)",
				},
				cli.StringFlag{
					Name:  "algorithm, alg",
					Usage: "Key management algorithm (e.g. RSA-OAEP)",
				},
				cli.StringFlag{
					Name:  "encryption, enc",
					Usage: "Content encryption algorithm (e.g. A128GCM)",
				},
				cli.BoolFlag{
					Name:  "full, f",
					Usage: "Use full serialization format (instead of compact)",
				},
			},
			Action: func(c *cli.Context) {
				keyBytes, err := ioutil.ReadFile(requiredFlag(c, "key"))
				exitOnError(err, "unable to read key file")

				pub, err := jose.LoadPublicKey(keyBytes)
				exitOnError(err, "unable to read public key")

				alg := jose.KeyAlgorithm(requiredFlag(c, "alg"))
				enc := jose.ContentEncryption(requiredFlag(c, "enc"))

				crypter, err := jose.NewEncrypter(alg, enc, pub)
				exitOnError(err, "unable to instantiate encrypter")

				obj, err := crypter.Encrypt(readInput(c.String("input")))
				exitOnError(err, "unable to encrypt")

				var msg string
				if c.Bool("full") {
					msg = obj.FullSerialize()
				} else {
					msg, err = obj.CompactSerialize()
					exitOnError(err, "unable to serialize message")
				}

				writeOutput(c.String("output"), []byte(msg))
			},
		},
		{
			Name:  "decrypt",
			Usage: "decrypt a plaintext",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "key, k",
					Usage: "Path to key file (PEM/DER)",
				},
				cli.StringFlag{
					Name:  "input, in",
					Usage: "Path to input file (stdin if missing)",
				},
				cli.StringFlag{
					Name:  "output, out",
					Usage: "Path to output file (stdout if missing)",
				},
			},
			Action: func(c *cli.Context) {
				keyBytes, err := ioutil.ReadFile(requiredFlag(c, "key"))
				exitOnError(err, "unable to read private key")

				priv, err := jose.LoadPrivateKey(keyBytes)
				exitOnError(err, "unable to read private key")

				obj, err := jose.ParseEncrypted(string(readInput(c.String("input"))))
				exitOnError(err, "unable to parse message")

				plaintext, err := obj.Decrypt(priv)
				exitOnError(err, "unable to decrypt message")

				writeOutput(c.String("output"), plaintext)
			},
		},
		{
			Name:  "dump",
			Usage: "parse & dump message in full serialization format",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "input, in",
					Usage: "Path to input file (stdin if missing)",
				},
				cli.StringFlag{
					Name:  "output, out",
					Usage: "Path to output file (stdout if missing)",
				},
				cli.StringFlag{
					Name:  "format, f",
					Usage: "Message format (JWE/JWS, defaults to JWE)",
				},
			},
			Action: func(c *cli.Context) {
				input := string(readInput(c.String("input")))

				var serialized string
				var err error
				switch c.String("format") {
				case "", "JWE":
					var jwe *jose.JsonWebEncryption
					jwe, err = jose.ParseEncrypted(input)
					if err == nil {
						serialized = jwe.FullSerialize()
					}
				case "JWS":
					var jws *jose.JsonWebSignature
					jws, err = jose.ParseSigned(input)
					if err == nil {
						serialized = jws.FullSerialize()
					}
				}

				exitOnError(err, "unable to parse message")

				var raw map[string]interface{}
				err = json.Unmarshal([]byte(serialized), &raw)
				exitOnError(err, "unable to parse message")

				output, err := json.MarshalIndent(&raw, "", "\t")
				exitOnError(err, "unable to serialize message")

				writeOutput(c.String("output"), output)
			},
		},
		{
			Name:  "sign",
			Usage: "sign a text",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "key, k",
					Usage: "Path to key file (PEM/DER)",
				},
				cli.StringFlag{
					Name:  "input, in",
					Usage: "Path to input file (stdin if missing)",
				},
				cli.StringFlag{
					Name:  "output, out",
					Usage: "Path to output file (stdout if missing)",
				},
				cli.BoolFlag{
					Name:  "full, f",
					Usage: "Use full serialization format (instead of compact)",
				},
			},
			Action: func(c *cli.Context) {
				keyBytes, err := ioutil.ReadFile(requiredFlag(c, "key"))
				exitOnError(err, "unable to read key file")

				signingKey, err := jose.LoadPrivateKey(keyBytes)
				exitOnError(err, "unable to read private key")

				signer, err := jose.NewSigner(jose.SignatureAlgorithm("RS256"), signingKey)
				exitOnError(err, "unable to make signer")
				obj, err := signer.Sign(readInput(c.String("input")))
				exitOnError(err, "unable to sign")

				var msg string
				if c.Bool("full") {
					msg = obj.FullSerialize()
				} else {
					msg, err = obj.CompactSerialize()
					exitOnError(err, "unable to serialize message")
				}

				writeOutput(c.String("output"), []byte(msg))
			},
		},
	}

	err := app.Run(os.Args)
	exitOnError(err, "unable to run application")
}

// Retrieve value of a required flag
func requiredFlag(c *cli.Context, flag string) string {
	value := c.String(flag)
	if value == "" {
		fmt.Fprintf(os.Stderr, "missing required flag --%s\n", flag)
		os.Exit(1)
	}
	return value
}

// Exit and print error message if we encountered a problem
func exitOnError(err error, msg string) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s: %s\n", msg, err)
		os.Exit(1)
	}
}

// Read input from file or stdin
func readInput(path string) []byte {
	var bytes []byte
	var err error

	if path != "" {
		bytes, err = ioutil.ReadFile(path)
	} else {
		bytes, err = ioutil.ReadAll(os.Stdin)
	}

	exitOnError(err, "unable to read input")
	return bytes
}

// Write output to file or stdin
func writeOutput(path string, data []byte) {
	var err error

	if path != "" {
		err = ioutil.WriteFile(path, data, 0644)
	} else {
		_, err = os.Stdout.Write(data)
	}

	exitOnError(err, "unable to write output")
}
<commit_msg>Add algorithm flag to sign util.<commit_after>/*-
 * Copyright 2014 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/codegangsta/cli"
	"github.com/square/go-jose"
)

func main() {
	app := cli.NewApp()
	app.Name = "jose-util"
	app.Usage = "command-line utility to deal with JOSE objects"
	app.Version = "0.0.2"
	app.Author = ""
	app.Email = ""

	app.Commands = []cli.Command{
		{
			Name:  "encrypt",
			Usage: "encrypt a plaintext",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "key, k",
					Usage: "Path to key file (PEM/DER)",
				},
				cli.StringFlag{
					Name:  "input, in",
					Usage: "Path to input file (stdin if missing)",
				},
				cli.StringFlag{
					Name:  "output, out",
					Usage: "Path to output file (stdout if missing)",
				},
				cli.StringFlag{
					Name:  "algorithm, alg",
					Usage: "Key management algorithm (e.g. RSA-OAEP)",
				},
				cli.StringFlag{
					Name:  "encryption, enc",
					Usage: "Content encryption algorithm (e.g. A128GCM)",
				},
				cli.BoolFlag{
					Name:  "full, f",
					Usage: "Use full serialization format (instead of compact)",
				},
			},
			Action: func(c *cli.Context) {
				keyBytes, err := ioutil.ReadFile(requiredFlag(c, "key"))
				exitOnError(err, "unable to read key file")

				pub, err := jose.LoadPublicKey(keyBytes)
				exitOnError(err, "unable to read public key")

				alg := jose.KeyAlgorithm(requiredFlag(c, "alg"))
				enc := jose.ContentEncryption(requiredFlag(c, "enc"))

				crypter, err := jose.NewEncrypter(alg, enc, pub)
				exitOnError(err, "unable to instantiate encrypter")

				obj, err := crypter.Encrypt(readInput(c.String("input")))
				exitOnError(err, "unable to encrypt")

				var msg string
				if c.Bool("full") {
					msg = obj.FullSerialize()
				} else {
					msg, err = obj.CompactSerialize()
					exitOnError(err, "unable to serialize message")
				}

				writeOutput(c.String("output"), []byte(msg))
			},
		},
		{
			Name:  "decrypt",
			Usage: "decrypt a plaintext",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "key, k",
					Usage: "Path to key file (PEM/DER)",
				},
				cli.StringFlag{
					Name:  "input, in",
					Usage: "Path to input file (stdin if missing)",
				},
				cli.StringFlag{
					Name:  "output, out",
					Usage: "Path to output file (stdout if missing)",
				},
			},
			Action: func(c *cli.Context) {
				keyBytes, err := ioutil.ReadFile(requiredFlag(c, "key"))
				exitOnError(err, "unable to read private key")

				priv, err := jose.LoadPrivateKey(keyBytes)
				exitOnError(err, "unable to read private key")

				obj, err := jose.ParseEncrypted(string(readInput(c.String("input"))))
				exitOnError(err, "unable to parse message")

				plaintext, err := obj.Decrypt(priv)
				exitOnError(err, "unable to decrypt message")

				writeOutput(c.String("output"), plaintext)
			},
		},
		{
			Name:  "dump",
			Usage: "parse & dump message in full serialization format",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "input, in",
					Usage: "Path to input file (stdin if missing)",
				},
				cli.StringFlag{
					Name:  "output, out",
					Usage: "Path to output file (stdout if missing)",
				},
				cli.StringFlag{
					Name:  "format, f",
					Usage: "Message format (JWE/JWS, defaults to JWE)",
				},
			},
			Action: func(c *cli.Context) {
				input := string(readInput(c.String("input")))

				var serialized string
				var err error
				switch c.String("format") {
				case "", "JWE":
					var jwe *jose.JsonWebEncryption
					jwe, err = jose.ParseEncrypted(input)
					if err == nil {
						serialized = jwe.FullSerialize()
					}
				case "JWS":
					var jws *jose.JsonWebSignature
					jws, err = jose.ParseSigned(input)
					if err == nil {
						serialized = jws.FullSerialize()
					}
				}

				exitOnError(err, "unable to parse message")

				var raw map[string]interface{}
				err = json.Unmarshal([]byte(serialized), &raw)
				exitOnError(err, "unable to parse message")

				output, err := json.MarshalIndent(&raw, "", "\t")
				exitOnError(err, "unable to serialize message")

				writeOutput(c.String("output"), output)
			},
		},
		{
			Name:  "sign",
			Usage: "sign a text",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "algorithm, alg",
					Usage: "Key management algorithm (e.g. RSA-OAEP)",
				},
				cli.StringFlag{
					Name:  "key, k",
					Usage: "Path to key file (PEM/DER)",
				},
				cli.StringFlag{
					Name:  "input, in",
					Usage: "Path to input file (stdin if missing)",
				},
				cli.StringFlag{
					Name:  "output, out",
					Usage: "Path to output file (stdout if missing)",
				},
				cli.BoolFlag{
					Name:  "full, f",
					Usage: "Use full serialization format (instead of compact)",
				},
			},
			Action: func(c *cli.Context) {
				keyBytes, err := ioutil.ReadFile(requiredFlag(c, "key"))
				exitOnError(err, "unable to read key file")

				signingKey, err := jose.LoadPrivateKey(keyBytes)
				exitOnError(err, "unable to read private key")

				alg := jose.SignatureAlgorithm(requiredFlag(c, "algorithm"))
				signer, err := jose.NewSigner(alg, signingKey)
				exitOnError(err, "unable to make signer")
				obj, err := signer.Sign(readInput(c.String("input")))
				exitOnError(err, "unable to sign")

				var msg string
				if c.Bool("full") {
					msg = obj.FullSerialize()
				} else {
					msg, err = obj.CompactSerialize()
					exitOnError(err, "unable to serialize message")
				}

				writeOutput(c.String("output"), []byte(msg))
			},
		},
	}

	err := app.Run(os.Args)
	exitOnError(err, "unable to run application")
}

// Retrieve value of a required flag
func requiredFlag(c *cli.Context, flag string) string {
	value := c.String(flag)
	if value == "" {
		fmt.Fprintf(os.Stderr, "missing required flag --%s\n", flag)
		os.Exit(1)
	}
	return value
}

// Exit and print error message if we encountered a problem
func exitOnError(err error, msg string) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s: %s\n", msg, err)
		os.Exit(1)
	}
}

// Read input from file or stdin
func readInput(path string) []byte {
	var bytes []byte
	var err error

	if path != "" {
		bytes, err = ioutil.ReadFile(path)
	} else {
		bytes, err = ioutil.ReadAll(os.Stdin)
	}

	exitOnError(err, "unable to read input")
	return bytes
}

// Write output to file or stdin
func writeOutput(path string, data []byte) {
	var err error

	if path != "" {
		err = ioutil.WriteFile(path, data, 0644)
	} else {
		_, err = os.Stdout.Write(data)
	}

	exitOnError(err, "unable to write output")
}
<|endoftext|>


<commit_before>// Package alertmanager implements a Service capable of processing webhooks from prometheus alertmanager.
package alertmanager

import (
	"bytes"
	"encoding/json"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/matrix-org/go-neb/database"
	"github.com/matrix-org/go-neb/types"
	"github.com/matrix-org/gomatrix"
	html "html/template"
	"net/http"
	text "text/template"
)

// ServiceType of the Alertmanager service.
const ServiceType = "alertmanager"

// Service contains the Config fields for the Alertmanager service.
//
// This service will send notifications into a Matrix room when Alertmanager sends
// webhook events to it. It requires a public domain which Alertmanager can reach.
// Notices will be sent as the service user ID.
//
// For the template strings, take a look at https://golang.org/pkg/text/template/
// and the html variant https://golang.org/pkg/html/template/.
// The data they get is a webhookNotification
//
// You can set msg_type to either m.text or m.notice
//
// Example JSON request:
//    {
//        rooms: {
//            "!ewfug483gsfe:localhost": {
//                "text_template": "your plain text template goes here",
//                "html_template": "your html template goes here",
//                "msg_type": "m.text"
//            },
//        }
//    }
type Service struct {
	types.DefaultService
	webhookEndpointURL string
	// The URL which should be added to alertmanagers config - Populated by Go-NEB after Service registration.
	WebhookURL string `json:"webhook_url"`
	// A map of matrix rooms to templates
	Rooms map[string]struct {
		TextTemplate string `json:"text_template"`
		HTMLTemplate string `json:"html_template"`
		MsgType      string `json:"msg_type"`
	} `json:"rooms"`
}

// WebhookNotification is the payload from Alertmanager
type WebhookNotification struct {
	Version           string            `json:"version"`
	GroupKey          string            `json:"groupKey"`
	Status            string            `json:"status"`
	Receiver          string            `json:"receiver"`
	GroupLabels       map[string]string `json:"groupLabels"`
	CommonLabels      map[string]string `json:"commonLabels"`
	CommonAnnotations map[string]string `json:"commonAnnotations"`
	ExternalURL       string            `json:"externalURL"`
	Alerts            []struct {
		Status       string            `json:"status"`
		Labels       map[string]string `json:"labels"`
		Annotations  map[string]string `json:"annotations"`
		StartsAt     string            `json:"startsAt"`
		EndsAt       string            `json:"endsAt"`
		GeneratorURL string            `json:"generatorURL"`
	} `json:"alerts"`
}

// OnReceiveWebhook receives requests from Alertmanager and sends requests to Matrix as a result.
func (s *Service) OnReceiveWebhook(w http.ResponseWriter, req *http.Request, cli *gomatrix.Client) {
	decoder := json.NewDecoder(req.Body)
	var notif WebhookNotification
	if err := decoder.Decode(&notif); err != nil {
		log.WithError(err).Error("Alertmanager webhook received an invalid JSON payload")
		w.WriteHeader(400)
		return
	}

	for roomID, templates := range s.Rooms {
		var msg interface{}
		// we don't check whether the templates parse because we already did when storing them in the db
		textTemplate, _ := text.New("textTemplate").Parse(templates.TextTemplate)
		var bodyBuffer bytes.Buffer
		textTemplate.Execute(&bodyBuffer, notif)
		if templates.HTMLTemplate != "" {
			// we don't check whether the templates parse because we already did when storing them in the db
			htmlTemplate, _ := html.New("htmlTemplate").Parse(templates.HTMLTemplate)
			var formattedBodyBuffer bytes.Buffer
			htmlTemplate.Execute(&formattedBodyBuffer, notif)
			msg = gomatrix.HTMLMessage{
				Body:          bodyBuffer.String(),
				MsgType:       templates.MsgType,
				Format:        "org.matrix.custom.html",
				FormattedBody: formattedBodyBuffer.String(),
			}
		} else {
			msg = gomatrix.TextMessage{
				Body:    bodyBuffer.String(),
				MsgType: templates.MsgType,
			}
		}

		log.WithFields(log.Fields{
			"message": msg,
			"room_id": roomID,
		}).Print("Sending Alertmanager notification to room")
		if _, e := cli.SendMessageEvent(roomID, "m.room.message", msg); e != nil {
			log.WithError(e).WithField("room_id", roomID).Print(
				"Failed to send Alertmanager notification to room.")
		}
	}
	w.WriteHeader(200)
}

// Register makes sure the Config information supplied is valid.
func (s *Service) Register(oldService types.Service, client *gomatrix.Client) error {
	s.WebhookURL = s.webhookEndpointURL
	for _, templates := range s.Rooms {
		// validate that we have at least a plain text template
		if templates.TextTemplate == "" {
			return fmt.Errorf("plain text template missing")
		}

		// validate the plain text template is valid
		_, err := text.New("textTemplate").Parse(templates.TextTemplate)
		if err != nil {
			return fmt.Errorf("plain text template is invalid")
		}

		if templates.HTMLTemplate != "" {
			// validate that the html template is valid
			_, err := html.New("htmlTemplate").Parse(templates.HTMLTemplate)
			if err != nil {
				return fmt.Errorf("html template is invalid")
			}
		}
		// validate that the msgtype is either m.notice or m.text
		if templates.MsgType != "m.notice" && templates.MsgType != "m.text" {
			return fmt.Errorf("msg_type is neither 'm.notice' nor 'm.text'")
		}
	}
	s.joinRooms(client)
	return nil
}

// PostRegister deletes this service if there are no registered repos.
func (s *Service) PostRegister(oldService types.Service) {
	// At least one room still active
	if len(s.Rooms) > 0 {
		return
	}
	// Delete this service since no repos are configured
	logger := log.WithFields(log.Fields{
		"service_type": s.ServiceType(),
		"service_id":   s.ServiceID(),
	})
	logger.Info("Removing service as no repositories are registered.")
	if err := database.GetServiceDB().DeleteService(s.ServiceID()); err != nil {
		logger.WithError(err).Error("Failed to delete service")
	}
}

func (s *Service) joinRooms(client *gomatrix.Client) {
	for roomID := range s.Rooms {
		if _, err := client.JoinRoom(roomID, "", nil); err != nil {
			log.WithFields(log.Fields{
				log.ErrorKey: err,
				"room_id":    roomID,
				"user_id":    client.UserID,
			}).Error("Failed to join room")
		}
	}
}

func init() {
	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
		return &Service{
			DefaultService:     types.NewDefaultService(serviceID, serviceUserID, ServiceType),
			webhookEndpointURL: webhookEndpointURL,
		}
	})
}
<commit_msg>Better error handling for the alertmanager service (#306)<commit_after>// Package alertmanager implements a Service capable of processing webhooks from prometheus alertmanager.
package alertmanager

import (
	"bytes"
	"encoding/json"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/matrix-org/go-neb/database"
	"github.com/matrix-org/go-neb/types"
	"github.com/matrix-org/gomatrix"
	html "html/template"
	"net/http"
	text "text/template"
)

// ServiceType of the Alertmanager service.
const ServiceType = "alertmanager"

// Service contains the Config fields for the Alertmanager service.
//
// This service will send notifications into a Matrix room when Alertmanager sends
// webhook events to it. It requires a public domain which Alertmanager can reach.
// Notices will be sent as the service user ID.
//
// For the template strings, take a look at https://golang.org/pkg/text/template/
// and the html variant https://golang.org/pkg/html/template/.
// The data they get is a webhookNotification
//
// You can set msg_type to either m.text or m.notice
//
// Example JSON request:
//    {
//        rooms: {
//            "!ewfug483gsfe:localhost": {
//                "text_template": "your plain text template goes here",
//                "html_template": "your html template goes here",
//                "msg_type": "m.text"
//            },
//        }
//    }
type Service struct {
	types.DefaultService
	webhookEndpointURL string
	// The URL which should be added to alertmanagers config - Populated by Go-NEB after Service registration.
	WebhookURL string `json:"webhook_url"`
	// A map of matrix rooms to templates
	Rooms map[string]struct {
		TextTemplate string `json:"text_template"`
		HTMLTemplate string `json:"html_template"`
		MsgType      string `json:"msg_type"`
	} `json:"rooms"`
}

// WebhookNotification is the payload from Alertmanager
type WebhookNotification struct {
	Version           string            `json:"version"`
	GroupKey          string            `json:"groupKey"`
	Status            string            `json:"status"`
	Receiver          string            `json:"receiver"`
	GroupLabels       map[string]string `json:"groupLabels"`
	CommonLabels      map[string]string `json:"commonLabels"`
	CommonAnnotations map[string]string `json:"commonAnnotations"`
	ExternalURL       string            `json:"externalURL"`
	Alerts            []struct {
		Status       string            `json:"status"`
		Labels       map[string]string `json:"labels"`
		Annotations  map[string]string `json:"annotations"`
		StartsAt     string            `json:"startsAt"`
		EndsAt       string            `json:"endsAt"`
		GeneratorURL string            `json:"generatorURL"`
	} `json:"alerts"`
}

// OnReceiveWebhook receives requests from Alertmanager and sends requests to Matrix as a result.
func (s *Service) OnReceiveWebhook(w http.ResponseWriter, req *http.Request, cli *gomatrix.Client) {
	decoder := json.NewDecoder(req.Body)
	var notif WebhookNotification
	if err := decoder.Decode(&notif); err != nil {
		log.WithError(err).Error("Alertmanager webhook received an invalid JSON payload")
		w.WriteHeader(400)
		return
	}

	for roomID, templates := range s.Rooms {
		var msg interface{}
		// we don't check whether the templates parse because we already did when storing them in the db
		textTemplate, _ := text.New("textTemplate").Parse(templates.TextTemplate)
		var bodyBuffer bytes.Buffer
		if err := textTemplate.Execute(&bodyBuffer, notif); err != nil {
			log.WithError(err).Error("Alertmanager webhook failed to execute text template")
			w.WriteHeader(500)
			return
		}
		if templates.HTMLTemplate != "" {
			// we don't check whether the templates parse because we already did when storing them in the db
			htmlTemplate, _ := html.New("htmlTemplate").Parse(templates.HTMLTemplate)
			var formattedBodyBuffer bytes.Buffer
			if err := htmlTemplate.Execute(&formattedBodyBuffer, notif); err != nil {
				log.WithError(err).Error("Alertmanager webhook failed to execute HTML template")
				w.WriteHeader(500)
				return
			}
			msg = gomatrix.HTMLMessage{
				Body:          bodyBuffer.String(),
				MsgType:       templates.MsgType,
				Format:        "org.matrix.custom.html",
				FormattedBody: formattedBodyBuffer.String(),
			}
		} else {
			msg = gomatrix.TextMessage{
				Body:    bodyBuffer.String(),
				MsgType: templates.MsgType,
			}
		}

		log.WithFields(log.Fields{
			"message": msg,
			"room_id": roomID,
		}).Print("Sending Alertmanager notification to room")
		if _, e := cli.SendMessageEvent(roomID, "m.room.message", msg); e != nil {
			log.WithError(e).WithField("room_id", roomID).Print(
				"Failed to send Alertmanager notification to room.")
		}
	}
	w.WriteHeader(200)
}

// Register makes sure the Config information supplied is valid.
func (s *Service) Register(oldService types.Service, client *gomatrix.Client) error {
	s.WebhookURL = s.webhookEndpointURL
	for _, templates := range s.Rooms {
		// validate that we have at least a plain text template
		if templates.TextTemplate == "" {
			return fmt.Errorf("plain text template missing")
		}

		// validate the plain text template is valid
		_, err := text.New("textTemplate").Parse(templates.TextTemplate)
		if err != nil {
			return fmt.Errorf("plain text template is invalid: %v", err)
		}

		if templates.HTMLTemplate != "" {
			// validate that the html template is valid
			_, err := html.New("htmlTemplate").Parse(templates.HTMLTemplate)
			if err != nil {
				return fmt.Errorf("html template is invalid: %v", err)
			}
		}
		// validate that the msgtype is either m.notice or m.text
		if templates.MsgType != "m.notice" && templates.MsgType != "m.text" {
			return fmt.Errorf("msg_type is neither 'm.notice' nor 'm.text'")
		}
	}
	s.joinRooms(client)
	return nil
}

// PostRegister deletes this service if there are no registered repos.
func (s *Service) PostRegister(oldService types.Service) {
	// At least one room still active
	if len(s.Rooms) > 0 {
		return
	}
	// Delete this service since no repos are configured
	logger := log.WithFields(log.Fields{
		"service_type": s.ServiceType(),
		"service_id":   s.ServiceID(),
	})
	logger.Info("Removing service as no repositories are registered.")
	if err := database.GetServiceDB().DeleteService(s.ServiceID()); err != nil {
		logger.WithError(err).Error("Failed to delete service")
	}
}

func (s *Service) joinRooms(client *gomatrix.Client) {
	for roomID := range s.Rooms {
		if _, err := client.JoinRoom(roomID, "", nil); err != nil {
			log.WithFields(log.Fields{
				log.ErrorKey: err,
				"room_id":    roomID,
				"user_id":    client.UserID,
			}).Error("Failed to join room")
		}
	}
}

func init() {
	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
		return &Service{
			DefaultService:     types.NewDefaultService(serviceID, serviceUserID, ServiceType),
			webhookEndpointURL: webhookEndpointURL,
		}
	})
}
<|endoftext|>


<commit_before>package azurerm

import (
	"fmt"
	"testing"

	"github.com/Azure/azure-sdk-for-go/arm/network"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAzureRMLoadBalancerNatRule_basic(t *testing.T) {
	var lb network.LoadBalancer
	ri := acctest.RandInt()
	natRuleName := fmt.Sprintf("NatRule-%d", ri)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMLoadBalancerDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMLoadBalancerNatRule_basic(ri, natRuleName),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMLoadBalancerExists("azurerm_lb.test", &lb),
					testCheckAzureRMLoadBalancerNatRuleExists(natRuleName, &lb),
				),
			},
		},
	})
}

func TestAccAzureRMLoadBalancerNatRule_removal(t *testing.T) {
	var lb network.LoadBalancer
	ri := acctest.RandInt()
	natRuleName := fmt.Sprintf("NatRule-%d", ri)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMLoadBalancerDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMLoadBalancerNatRule_basic(ri, natRuleName),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMLoadBalancerExists("azurerm_lb.test", &lb),
					testCheckAzureRMLoadBalancerNatRuleExists(natRuleName, &lb),
				),
			},
			{
				Config: testAccAzureRMLoadBalancerNatRule_removal(ri),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMLoadBalancerExists("azurerm_lb.test", &lb),
					testCheckAzureRMLoadBalancerNatRuleNotExists(natRuleName, &lb),
				),
			},
		},
	})
}

func testCheckAzureRMLoadBalancerNatRuleExists(natRuleName string, lb *network.LoadBalancer) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		_, _, exists := findLoadBalancerNatRuleByName(lb, natRuleName)
		if !exists {
			return fmt.Errorf("A NAT Rule with name %q cannot be found.", natRuleName)
		}

		return nil
	}
}

func testCheckAzureRMLoadBalancerNatRuleNotExists(natRuleName string, lb *network.LoadBalancer) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		_, _, exists := findLoadBalancerNatRuleByName(lb, natRuleName)
		if exists {
			return fmt.Errorf("A NAT Rule with name %q has been found.", natRuleName)
		}

		return nil
	}
}

func testAccAzureRMLoadBalancerNatRule_basic(rInt int, natRuleName string) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test" {
    name = "acctestrg-%d"
    location = "West US"
}

resource "azurerm_public_ip" "test" {
    name = "test-ip-%d"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"
    public_ip_address_allocation = "static"
}

resource "azurerm_lb" "test" {
    name = "arm-test-loadbalancer-%d"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"

    frontend_ip_configuration {
      name = "one-%d"
      public_ip_address_id = "${azurerm_public_ip.test.id}"
    }
}

resource "azurerm_lb_nat_rule" "test" {
  location = "West US"
  resource_group_name = "${azurerm_resource_group.test.name}"
  loadbalancer_id = "${azurerm_lb.test.id}"
  name = "%s"
  protocol = "Tcp"
  frontend_port = 3389
  backend_port = 3389
  frontend_ip_configuration_name = "one-%d"
}

`, rInt, rInt, rInt, rInt, natRuleName, rInt)
}

func testAccAzureRMLoadBalancerNatRule_removal(rInt int) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test" {
    name = "acctestrg-%d"
    location = "West US"
}

resource "azurerm_public_ip" "test" {
    name = "test-ip-%d"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"
    public_ip_address_allocation = "static"
}

resource "azurerm_lb" "test" {
    name = "arm-test-loadbalancer-%d"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"

    frontend_ip_configuration {
      name = "one-%d"
      public_ip_address_id = "${azurerm_public_ip.test.id}"
    }
}
`, rInt, rInt, rInt, rInt)
}
<commit_msg>Update test for azurerm_lb_nat_rule id.<commit_after>package azurerm

import (
	"fmt"
	"os"
	"testing"

	"github.com/Azure/azure-sdk-for-go/arm/network"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAzureRMLoadBalancerNatRule_basic(t *testing.T) {
	var lb network.LoadBalancer
	ri := acctest.RandInt()
	natRuleName := fmt.Sprintf("NatRule-%d", ri)

	testAccPreCheck(t)
	subscriptionID := os.Getenv("ARM_SUBSCRIPTION_ID")
	natRule_id := fmt.Sprintf(
		"/subscriptions/%s/resourceGroups/acctestrg-%d/providers/Microsoft.Network/loadBalancers/arm-test-loadbalancer-%d/inboundNatRules/%s",
		subscriptionID, ri, ri, natRuleName)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMLoadBalancerDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMLoadBalancerNatRule_basic(ri, natRuleName),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMLoadBalancerExists("azurerm_lb.test", &lb),
					testCheckAzureRMLoadBalancerNatRuleExists(natRuleName, &lb),
					resource.TestCheckResourceAttr(
						"azurerm_lb_nat_rule.test", "id", natRule_id),
				),
			},
		},
	})
}

func TestAccAzureRMLoadBalancerNatRule_removal(t *testing.T) {
	var lb network.LoadBalancer
	ri := acctest.RandInt()
	natRuleName := fmt.Sprintf("NatRule-%d", ri)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMLoadBalancerDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMLoadBalancerNatRule_basic(ri, natRuleName),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMLoadBalancerExists("azurerm_lb.test", &lb),
					testCheckAzureRMLoadBalancerNatRuleExists(natRuleName, &lb),
				),
			},
			{
				Config: testAccAzureRMLoadBalancerNatRule_removal(ri),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMLoadBalancerExists("azurerm_lb.test", &lb),
					testCheckAzureRMLoadBalancerNatRuleNotExists(natRuleName, &lb),
				),
			},
		},
	})
}

func testCheckAzureRMLoadBalancerNatRuleExists(natRuleName string, lb *network.LoadBalancer) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		_, _, exists := findLoadBalancerNatRuleByName(lb, natRuleName)
		if !exists {
			return fmt.Errorf("A NAT Rule with name %q cannot be found.", natRuleName)
		}

		return nil
	}
}

func testCheckAzureRMLoadBalancerNatRuleNotExists(natRuleName string, lb *network.LoadBalancer) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		_, _, exists := findLoadBalancerNatRuleByName(lb, natRuleName)
		if exists {
			return fmt.Errorf("A NAT Rule with name %q has been found.", natRuleName)
		}

		return nil
	}
}

func testAccAzureRMLoadBalancerNatRule_basic(rInt int, natRuleName string) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test" {
    name = "acctestrg-%d"
    location = "West US"
}

resource "azurerm_public_ip" "test" {
    name = "test-ip-%d"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"
    public_ip_address_allocation = "static"
}

resource "azurerm_lb" "test" {
    name = "arm-test-loadbalancer-%d"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"

    frontend_ip_configuration {
      name = "one-%d"
      public_ip_address_id = "${azurerm_public_ip.test.id}"
    }
}

resource "azurerm_lb_nat_rule" "test" {
  location = "West US"
  resource_group_name = "${azurerm_resource_group.test.name}"
  loadbalancer_id = "${azurerm_lb.test.id}"
  name = "%s"
  protocol = "Tcp"
  frontend_port = 3389
  backend_port = 3389
  frontend_ip_configuration_name = "one-%d"
}

`, rInt, rInt, rInt, rInt, natRuleName, rInt)
}

func testAccAzureRMLoadBalancerNatRule_removal(rInt int) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test" {
    name = "acctestrg-%d"
    location = "West US"
}

resource "azurerm_public_ip" "test" {
    name = "test-ip-%d"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"
    public_ip_address_allocation = "static"
}

resource "azurerm_lb" "test" {
    name = "arm-test-loadbalancer-%d"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"

    frontend_ip_configuration {
      name = "one-%d"
      public_ip_address_id = "${azurerm_public_ip.test.id}"
    }
}
`, rInt, rInt, rInt, rInt)
}
<|endoftext|>


<commit_before>package internet_speed

import (
	"fmt"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/showwin/speedtest-go/speedtest"
)

// InternetSpeed is used to store configuration values.
type InternetSpeed struct {
	EnableFileDownload bool            `toml:"enable_file_download"`
	Log                telegraf.Logger `toml:"-"`
}

const sampleConfig = `
  ## Sets if runs file download test
  ## Default: false
  enable_file_download = false
`

// Description returns information about the plugin.
func (is *InternetSpeed) Description() string {
	return "Monitors internet speed using speedtest.net service"
}

// SampleConfig displays configuration instructions.
func (is *InternetSpeed) SampleConfig() string {
	return sampleConfig
}

const measurement = "internet_speed"

func (is *InternetSpeed) Gather(acc telegraf.Accumulator) error {
	user, err := speedtest.FetchUserInfo()
	if err != nil {
		return fmt.Errorf("fetching user info failed: %v", err)
	}
	serverList, err := speedtest.FetchServerList(user)
	if err != nil {
		return fmt.Errorf("fetching server list failed: %v", err)
	}

	if len(serverList.Servers) < 1 {
		return fmt.Errorf("no servers found")
	}
	s := serverList.Servers[0]
	is.Log.Debug("Starting Speed Test")
	is.Log.Debug("Running Ping...")
	err = s.PingTest()
	if err != nil {
		return fmt.Errorf("ping test failed: %v", err)
	}
	is.Log.Debug("Running Download...")
	err = s.DownloadTest(is.EnableFileDownload)
	if err != nil {
		return fmt.Errorf("download test failed: %v", err)
	}
	is.Log.Debug("Running Upload...")
	err = s.UploadTest(is.EnableFileDownload)
	if err != nil {
		return fmt.Errorf("upload test failed failed: %v", err)
	}

	is.Log.Debug("Test finished.")

	fields := make(map[string]interface{})
	fields["download"] = s.DLSpeed
	fields["upload"] = s.ULSpeed
	fields["latency"] = s.Latency

	tags := make(map[string]string)

	acc.AddFields(measurement, fields, tags)
	return nil
}
func init() {
	inputs.Add("internet_speed", func() telegraf.Input {
		return &InternetSpeed{}
	})
}
<commit_msg>fix: internet_speed input plugin not collecting/reporting latency (#9957)<commit_after>package internet_speed

import (
	"fmt"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
	"github.com/showwin/speedtest-go/speedtest"
)

// InternetSpeed is used to store configuration values.
type InternetSpeed struct {
	EnableFileDownload bool            `toml:"enable_file_download"`
	Log                telegraf.Logger `toml:"-"`
}

const sampleConfig = `
  ## Sets if runs file download test
  ## Default: false
  enable_file_download = false
`

// Description returns information about the plugin.
func (is *InternetSpeed) Description() string {
	return "Monitors internet speed using speedtest.net service"
}

// SampleConfig displays configuration instructions.
func (is *InternetSpeed) SampleConfig() string {
	return sampleConfig
}

const measurement = "internet_speed"

func (is *InternetSpeed) Gather(acc telegraf.Accumulator) error {
	user, err := speedtest.FetchUserInfo()
	if err != nil {
		return fmt.Errorf("fetching user info failed: %v", err)
	}
	serverList, err := speedtest.FetchServerList(user)
	if err != nil {
		return fmt.Errorf("fetching server list failed: %v", err)
	}

	if len(serverList.Servers) < 1 {
		return fmt.Errorf("no servers found")
	}
	s := serverList.Servers[0]
	is.Log.Debug("Starting Speed Test")
	is.Log.Debug("Running Ping...")
	err = s.PingTest()
	if err != nil {
		return fmt.Errorf("ping test failed: %v", err)
	}
	is.Log.Debug("Running Download...")
	err = s.DownloadTest(is.EnableFileDownload)
	if err != nil {
		return fmt.Errorf("download test failed: %v", err)
	}
	is.Log.Debug("Running Upload...")
	err = s.UploadTest(is.EnableFileDownload)
	if err != nil {
		return fmt.Errorf("upload test failed failed: %v", err)
	}

	is.Log.Debug("Test finished.")

	fields := make(map[string]interface{})
	fields["download"] = s.DLSpeed
	fields["upload"] = s.ULSpeed
	fields["latency"] = timeDurationMillisecondToFloat64(s.Latency)

	tags := make(map[string]string)

	acc.AddFields(measurement, fields, tags)
	return nil
}
func init() {
	inputs.Add("internet_speed", func() telegraf.Input {
		return &InternetSpeed{}
	})
}

func timeDurationMillisecondToFloat64(d time.Duration) float64 {
	return float64(d) / float64(time.Millisecond)
}
<|endoftext|>