<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 exec

import (
	"context"
	"reflect"
	"testing"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/funcx"
	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/mtime"
	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/window"
	"github.com/apache/beam/sdks/go/pkg/beam/core/typex"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
)

// TestInvoke verifies the the various forms of input to Invoke are handled correctly.
func TestInvoke(t *testing.T) {
	tests := []struct {
		Fn       interface{}
		Opt      *MainInput
		Args     []interface{}
		Expected interface{}
	}{
		{
			// Void function
			Fn: func() {},
		},
		{
			// Contexts and errors are allowed and handled.
			Fn: func(ctx context.Context) error { return nil },
		},
		{
			// Sum
			Fn:       func(a, b, c int) int { return a + b + c },
			Args:     []interface{}{1, 2, 3},
			Expected: 6,
		},
		{
			// Concat
			Fn:       func(a, b, c string) string { return a + b + c },
			Args:     []interface{}{"a", "b", "c"},
			Expected: "abc",
		},
		{
			// Length (slice type)
			Fn:       func(list []int) (int, error) { return len(list), nil },
			Args:     []interface{}{[]int{1, 2, 3}},
			Expected: 3,
		},
		{
			// Emitter
			Fn:   func(emit func(int)) { emit(1) },
			Args: []interface{}{func(int) {}},
		},
		{
			// Side input
			Fn: func(a int, get func(*int) bool) int {
				var ret int
				if !get(&ret) {
					return a
				}
				return ret
			},
			Args:     []interface{}{1, func(out *int) bool { *out = 2; return true }},
			Expected: 2,
		},
		{
			// Sum as Main
			Fn:       func(a, b, c int) int { return a + b + c },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Args:     []interface{}{2, 3},
			Expected: 6,
		},
		{
			// Sum as Main KV
			Fn:       func(a, b, c int) int { return a + b + c },
			Opt:      &MainInput{Key: FullValue{Elm: 1, Elm2: 2}},
			Args:     []interface{}{3},
			Expected: 6,
		},
		{
			// EventTime
			Fn:       func(ts typex.EventTime, a int) int { return int(ts.Milliseconds()) + a },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Expected: 3,
		},
		{
			// Window
			Fn:       func(w typex.Window, a int) int64 { return w.MaxTimestamp().Milliseconds() },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Expected: mtime.EndOfGlobalWindowTime.Milliseconds(),
		},
	}

	for _, test := range tests {
		fn, err := funcx.New(reflectx.MakeFunc(test.Fn))
		if err != nil {
			t.Fatalf("function not valid: %v", err)
		}

		ts := mtime.ZeroTimestamp.Add(2 * time.Millisecond)

		val, err := Invoke(context.Background(), window.SingleGlobalWindow, ts, fn, test.Opt, test.Args...)
		if err != nil {
			t.Fatalf("Invoke(%v,%v) failed: %v", fn.Fn.Name(), test.Args, err)
		}
		if val != nil && val.Elm != test.Expected {
			t.Errorf("Invoke(%v,%v) = %v, want %v", fn.Fn.Name(), test.Args, val.Elm, test.Expected)
		}
	}
}

// Benchmarks

// NOTE(herohde) 12/19/2017: example run on a laptop
//
// BenchmarkDirectCall-4              2000000000	         0.26 ns/op
// BenchmarkIndirectCall-4            1000000000	         2.28 ns/op
// BenchmarkReflectedAndBackCall-4    1000000000	         2.40 ns/op
// BenchmarkReflectCall-4            	10000000	       197    ns/op
// BenchmarkReflectCallNewArgs-4     	10000000	       200    ns/op
// BenchmarkReflectCallReuseArgs-4   	10000000	       196    ns/op
// BenchmarkInvokeCall-4             	 3000000	       452    ns/op
// BenchmarkInvokeCallExtra-4        	 3000000	       419    ns/op
// BenchmarkReflectFnCall-4          	 5000000	       340    ns/op
// BenchmarkInvokeFnCall-4           	 2000000	       604    ns/op
// BenchmarkInvokeFnCallExtra-4      	 3000000	       559    ns/op

func inc(n int) int {
	return n + 1
}

func BenchmarkDirectCall(b *testing.B) {
	n := 0
	for i := 0; i < b.N; i++ {
		n = inc(n)
	}
	b.Log(n)
}

func BenchmarkIndirectCall(b *testing.B) {
	fn := func(int) int { return 0 }
	if b.N > 0 {
		fn = inc // this is always set, but the compiler doesn't know
	}

	n := 0
	for i := 0; i < b.N; i++ {
		n = fn(n)
	}
	b.Log(n)
}

func BenchmarkReflectedAndBackCall(b *testing.B) {
	fn := reflect.ValueOf(inc).Interface().(func(int) int)
	n := 0
	for i := 0; i < b.N; i++ {
		n = fn(n)
	}
	b.Log(n)
}

func BenchmarkReflectCall(b *testing.B) {
	fn := reflect.ValueOf(inc)
	n := reflect.ValueOf(0)
	for i := 0; i < b.N; i++ {
		n = fn.Call([]reflect.Value{n})[0]
	}
	b.Log(n.Interface())
}

func BenchmarkReflectCallNewArgs(b *testing.B) {
	fn := reflect.ValueOf(inc)
	n := reflect.ValueOf(0)
	for i := 0; i < b.N; i++ {
		args := make([]reflect.Value, 1, 1)
		args[0] = n
		n = fn.Call(args)[0]
	}
	b.Log(n.Interface())
}

func BenchmarkReflectCallReuseArgs(b *testing.B) {
	fn := reflect.ValueOf(inc)
	n := reflect.ValueOf(0)
	args := make([]reflect.Value, 1, 1)
	for i := 0; i < b.N; i++ {
		args[0] = n
		n = fn.Call(args)[0]
	}
	b.Log(n.Interface())
}

func BenchmarkInvokeCall(b *testing.B) {
	fn, _ := funcx.New(reflectx.MakeFunc(inc))
	ctx := context.Background()
	n := 0
	for i := 0; i < b.N; i++ {
		ret, _ := InvokeWithoutEventTime(ctx, fn, &MainInput{Key: FullValue{Elm: n}})
		n = ret.Elm.(int)
	}
	b.Log(n)
}

func BenchmarkInvokeCallExtra(b *testing.B) {
	fn, _ := funcx.New(reflectx.MakeFunc(inc))
	ctx := context.Background()
	n := 0
	for i := 0; i < b.N; i++ {
		ret, _ := InvokeWithoutEventTime(ctx, fn, nil, n)
		n = ret.Elm.(int)
	}
	b.Log(n)
}

// The below take the additional overhead of MakeFunc.

func incFn(args []reflect.Value) []reflect.Value {
	return []reflect.Value{reflect.ValueOf(args[0].Interface().(int) + 1)}
}

func BenchmarkReflectFnCall(b *testing.B) {
	fn := reflect.MakeFunc(reflect.TypeOf(inc), incFn)
	n := reflect.ValueOf(0)
	for i := 0; i < b.N; i++ {
		n = fn.Call([]reflect.Value{n})[0]
	}
	b.Log(n.Interface())
}

func BenchmarkInvokeFnCall(b *testing.B) {
	fn, _ := funcx.New(reflectx.MakeFunc(reflect.MakeFunc(reflect.TypeOf(inc), incFn).Interface().(func(int) int)))
	ctx := context.Background()
	n := 0
	for i := 0; i < b.N; i++ {
		ret, _ := InvokeWithoutEventTime(ctx, fn, &MainInput{Key: FullValue{Elm: n}})
		n = ret.Elm.(int)
	}
	b.Log(n)
}

func BenchmarkInvokeFnCallExtra(b *testing.B) {
	fn, _ := funcx.New(reflectx.MakeFunc(reflect.MakeFunc(reflect.TypeOf(inc), incFn).Interface().(func(int) int)))
	ctx := context.Background()
	n := 0
	for i := 0; i < b.N; i++ {
		ret, _ := InvokeWithoutEventTime(ctx, fn, nil, n)
		n = ret.Elm.(int)
	}
	b.Log(n)
}
<commit_msg>Add comparitive benchmark between cached and single use invoker.<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 exec

import (
	"context"
	"fmt"
	"reflect"
	"testing"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/funcx"
	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/mtime"
	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/window"
	"github.com/apache/beam/sdks/go/pkg/beam/core/typex"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
)

// TestInvoke verifies the the various forms of input to Invoke are handled correctly.
func TestInvoke(t *testing.T) {
	tests := []struct {
		Fn       interface{}
		Opt      *MainInput
		Args     []interface{}
		Expected interface{}
	}{
		{
			// Void function
			Fn: func() {},
		},
		{
			// Contexts and errors are allowed and handled.
			Fn: func(ctx context.Context) error { return nil },
		},
		{
			// Sum
			Fn:       func(a, b, c int) int { return a + b + c },
			Args:     []interface{}{1, 2, 3},
			Expected: 6,
		},
		{
			// Concat
			Fn:       func(a, b, c string) string { return a + b + c },
			Args:     []interface{}{"a", "b", "c"},
			Expected: "abc",
		},
		{
			// Length (slice type)
			Fn:       func(list []int) (int, error) { return len(list), nil },
			Args:     []interface{}{[]int{1, 2, 3}},
			Expected: 3,
		},
		{
			// Emitter
			Fn:   func(emit func(int)) { emit(1) },
			Args: []interface{}{func(int) {}},
		},
		{
			// Side input
			Fn: func(a int, get func(*int) bool) int {
				var ret int
				if !get(&ret) {
					return a
				}
				return ret
			},
			Args:     []interface{}{1, func(out *int) bool { *out = 2; return true }},
			Expected: 2,
		},
		{
			// Sum as Main
			Fn:       func(a, b, c int) int { return a + b + c },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Args:     []interface{}{2, 3},
			Expected: 6,
		},
		{
			// Sum as Main KV
			Fn:       func(a, b, c int) int { return a + b + c },
			Opt:      &MainInput{Key: FullValue{Elm: 1, Elm2: 2}},
			Args:     []interface{}{3},
			Expected: 6,
		},
		{
			// EventTime
			Fn:       func(ts typex.EventTime, a int) int { return int(ts.Milliseconds()) + a },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Expected: 3,
		},
		{
			// Window
			Fn:       func(w typex.Window, a int) int64 { return w.MaxTimestamp().Milliseconds() },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Expected: mtime.EndOfGlobalWindowTime.Milliseconds(),
		},
	}

	for _, test := range tests {
		fn, err := funcx.New(reflectx.MakeFunc(test.Fn))
		if err != nil {
			t.Fatalf("function not valid: %v", err)
		}

		ts := mtime.ZeroTimestamp.Add(2 * time.Millisecond)

		val, err := Invoke(context.Background(), window.SingleGlobalWindow, ts, fn, test.Opt, test.Args...)
		if err != nil {
			t.Fatalf("Invoke(%v,%v) failed: %v", fn.Fn.Name(), test.Args, err)
		}
		if val != nil && val.Elm != test.Expected {
			t.Errorf("Invoke(%v,%v) = %v, want %v", fn.Fn.Name(), test.Args, val.Elm, test.Expected)
		}
	}
}

// Benchmarks

// Invoke is implemented as a single use of a cached invoker, so a measure of
// correctness is that the cached version takes less time per operation than the
// single use version.
//
// NOTE(lostluck) 2018/07/24: run on a amd64 linux desktop
//
// BenchmarkInvoke/SingleInvoker_Void_function-12         	 5000000	       300 ns/op	     128 B/op	       1 allocs/op
// BenchmarkInvoke/CachedInvoker_Void_function-12         	10000000	       195 ns/op	       0 B/op	       0 allocs/op
// BenchmarkInvoke/SingleInvoker_Contexts_and_errors_are_allowed_and_handled-12         	 2000000	       967 ns/op	     256 B/op	       6 allocs/op
// BenchmarkInvoke/CachedInvoker_Contexts_and_errors_are_allowed_and_handled-12         	 2000000	       854 ns/op	     112 B/op	       4 allocs/op
// BenchmarkInvoke/SingleInvoker_Sum-12                                                 	 2000000	       805 ns/op	     464 B/op	      11 allocs/op
// BenchmarkInvoke/CachedInvoker_Sum-12                                                 	 3000000	       523 ns/op	     224 B/op	       5 allocs/op
// BenchmarkInvoke/SingleInvoker_Concat-12                                              	 2000000	       892 ns/op	     504 B/op	      12 allocs/op
// BenchmarkInvoke/CachedInvoker_Concat-12                                              	 2000000	       594 ns/op	     259 B/op	       6 allocs/op
// BenchmarkInvoke/SingleInvoker_Length_(slice_type)-12                                 	 2000000	       661 ns/op	     384 B/op	       9 allocs/op
// BenchmarkInvoke/CachedInvoker_Length_(slice_type)-12                                 	 3000000	       492 ns/op	     224 B/op	       5 allocs/op
// BenchmarkInvoke/SingleInvoker_Emitter-12                                             	 3000000	       412 ns/op	     184 B/op	       4 allocs/op
// BenchmarkInvoke/CachedInvoker_Emitter-12                                             	 5000000	       268 ns/op	      32 B/op	       1 allocs/op
// BenchmarkInvoke/SingleInvoker_Side_input-12                                          	 2000000	       743 ns/op	     392 B/op	      11 allocs/op
// BenchmarkInvoke/CachedInvoker_Side_input-12                                          	 3000000	       506 ns/op	     200 B/op	       6 allocs/op
// BenchmarkInvoke/SingleInvoker_Sum_as_Main-12                                         	 2000000	       810 ns/op	     464 B/op	      11 allocs/op
// BenchmarkInvoke/CachedInvoker_Sum_as_Main-12                                         	 3000000	       548 ns/op	     224 B/op	       5 allocs/op
// BenchmarkInvoke/SingleInvoker_Sum_as_Main_KV-12                                      	 2000000	       823 ns/op	     464 B/op	      11 allocs/op
// BenchmarkInvoke/CachedInvoker_Sum_as_Main_KV-12                                      	 3000000	       547 ns/op	     224 B/op	       5 allocs/op
// BenchmarkInvoke/SingleInvoker_EventTime-12                                           	 2000000	       711 ns/op	     376 B/op	      10 allocs/op
// BenchmarkInvoke/CachedInvoker_EventTime-12                                           	 3000000	       513 ns/op	     200 B/op	       6 allocs/op
// BenchmarkInvoke/SingleInvoker_Window-12                                              	 1000000	      1023 ns/op	     368 B/op	       9 allocs/op
// BenchmarkInvoke/CachedInvoker_Window-12                                              	 2000000	       838 ns/op	     192 B/op	       5 allocs/op

func BenchmarkInvoke(b *testing.B) {
	tests := []struct {
		Name     string
		Fn       interface{}
		Opt      *MainInput
		Args     []interface{}
		Expected interface{}
	}{
		{
			Name: "Void function",
			Fn:   func() {},
		},
		{
			Name: "Contexts and errors are allowed and handled",
			Fn:   func(ctx context.Context) error { return nil },
		},
		{
			Name:     "Sum",
			Fn:       func(a, b, c int) int { return a + b + c },
			Args:     []interface{}{1, 2, 3},
			Expected: 6,
		},
		{
			Name:     "Concat",
			Fn:       func(a, b, c string) string { return a + b + c },
			Args:     []interface{}{"a", "b", "c"},
			Expected: "abc",
		},
		{
			Name:     "Length (slice type)",
			Fn:       func(list []int) (int, error) { return len(list), nil },
			Args:     []interface{}{[]int{1, 2, 3}},
			Expected: 3,
		},
		{
			Name: "Emitter",
			Fn:   func(emit func(int)) { emit(1) },
			Args: []interface{}{func(int) {}},
		},
		{
			Name: "Side input",
			Fn: func(a int, get func(*int) bool) int {
				var ret int
				if !get(&ret) {
					return a
				}
				return ret
			},
			Args:     []interface{}{1, func(out *int) bool { *out = 2; return true }},
			Expected: 2,
		},
		{
			Name:     "Sum as Main",
			Fn:       func(a, b, c int) int { return a + b + c },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Args:     []interface{}{2, 3},
			Expected: 6,
		},
		{
			Name:     "Sum as Main KV",
			Fn:       func(a, b, c int) int { return a + b + c },
			Opt:      &MainInput{Key: FullValue{Elm: 1, Elm2: 2}},
			Args:     []interface{}{3},
			Expected: 6,
		},
		{
			Name:     "EventTime",
			Fn:       func(ts typex.EventTime, a int) int { return int(ts.Milliseconds()) + a },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Expected: 3,
		},
		{
			Name:     "Window",
			Fn:       func(w typex.Window, a int) int64 { return w.MaxTimestamp().Milliseconds() },
			Opt:      &MainInput{Key: FullValue{Elm: 1}},
			Expected: mtime.EndOfGlobalWindowTime.Milliseconds(),
		},
	}

	for _, test := range tests {
		fn, err := funcx.New(reflectx.MakeFunc(test.Fn))
		if err != nil {
			b.Fatalf("function not valid: %v", err)
		}

		ts := mtime.ZeroTimestamp.Add(2 * time.Millisecond)
		b.Run(fmt.Sprintf("SingleInvoker_%s", test.Name), func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				_, err := Invoke(context.Background(), window.SingleGlobalWindow, ts, fn, test.Opt, test.Args...)
				if err != nil {
					b.Fatalf("Invoke(%v,%v) failed: %v", fn.Fn.Name(), test.Args, err)
				}
			}
		})
		b.Run(fmt.Sprintf("CachedInvoker_%s", test.Name), func(b *testing.B) {
			inv := newInvoker(fn)
			for i := 0; i < b.N; i++ {
				_, err := inv.Invoke(context.Background(), window.SingleGlobalWindow, ts, test.Opt, test.Args...)
				if err != nil {
					b.Fatalf("Invoke(%v,%v) failed: %v", fn.Fn.Name(), test.Args, err)
				}
			}
		})
	}
}

// NOTE(herohde) 12/19/2017: example run on a laptop
//
// BenchmarkDirectCall-4              2000000000	         0.26 ns/op
// BenchmarkIndirectCall-4            1000000000	         2.28 ns/op
// BenchmarkReflectedAndBackCall-4    1000000000	         2.40 ns/op
// BenchmarkReflectCall-4            	10000000	       197    ns/op
// BenchmarkReflectCallNewArgs-4     	10000000	       200    ns/op
// BenchmarkReflectCallReuseArgs-4   	10000000	       196    ns/op
// BenchmarkInvokeCall-4             	 3000000	       452    ns/op
// BenchmarkInvokeCallExtra-4        	 3000000	       419    ns/op
// BenchmarkReflectFnCall-4          	 5000000	       340    ns/op
// BenchmarkInvokeFnCall-4           	 2000000	       604    ns/op
// BenchmarkInvokeFnCallExtra-4      	 3000000	       559    ns/op

func inc(n int) int {
	return n + 1
}

func BenchmarkDirectCall(b *testing.B) {
	n := 0
	for i := 0; i < b.N; i++ {
		n = inc(n)
	}
	b.Log(n)
}

func BenchmarkIndirectCall(b *testing.B) {
	fn := func(int) int { return 0 }
	if b.N > 0 {
		fn = inc // this is always set, but the compiler doesn't know
	}

	n := 0
	for i := 0; i < b.N; i++ {
		n = fn(n)
	}
	b.Log(n)
}

func BenchmarkReflectedAndBackCall(b *testing.B) {
	fn := reflect.ValueOf(inc).Interface().(func(int) int)
	n := 0
	for i := 0; i < b.N; i++ {
		n = fn(n)
	}
	b.Log(n)
}

func BenchmarkReflectCall(b *testing.B) {
	fn := reflect.ValueOf(inc)
	n := reflect.ValueOf(0)
	for i := 0; i < b.N; i++ {
		n = fn.Call([]reflect.Value{n})[0]
	}
	b.Log(n.Interface())
}

func BenchmarkReflectCallNewArgs(b *testing.B) {
	fn := reflect.ValueOf(inc)
	n := reflect.ValueOf(0)
	for i := 0; i < b.N; i++ {
		args := make([]reflect.Value, 1, 1)
		args[0] = n
		n = fn.Call(args)[0]
	}
	b.Log(n.Interface())
}

func BenchmarkReflectCallReuseArgs(b *testing.B) {
	fn := reflect.ValueOf(inc)
	n := reflect.ValueOf(0)
	args := make([]reflect.Value, 1, 1)
	for i := 0; i < b.N; i++ {
		args[0] = n
		n = fn.Call(args)[0]
	}
	b.Log(n.Interface())
}

func BenchmarkInvokeCall(b *testing.B) {
	fn, _ := funcx.New(reflectx.MakeFunc(inc))
	ctx := context.Background()
	n := 0
	for i := 0; i < b.N; i++ {
		ret, _ := InvokeWithoutEventTime(ctx, fn, &MainInput{Key: FullValue{Elm: n}})
		n = ret.Elm.(int)
	}
	b.Log(n)
}

func BenchmarkInvokeCallExtra(b *testing.B) {
	fn, _ := funcx.New(reflectx.MakeFunc(inc))
	ctx := context.Background()
	n := 0
	for i := 0; i < b.N; i++ {
		ret, _ := InvokeWithoutEventTime(ctx, fn, nil, n)
		n = ret.Elm.(int)
	}
	b.Log(n)
}

// The below take the additional overhead of MakeFunc.

func incFn(args []reflect.Value) []reflect.Value {
	return []reflect.Value{reflect.ValueOf(args[0].Interface().(int) + 1)}
}

func BenchmarkReflectFnCall(b *testing.B) {
	fn := reflect.MakeFunc(reflect.TypeOf(inc), incFn)
	n := reflect.ValueOf(0)
	for i := 0; i < b.N; i++ {
		n = fn.Call([]reflect.Value{n})[0]
	}
	b.Log(n.Interface())
}

func BenchmarkInvokeFnCall(b *testing.B) {
	fn, _ := funcx.New(reflectx.MakeFunc(reflect.MakeFunc(reflect.TypeOf(inc), incFn).Interface().(func(int) int)))
	ctx := context.Background()
	n := 0
	for i := 0; i < b.N; i++ {
		ret, _ := InvokeWithoutEventTime(ctx, fn, &MainInput{Key: FullValue{Elm: n}})
		n = ret.Elm.(int)
	}
	b.Log(n)
}

func BenchmarkInvokeFnCallExtra(b *testing.B) {
	fn, _ := funcx.New(reflectx.MakeFunc(reflect.MakeFunc(reflect.TypeOf(inc), incFn).Interface().(func(int) int)))
	ctx := context.Background()
	n := 0
	for i := 0; i < b.N; i++ {
		ret, _ := InvokeWithoutEventTime(ctx, fn, nil, n)
		n = ret.Elm.(int)
	}
	b.Log(n)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"regexp"

	"github.com/majormjr/rcon"
)

type FactorioServer struct {
	Cmd            *exec.Cmd              `json:"-"`
	Savefile       string                 `json:"savefile"`
	Latency        int                    `json:"latency"`
	BindIP         string                 `json:"bindip"`
	Port           int                    `json:"port"`
	Running        bool                   `json:"running"`
	Version        Version                `json:"fac_version"`
	BaseModVersion string                 `json:"base_mod_version"`
	StdOut         io.ReadCloser          `json:"-"`
	StdErr         io.ReadCloser          `json:"-"`
	StdIn          io.WriteCloser         `json:"-"`
	Settings       map[string]interface{} `json:"-"`
	Rcon           *rcon.RemoteConsole    `json:"-"`
	LogChan        chan []string          `json:"-"`
}

func randomPort() int {
	// Returns random port to use for rcon connection
	return rand.Intn(45000-40000) + 40000
}

func initFactorio() (f *FactorioServer, err error) {
	f = new(FactorioServer)
	f.Settings = make(map[string]interface{})

	if err = os.MkdirAll(config.FactorioConfigDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create config directory: %v", err)
	}

	settingsPath := filepath.Join(config.FactorioConfigDir, config.SettingsFile)
	var settings *os.File

	if _, err := os.Stat(settingsPath); os.IsNotExist(err) {
		// copy example settings to supplied settings file, if not exists
		log.Printf("Server settings at %s not found, copying example server settings.\n", settingsPath)

		examplePath := filepath.Join(config.FactorioDir, "data", "server-settings.example.json")

		example, err := os.Open(examplePath)
		if err != nil {
			return nil, fmt.Errorf("failed to open example server settings: %v", err)
		}
		defer example.Close()

		settings, err = os.Create(settingsPath)
		if err != nil {
			return nil, fmt.Errorf("failed to create server settings file: %v", err)
		}
		defer settings.Close()

		_, err = io.Copy(settings, example)
		if err != nil {
			return nil, fmt.Errorf("failed to copy example server settings: %v", err)
		}
	} else {
		// otherwise, open file normally
		settings, err = os.Open(settingsPath)
		if err != nil {
			return nil, fmt.Errorf("failed to open server settings file: %v", err)
		}
		defer settings.Close()
	}

	// before reading reset offset
	if _, err = settings.Seek(0, 0); err != nil {
		return nil, fmt.Errorf("error while seeking in settings file: %v", err)
	}

	if err = json.NewDecoder(settings).Decode(&f.Settings); err != nil {
		return nil, fmt.Errorf("error reading %s: %v", settingsPath, err)
	}

	log.Printf("Loaded Factorio settings from %s\n", settingsPath)

	//Load factorio version
	out, err := exec.Command(config.FactorioBinary, "--version").Output()
	if err != nil {
		log.Printf("error on loading factorio version: %s", err)
		return
	}

	reg := regexp.MustCompile("Version.*?((\\d+\\.)?(\\d+\\.)?(\\*|\\d+)+)")
	found := reg.FindStringSubmatch(string(out))
	err = f.Version.UnmarshalText([]byte(found[1]))
	if err != nil {
		log.Printf("could not parse version: %v", err)
		return
	}

	//Load baseMod version
	baseModInfoFile := filepath.Join(config.FactorioDir, "data", "base", "info.json")
	bmifBa, err := ioutil.ReadFile(baseModInfoFile)
	if err != nil {
		log.Printf("couldn't open baseMods info.json: %s", err)
		return
	}
	var modInfo ModInfo
	err = json.Unmarshal(bmifBa, &modInfo)
	if err != nil {
		log.Printf("error unmarshalling baseMods info.json to a modInfo: %s", err)
		return
	}

	f.BaseModVersion = modInfo.Version

	// load admins from additional file
	if(f.Version.Greater(Version{0,17,0})) {
		if _, err := os.Stat(filepath.Join(config.FactorioConfigDir, config.FactorioAdminFile)); os.IsNotExist(err) {
			//save empty admins-file
			ioutil.WriteFile(filepath.Join(config.FactorioConfigDir, config.FactorioAdminFile), []byte("[]"), 0664)
		} else {
			data, err := ioutil.ReadFile(filepath.Join(config.FactorioConfigDir, config.FactorioAdminFile))
			if err != nil {
				log.Printf("Error loading FactorioAdminFile: %s", err)
				return f, err
			}

			var jsonData interface{}
			err = json.Unmarshal(data, &jsonData)
			if err != nil {
				log.Printf("Error unmarshalling FactorioAdminFile: %s", err)
				return f, err
			}

			f.Settings["admins"] = jsonData
		}
	}

	return
}

func (f *FactorioServer) Run() error {
	var err error

	data, err := json.MarshalIndent(f.Settings, "", "  ")
	if err != nil {
		log.Println("Failed to marshal FactorioServerSettings: ", err)
	} else {
		ioutil.WriteFile(filepath.Join(config.FactorioConfigDir, config.SettingsFile), data, 0644)
	}

	args := []string{
		"--bind", (f.BindIP),
		"--port", strconv.Itoa(f.Port),
		"--server-settings", filepath.Join(config.FactorioConfigDir, config.SettingsFile),
		"--server-adminlist", filepath.Join(config.FactorioConfigDir, config.FactorioAdminFile),
		"--rcon-port", strconv.Itoa(config.FactorioRconPort),
		"--rcon-password", config.FactorioRconPass,
	}

	if f.Savefile == "Load Latest" {
		args = append(args, "--start-server-load-latest")
	} else {
		args = append(args, "--start-server", filepath.Join(config.FactorioSavesDir, f.Savefile))
	}

	log.Println("Starting server with command: ", config.FactorioBinary, args)

	f.Cmd = exec.Command(config.FactorioBinary, args...)

	f.StdOut, err = f.Cmd.StdoutPipe()
	if err != nil {
		log.Printf("Error opening stdout pipe: %s", err)
		return err
	}

	f.StdIn, err = f.Cmd.StdinPipe()
	if err != nil {
		log.Printf("Error opening stdin pipe: %s", err)
		return err
	}

	f.StdErr, err = f.Cmd.StderrPipe()
	if err != nil {
		log.Printf("Error opening stderr pipe: %s", err)
		return err
	}

	go f.parseRunningCommand(f.StdOut)
	go f.parseRunningCommand(f.StdErr)

	err = f.Cmd.Start()
	if err != nil {
		log.Printf("Factorio process failed to start: %s", err)
		return err
	}
	f.Running = true

	err = f.Cmd.Wait()
	if err != nil {
		log.Printf("Factorio process exited with error: %s", err)
		f.Running = false
		return err
	}

	return nil
}

func (f *FactorioServer) parseRunningCommand(std io.ReadCloser) (err error) {
	stdScanner := bufio.NewScanner(std)
	for stdScanner.Scan() {
		log.Printf("Factorio Server: %s", stdScanner.Text())
		if err := f.writeLog(stdScanner.Text()); err != nil {
			log.Printf("Error: %s", err)
		}

		line := strings.Fields(stdScanner.Text())
		// Ensure logline slice is in bounds
		if len(line) > 1 {
			// Check if Factorio Server reports any errors if so handle it
			if line[1] == "Error" {
				err := f.checkLogError(line)
				if err != nil {
					log.Printf("Error checking Factorio Server Error: %s", err)
				}
			}
			// If rcon port opens indicated in log connect to rcon
			rconLog := "Starting RCON interface at port " + strconv.Itoa(config.FactorioRconPort)
			// check if slice index is greater than 2 to prevent panic
			if len(line) > 2 {
				// log line for opened rcon connection
				if strings.Join(line[3:], " ") == rconLog {
					log.Printf("Rcon running on Factorio Server")
					err = connectRC()
					if err != nil {
						log.Printf("Error: %s", err)
					}
				}
			}
		}
	}
	if err := stdScanner.Err(); err != nil {
		log.Printf("Error reading std buffer: %s", err)
		return err
	}
	return nil
}

func (f *FactorioServer) writeLog(logline string) error {
	logfileName := config.FactorioDir + "factorio-server-console.log"
	file, err := os.OpenFile(logfileName, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		log.Printf("Cannot open logfile for appending Factorio Server output: %s", err)
		return err
	}
	defer file.Close()

	logline = logline + "\n"

	if _, err = file.WriteString(logline); err != nil {
		log.Printf("Error appending to factorio-server-console.log: %s", err)
		return err
	}

	return nil
}

func (f *FactorioServer) checkLogError(logline []string) error {
	// TODO Handle errors generated by running Factorio Server
	log.Println(logline)

	return nil
}

func (f *FactorioServer) Stop() error {
	if runtime.GOOS == "windows" {

		// Disable our own handling of CTRL+C, so we don't close when we send it to the console.
		setCtrlHandlingIsDisabledForThisProcess(true)

		// Send CTRL+C to all processes attached to the console (ourself, and the factorio server instance)
		sendCtrlCToPid(0)
		log.Println("Sent SIGINT to Factorio process. Factorio shutting down...")

		// Somehow, the Factorio devs managed to code the game to react appropriately to CTRL+C, including
		// saving the game, but not actually exit. So, we still have to manually kill the process, and
		// for extra fun, there's no way to know when the server save has actually completed (unless we want
		// to inject filesystem logic into what should be a process-level Stop() routine), so our best option
		// is to just wait an arbitrary amount of time and hope that the save is successful in that time.
		time.Sleep(2 * time.Second)
		f.Cmd.Process.Signal(os.Kill)

		// Re-enable handling of CTRL+C after we're sure that the factrio server is shut down.
		setCtrlHandlingIsDisabledForThisProcess(false)

		f.Running = false
		return nil
	}

	err := f.Cmd.Process.Signal(os.Interrupt)
	if err != nil {
		if err.Error() == "os: process already finished" {
			f.Running = false
			return err
		}
		log.Printf("Error sending SIGINT to Factorio process: %s", err)
		return err
	}
	f.Running = false
	log.Printf("Sent SIGINT to Factorio process. Factorio shutting down...")

	err = f.Rcon.Close()
	if err != nil {
		log.Printf("Error close rcon connection: %s", err)
	}

	return nil
}

func (f *FactorioServer) Kill() error {
	if runtime.GOOS == "windows" {

		err := f.Cmd.Process.Signal(os.Kill)
		if err != nil {
			if err.Error() == "os: process already finished" {
				f.Running = false
				return err
			}
			log.Printf("Error sending SIGKILL to Factorio process: %s", err)
			return err
		}
		f.Running = false
		log.Println("Sent SIGKILL to Factorio process. Factorio forced to exit.")

		return nil
	}

	err := f.Cmd.Process.Signal(os.Kill)
	if err != nil {
		if err.Error() == "os: process already finished" {
			f.Running = false
			return err
		}
		log.Printf("Error sending SIGKILL to Factorio process: %s", err)
		return err
	}
	f.Running = false
	log.Printf("Sent SIGKILL to Factorio process. Factorio forced to exit.")

	err = f.Rcon.Close()
	if err != nil {
		log.Printf("Error close rcon connection: %s", err)
	}

	return nil
}
<commit_msg>only append server-adminlist, when 0.17 or newer<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"regexp"

	"github.com/majormjr/rcon"
)

type FactorioServer struct {
	Cmd            *exec.Cmd              `json:"-"`
	Savefile       string                 `json:"savefile"`
	Latency        int                    `json:"latency"`
	BindIP         string                 `json:"bindip"`
	Port           int                    `json:"port"`
	Running        bool                   `json:"running"`
	Version        Version                `json:"fac_version"`
	BaseModVersion string                 `json:"base_mod_version"`
	StdOut         io.ReadCloser          `json:"-"`
	StdErr         io.ReadCloser          `json:"-"`
	StdIn          io.WriteCloser         `json:"-"`
	Settings       map[string]interface{} `json:"-"`
	Rcon           *rcon.RemoteConsole    `json:"-"`
	LogChan        chan []string          `json:"-"`
}

func randomPort() int {
	// Returns random port to use for rcon connection
	return rand.Intn(45000-40000) + 40000
}

func initFactorio() (f *FactorioServer, err error) {
	f = new(FactorioServer)
	f.Settings = make(map[string]interface{})

	if err = os.MkdirAll(config.FactorioConfigDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create config directory: %v", err)
	}

	settingsPath := filepath.Join(config.FactorioConfigDir, config.SettingsFile)
	var settings *os.File

	if _, err := os.Stat(settingsPath); os.IsNotExist(err) {
		// copy example settings to supplied settings file, if not exists
		log.Printf("Server settings at %s not found, copying example server settings.\n", settingsPath)

		examplePath := filepath.Join(config.FactorioDir, "data", "server-settings.example.json")

		example, err := os.Open(examplePath)
		if err != nil {
			return nil, fmt.Errorf("failed to open example server settings: %v", err)
		}
		defer example.Close()

		settings, err = os.Create(settingsPath)
		if err != nil {
			return nil, fmt.Errorf("failed to create server settings file: %v", err)
		}
		defer settings.Close()

		_, err = io.Copy(settings, example)
		if err != nil {
			return nil, fmt.Errorf("failed to copy example server settings: %v", err)
		}
	} else {
		// otherwise, open file normally
		settings, err = os.Open(settingsPath)
		if err != nil {
			return nil, fmt.Errorf("failed to open server settings file: %v", err)
		}
		defer settings.Close()
	}

	// before reading reset offset
	if _, err = settings.Seek(0, 0); err != nil {
		return nil, fmt.Errorf("error while seeking in settings file: %v", err)
	}

	if err = json.NewDecoder(settings).Decode(&f.Settings); err != nil {
		return nil, fmt.Errorf("error reading %s: %v", settingsPath, err)
	}

	log.Printf("Loaded Factorio settings from %s\n", settingsPath)

	//Load factorio version
	out, err := exec.Command(config.FactorioBinary, "--version").Output()
	if err != nil {
		log.Printf("error on loading factorio version: %s", err)
		return
	}

	reg := regexp.MustCompile("Version.*?((\\d+\\.)?(\\d+\\.)?(\\*|\\d+)+)")
	found := reg.FindStringSubmatch(string(out))
	err = f.Version.UnmarshalText([]byte(found[1]))
	if err != nil {
		log.Printf("could not parse version: %v", err)
		return
	}

	//Load baseMod version
	baseModInfoFile := filepath.Join(config.FactorioDir, "data", "base", "info.json")
	bmifBa, err := ioutil.ReadFile(baseModInfoFile)
	if err != nil {
		log.Printf("couldn't open baseMods info.json: %s", err)
		return
	}
	var modInfo ModInfo
	err = json.Unmarshal(bmifBa, &modInfo)
	if err != nil {
		log.Printf("error unmarshalling baseMods info.json to a modInfo: %s", err)
		return
	}

	f.BaseModVersion = modInfo.Version

	// load admins from additional file
	if(f.Version.Greater(Version{0,17,0})) {
		if _, err := os.Stat(filepath.Join(config.FactorioConfigDir, config.FactorioAdminFile)); os.IsNotExist(err) {
			//save empty admins-file
			ioutil.WriteFile(filepath.Join(config.FactorioConfigDir, config.FactorioAdminFile), []byte("[]"), 0664)
		} else {
			data, err := ioutil.ReadFile(filepath.Join(config.FactorioConfigDir, config.FactorioAdminFile))
			if err != nil {
				log.Printf("Error loading FactorioAdminFile: %s", err)
				return f, err
			}

			var jsonData interface{}
			err = json.Unmarshal(data, &jsonData)
			if err != nil {
				log.Printf("Error unmarshalling FactorioAdminFile: %s", err)
				return f, err
			}

			f.Settings["admins"] = jsonData
		}
	}

	return
}

func (f *FactorioServer) Run() error {
	var err error

	data, err := json.MarshalIndent(f.Settings, "", "  ")
	if err != nil {
		log.Println("Failed to marshal FactorioServerSettings: ", err)
	} else {
		ioutil.WriteFile(filepath.Join(config.FactorioConfigDir, config.SettingsFile), data, 0644)
	}

	args := []string{
		"--bind", (f.BindIP),
		"--port", strconv.Itoa(f.Port),
		"--server-settings", filepath.Join(config.FactorioConfigDir, config.SettingsFile),
		"--rcon-port", strconv.Itoa(config.FactorioRconPort),
		"--rcon-password", config.FactorioRconPass,
	}

	if(f.Version.Greater(Version{0,17,0})) {
		args = append(args, "--server-adminlist", filepath.Join(config.FactorioConfigDir, config.FactorioAdminFile))
	}

	if f.Savefile == "Load Latest" {
		args = append(args, "--start-server-load-latest")
	} else {
		args = append(args, "--start-server", filepath.Join(config.FactorioSavesDir, f.Savefile))
	}

	log.Println("Starting server with command: ", config.FactorioBinary, args)

	f.Cmd = exec.Command(config.FactorioBinary, args...)

	f.StdOut, err = f.Cmd.StdoutPipe()
	if err != nil {
		log.Printf("Error opening stdout pipe: %s", err)
		return err
	}

	f.StdIn, err = f.Cmd.StdinPipe()
	if err != nil {
		log.Printf("Error opening stdin pipe: %s", err)
		return err
	}

	f.StdErr, err = f.Cmd.StderrPipe()
	if err != nil {
		log.Printf("Error opening stderr pipe: %s", err)
		return err
	}

	go f.parseRunningCommand(f.StdOut)
	go f.parseRunningCommand(f.StdErr)

	err = f.Cmd.Start()
	if err != nil {
		log.Printf("Factorio process failed to start: %s", err)
		return err
	}
	f.Running = true

	err = f.Cmd.Wait()
	if err != nil {
		log.Printf("Factorio process exited with error: %s", err)
		f.Running = false
		return err
	}

	return nil
}

func (f *FactorioServer) parseRunningCommand(std io.ReadCloser) (err error) {
	stdScanner := bufio.NewScanner(std)
	for stdScanner.Scan() {
		log.Printf("Factorio Server: %s", stdScanner.Text())
		if err := f.writeLog(stdScanner.Text()); err != nil {
			log.Printf("Error: %s", err)
		}

		line := strings.Fields(stdScanner.Text())
		// Ensure logline slice is in bounds
		if len(line) > 1 {
			// Check if Factorio Server reports any errors if so handle it
			if line[1] == "Error" {
				err := f.checkLogError(line)
				if err != nil {
					log.Printf("Error checking Factorio Server Error: %s", err)
				}
			}
			// If rcon port opens indicated in log connect to rcon
			rconLog := "Starting RCON interface at port " + strconv.Itoa(config.FactorioRconPort)
			// check if slice index is greater than 2 to prevent panic
			if len(line) > 2 {
				// log line for opened rcon connection
				if strings.Join(line[3:], " ") == rconLog {
					log.Printf("Rcon running on Factorio Server")
					err = connectRC()
					if err != nil {
						log.Printf("Error: %s", err)
					}
				}
			}
		}
	}
	if err := stdScanner.Err(); err != nil {
		log.Printf("Error reading std buffer: %s", err)
		return err
	}
	return nil
}

func (f *FactorioServer) writeLog(logline string) error {
	logfileName := config.FactorioDir + "factorio-server-console.log"
	file, err := os.OpenFile(logfileName, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		log.Printf("Cannot open logfile for appending Factorio Server output: %s", err)
		return err
	}
	defer file.Close()

	logline = logline + "\n"

	if _, err = file.WriteString(logline); err != nil {
		log.Printf("Error appending to factorio-server-console.log: %s", err)
		return err
	}

	return nil
}

func (f *FactorioServer) checkLogError(logline []string) error {
	// TODO Handle errors generated by running Factorio Server
	log.Println(logline)

	return nil
}

func (f *FactorioServer) Stop() error {
	if runtime.GOOS == "windows" {

		// Disable our own handling of CTRL+C, so we don't close when we send it to the console.
		setCtrlHandlingIsDisabledForThisProcess(true)

		// Send CTRL+C to all processes attached to the console (ourself, and the factorio server instance)
		sendCtrlCToPid(0)
		log.Println("Sent SIGINT to Factorio process. Factorio shutting down...")

		// Somehow, the Factorio devs managed to code the game to react appropriately to CTRL+C, including
		// saving the game, but not actually exit. So, we still have to manually kill the process, and
		// for extra fun, there's no way to know when the server save has actually completed (unless we want
		// to inject filesystem logic into what should be a process-level Stop() routine), so our best option
		// is to just wait an arbitrary amount of time and hope that the save is successful in that time.
		time.Sleep(2 * time.Second)
		f.Cmd.Process.Signal(os.Kill)

		// Re-enable handling of CTRL+C after we're sure that the factrio server is shut down.
		setCtrlHandlingIsDisabledForThisProcess(false)

		f.Running = false
		return nil
	}

	err := f.Cmd.Process.Signal(os.Interrupt)
	if err != nil {
		if err.Error() == "os: process already finished" {
			f.Running = false
			return err
		}
		log.Printf("Error sending SIGINT to Factorio process: %s", err)
		return err
	}
	f.Running = false
	log.Printf("Sent SIGINT to Factorio process. Factorio shutting down...")

	err = f.Rcon.Close()
	if err != nil {
		log.Printf("Error close rcon connection: %s", err)
	}

	return nil
}

func (f *FactorioServer) Kill() error {
	if runtime.GOOS == "windows" {

		err := f.Cmd.Process.Signal(os.Kill)
		if err != nil {
			if err.Error() == "os: process already finished" {
				f.Running = false
				return err
			}
			log.Printf("Error sending SIGKILL to Factorio process: %s", err)
			return err
		}
		f.Running = false
		log.Println("Sent SIGKILL to Factorio process. Factorio forced to exit.")

		return nil
	}

	err := f.Cmd.Process.Signal(os.Kill)
	if err != nil {
		if err.Error() == "os: process already finished" {
			f.Running = false
			return err
		}
		log.Printf("Error sending SIGKILL to Factorio process: %s", err)
		return err
	}
	f.Running = false
	log.Printf("Sent SIGKILL to Factorio process. Factorio forced to exit.")

	err = f.Rcon.Close()
	if err != nil {
		log.Printf("Error close rcon connection: %s", err)
	}

	return nil
}
<|endoftext|>


<commit_before>package kumo

import (
	"reflect"
	"strings"
	"testing"
)

func Test_Parse(t *testing.T) {
	content := strings.NewReader(`
<?xml version="1.0" encoding="utf-8" ?>
<!DOCTYPE nzb PUBLIC "-//newzBin//DTD NZB 1.0//EN" "http://www.newzbin.com/DTD/nzb/nzb-1.0.dtd">
    <nzb xmlns="http://www.newzbin.com/DTD/2003/nzb">
        <file poster="poster" date="123" subject="subject">
            <groups>
                <group>alt.group</group>
            </groups>
            <segments>
                <segment bytes="11" number="1">1@foo.com</segment>
            </segments>
        </file>
    </nzb>
`)
	nzb, err := Parse(content)
	if err != nil {
		t.Fatalf("Parse error: %v", err)
	}

	want := NZB{
		Files: []File{
			{Poster: "poster", Subject: "subject", Date: 123, Groups: []string{"alt.group"}, Segments: []Segment{{Bytes: 11, Number: 1, Segment: "1@foo.com"}}}}}
	if !reflect.DeepEqual(nzb, &want) {
		t.Errorf("Returned %+v, want %+v", nzb, &want)
	}
}

func Test_Size(t *testing.T) {
	nzb := NZB{
		Files: []File{
			{Poster: "poster", Subject: "subject", Date: 123, Groups: []string{"alt.group"}, Segments: []Segment{{Bytes: 11, Number: 1, Segment: "1@foo.com"}}}}}

	want := 11
	if want != nzb.Size() {
		t.Errorf("Returned %+v, want %+v", nzb, &want)
	}
}
<commit_msg>Updated the expected value to an int64.<commit_after>package kumo

import (
	"reflect"
	"strings"
	"testing"
)

func Test_Parse(t *testing.T) {
	content := strings.NewReader(`
<?xml version="1.0" encoding="utf-8" ?>
<!DOCTYPE nzb PUBLIC "-//newzBin//DTD NZB 1.0//EN" "http://www.newzbin.com/DTD/nzb/nzb-1.0.dtd">
    <nzb xmlns="http://www.newzbin.com/DTD/2003/nzb">
        <file poster="poster" date="123" subject="subject">
            <groups>
                <group>alt.group</group>
            </groups>
            <segments>
                <segment bytes="11" number="1">1@foo.com</segment>
            </segments>
        </file>
    </nzb>
`)
	nzb, err := Parse(content)
	if err != nil {
		t.Fatalf("Parse error: %v", err)
	}

	want := NZB{
		Files: []File{
			{Poster: "poster", Subject: "subject", Date: 123, Groups: []string{"alt.group"}, Segments: []Segment{{Bytes: 11, Number: 1, Segment: "1@foo.com"}}}}}
	if !reflect.DeepEqual(nzb, &want) {
		t.Errorf("Returned %+v, want %+v", nzb, &want)
	}
}

func Test_Size(t *testing.T) {
	nzb := NZB{
		Files: []File{
			{Poster: "poster", Subject: "subject", Date: 123, Groups: []string{"alt.group"}, Segments: []Segment{{Bytes: 11, Number: 1, Segment: "1@foo.com"}}}}}

	want := int64(11)
	if want != nzb.Size() {
		t.Errorf("Returned %+v, want %+v", nzb, &want)
	}
}
<|endoftext|>


<commit_before>package lago_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/codemodus/lago"
)

func TestIntegNewFuncNilOpts(t *testing.T) {
	l := lago.New(nil)
	if l == nil {
		t.Errorf("don't want nil, got %T", l)
	}

	l.Println("test")
}

func TestIntegNewFuncEmptyOpts(t *testing.T) {
	l := lago.New(&lago.Options{})
	if l == nil {
		t.Errorf("don't want nil, got %T", l)
	}
}

func TestIntegNewFuncCompleteOpts(t *testing.T) {
	f, err := ioutil.TempFile("./", "test")
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		_ = f.Close()

		err := os.Remove(f.Name())
		if err != nil {
			fmt.Printf("WARN: tmp file not removed: %s\n", err)
		}
	}()

	l := lago.New(&lago.Options{
		Filepath: f.Name(),
	})

	if l == nil {
		t.Fatalf("don't want nil, got %T", l)
	}

	l.Printf("test")

	tf, err := os.Open(f.Name())
	if err != nil {
		t.Fatal(err)
	}

	b, err := ioutil.ReadAll(tf)
	if err != nil {
		t.Fatal(err)
	}

	if string(b)[:len(b)-1] != "test" {
		t.Errorf("want %s, got %s", "test", string(b))
	}
}
<commit_msg>Replace test string with var, and correct expectation output in integ test.<commit_after>package lago_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/codemodus/lago"
)

var (
	ts = "test"
)

func TestIntegNewFuncNilOpts(t *testing.T) {
	l := lago.New(nil)
	if l == nil {
		t.Errorf("don't want nil, got %T", l)
	}

	l.Println("test")
}

func TestIntegNewFuncEmptyOpts(t *testing.T) {
	l := lago.New(&lago.Options{})
	if l == nil {
		t.Errorf("don't want nil, got %T", l)
	}
}

func TestIntegNewFuncCompleteOpts(t *testing.T) {
	f, err := ioutil.TempFile("./", "test")
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		_ = f.Close()

		err := os.Remove(f.Name())
		if err != nil {
			fmt.Printf("WARN: tmp file not removed: %s\n", err)
		}
	}()

	l := lago.New(&lago.Options{
		Filepath: f.Name(),
	})

	if l == nil {
		t.Fatalf("don't want nil, got %T", l)
	}

	l.Printf(ts)

	tf, err := os.Open(f.Name())
	if err != nil {
		t.Fatal(err)
	}

	b, err := ioutil.ReadAll(tf)
	if err != nil {
		t.Fatal(err)
	}

	if string(b)[:len(b)-1] != ts {
		t.Errorf("want %s, got %s", ts, string(b)[:len(b)-1])
	}
}
<|endoftext|>


<commit_before>package ovmanager

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/libnetwork/datastore"
	"github.com/docker/libnetwork/discoverapi"
	"github.com/docker/libnetwork/driverapi"
	"github.com/docker/libnetwork/idm"
	"github.com/docker/libnetwork/netlabel"
	"github.com/docker/libnetwork/types"
)

const (
	networkType  = "overlay"
	vxlanIDStart = 4096
	vxlanIDEnd   = (1 << 24) - 1
)

type networkTable map[string]*network

type driver struct {
	config   map[string]interface{}
	networks networkTable
	store    datastore.DataStore
	vxlanIdm *idm.Idm
	sync.Mutex
}

type subnet struct {
	subnetIP *net.IPNet
	gwIP     *net.IPNet
	vni      uint32
}

type network struct {
	id      string
	driver  *driver
	subnets []*subnet
	sync.Mutex
}

// Init registers a new instance of overlay driver
func Init(dc driverapi.DriverCallback, config map[string]interface{}) error {
	var err error
	c := driverapi.Capability{
		DataScope: datastore.GlobalScope,
	}

	d := &driver{
		networks: networkTable{},
		config:   config,
	}

	d.vxlanIdm, err = idm.New(nil, "vxlan-id", 1, vxlanIDEnd)
	if err != nil {
		return fmt.Errorf("failed to initialize vxlan id manager: %v", err)
	}

	return dc.RegisterDriver(networkType, d, c)
}

func (d *driver) NetworkAllocate(id string, option map[string]string, ipV4Data, ipV6Data []driverapi.IPAMData) (map[string]string, error) {
	if id == "" {
		return nil, fmt.Errorf("invalid network id for overlay network")
	}

	if ipV4Data == nil {
		return nil, fmt.Errorf("empty ipv4 data passed during overlay network creation")
	}

	n := &network{
		id:      id,
		driver:  d,
		subnets: []*subnet{},
	}

	opts := make(map[string]string)
	vxlanIDList := make([]uint32, 0, len(ipV4Data))
	for key, val := range option {
		if key == netlabel.OverlayVxlanIDList {
			logrus.Debugf("overlay network option: %s", val)
			valStrList := strings.Split(val, ",")
			for _, idStr := range valStrList {
				vni, err := strconv.Atoi(idStr)
				if err != nil {
					return nil, fmt.Errorf("invalid vxlan id value %q passed", idStr)
				}

				vxlanIDList = append(vxlanIDList, uint32(vni))
			}
		} else {
			opts[key] = val
		}
	}

	for i, ipd := range ipV4Data {
		s := &subnet{
			subnetIP: ipd.Pool,
			gwIP:     ipd.Gateway,
		}

		if len(vxlanIDList) > i {
			s.vni = vxlanIDList[i]
		}

		if err := n.obtainVxlanID(s); err != nil {
			n.releaseVxlanID()
			return nil, fmt.Errorf("could not obtain vxlan id for pool %s: %v", s.subnetIP, err)
		}

		n.subnets = append(n.subnets, s)
	}

	val := fmt.Sprintf("%d", n.subnets[0].vni)
	for _, s := range n.subnets[1:] {
		val = val + fmt.Sprintf(",%d", s.vni)
	}
	opts[netlabel.OverlayVxlanIDList] = val

	d.Lock()
	d.networks[id] = n
	d.Unlock()

	return opts, nil
}

func (d *driver) NetworkFree(id string) error {
	if id == "" {
		return fmt.Errorf("invalid network id passed while freeing overlay network")
	}

	d.Lock()
	n, ok := d.networks[id]
	d.Unlock()

	if !ok {
		return fmt.Errorf("overlay network with id %s not found", id)
	}

	// Release all vxlan IDs in one shot.
	n.releaseVxlanID()

	d.Lock()
	delete(d.networks, id)
	d.Unlock()

	return nil
}

func (n *network) obtainVxlanID(s *subnet) error {
	var (
		err error
		vni uint64
	)

	n.Lock()
	vni = uint64(s.vni)
	n.Unlock()

	if vni == 0 {
		vni, err = n.driver.vxlanIdm.GetIDInRange(vxlanIDStart, vxlanIDEnd)
		if err != nil {
			return err
		}

		n.Lock()
		s.vni = uint32(vni)
		n.Unlock()
		return nil
	}

	return n.driver.vxlanIdm.GetSpecificID(vni)
}

func (n *network) releaseVxlanID() {
	n.Lock()
	vnis := make([]uint32, 0, len(n.subnets))
	for _, s := range n.subnets {
		vnis = append(vnis, s.vni)
		s.vni = 0
	}
	n.Unlock()

	for _, vni := range vnis {
		n.driver.vxlanIdm.Release(uint64(vni))
	}
}

func (d *driver) CreateNetwork(id string, option map[string]interface{}, nInfo driverapi.NetworkInfo, ipV4Data, ipV6Data []driverapi.IPAMData) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) EventNotify(etype driverapi.EventType, nid, tableName, key string, value []byte) {
}

func (d *driver) DeleteNetwork(nid string) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) CreateEndpoint(nid, eid string, ifInfo driverapi.InterfaceInfo, epOptions map[string]interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) DeleteEndpoint(nid, eid string) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) EndpointOperInfo(nid, eid string) (map[string]interface{}, error) {
	return nil, types.NotImplementedErrorf("not implemented")
}

// Join method is invoked when a Sandbox is attached to an endpoint.
func (d *driver) Join(nid, eid string, sboxKey string, jinfo driverapi.JoinInfo, options map[string]interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

// Leave method is invoked when a Sandbox detaches from an endpoint.
func (d *driver) Leave(nid, eid string) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) Type() string {
	return networkType
}

// DiscoverNew is a notification for a new discovery event, such as a new node joining a cluster
func (d *driver) DiscoverNew(dType discoverapi.DiscoveryType, data interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

// DiscoverDelete is a notification for a discovery delete event, such as a node leaving a cluster
func (d *driver) DiscoverDelete(dType discoverapi.DiscoveryType, data interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) ProgramExternalConnectivity(nid, eid string, options map[string]interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) RevokeExternalConnectivity(nid, eid string) error {
	return types.NotImplementedErrorf("not implemented")
}
<commit_msg>Create vxlan-id space from 0 instead of starting from 1<commit_after>package ovmanager

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/libnetwork/datastore"
	"github.com/docker/libnetwork/discoverapi"
	"github.com/docker/libnetwork/driverapi"
	"github.com/docker/libnetwork/idm"
	"github.com/docker/libnetwork/netlabel"
	"github.com/docker/libnetwork/types"
)

const (
	networkType  = "overlay"
	vxlanIDStart = 4096
	vxlanIDEnd   = (1 << 24) - 1
)

type networkTable map[string]*network

type driver struct {
	config   map[string]interface{}
	networks networkTable
	store    datastore.DataStore
	vxlanIdm *idm.Idm
	sync.Mutex
}

type subnet struct {
	subnetIP *net.IPNet
	gwIP     *net.IPNet
	vni      uint32
}

type network struct {
	id      string
	driver  *driver
	subnets []*subnet
	sync.Mutex
}

// Init registers a new instance of overlay driver
func Init(dc driverapi.DriverCallback, config map[string]interface{}) error {
	var err error
	c := driverapi.Capability{
		DataScope: datastore.GlobalScope,
	}

	d := &driver{
		networks: networkTable{},
		config:   config,
	}

	d.vxlanIdm, err = idm.New(nil, "vxlan-id", 0, vxlanIDEnd)
	if err != nil {
		return fmt.Errorf("failed to initialize vxlan id manager: %v", err)
	}

	return dc.RegisterDriver(networkType, d, c)
}

func (d *driver) NetworkAllocate(id string, option map[string]string, ipV4Data, ipV6Data []driverapi.IPAMData) (map[string]string, error) {
	if id == "" {
		return nil, fmt.Errorf("invalid network id for overlay network")
	}

	if ipV4Data == nil {
		return nil, fmt.Errorf("empty ipv4 data passed during overlay network creation")
	}

	n := &network{
		id:      id,
		driver:  d,
		subnets: []*subnet{},
	}

	opts := make(map[string]string)
	vxlanIDList := make([]uint32, 0, len(ipV4Data))
	for key, val := range option {
		if key == netlabel.OverlayVxlanIDList {
			logrus.Debugf("overlay network option: %s", val)
			valStrList := strings.Split(val, ",")
			for _, idStr := range valStrList {
				vni, err := strconv.Atoi(idStr)
				if err != nil {
					return nil, fmt.Errorf("invalid vxlan id value %q passed", idStr)
				}

				vxlanIDList = append(vxlanIDList, uint32(vni))
			}
		} else {
			opts[key] = val
		}
	}

	for i, ipd := range ipV4Data {
		s := &subnet{
			subnetIP: ipd.Pool,
			gwIP:     ipd.Gateway,
		}

		if len(vxlanIDList) > i {
			s.vni = vxlanIDList[i]
		}

		if err := n.obtainVxlanID(s); err != nil {
			n.releaseVxlanID()
			return nil, fmt.Errorf("could not obtain vxlan id for pool %s: %v", s.subnetIP, err)
		}

		n.subnets = append(n.subnets, s)
	}

	val := fmt.Sprintf("%d", n.subnets[0].vni)
	for _, s := range n.subnets[1:] {
		val = val + fmt.Sprintf(",%d", s.vni)
	}
	opts[netlabel.OverlayVxlanIDList] = val

	d.Lock()
	d.networks[id] = n
	d.Unlock()

	return opts, nil
}

func (d *driver) NetworkFree(id string) error {
	if id == "" {
		return fmt.Errorf("invalid network id passed while freeing overlay network")
	}

	d.Lock()
	n, ok := d.networks[id]
	d.Unlock()

	if !ok {
		return fmt.Errorf("overlay network with id %s not found", id)
	}

	// Release all vxlan IDs in one shot.
	n.releaseVxlanID()

	d.Lock()
	delete(d.networks, id)
	d.Unlock()

	return nil
}

func (n *network) obtainVxlanID(s *subnet) error {
	var (
		err error
		vni uint64
	)

	n.Lock()
	vni = uint64(s.vni)
	n.Unlock()

	if vni == 0 {
		vni, err = n.driver.vxlanIdm.GetIDInRange(vxlanIDStart, vxlanIDEnd)
		if err != nil {
			return err
		}

		n.Lock()
		s.vni = uint32(vni)
		n.Unlock()
		return nil
	}

	return n.driver.vxlanIdm.GetSpecificID(vni)
}

func (n *network) releaseVxlanID() {
	n.Lock()
	vnis := make([]uint32, 0, len(n.subnets))
	for _, s := range n.subnets {
		vnis = append(vnis, s.vni)
		s.vni = 0
	}
	n.Unlock()

	for _, vni := range vnis {
		n.driver.vxlanIdm.Release(uint64(vni))
	}
}

func (d *driver) CreateNetwork(id string, option map[string]interface{}, nInfo driverapi.NetworkInfo, ipV4Data, ipV6Data []driverapi.IPAMData) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) EventNotify(etype driverapi.EventType, nid, tableName, key string, value []byte) {
}

func (d *driver) DeleteNetwork(nid string) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) CreateEndpoint(nid, eid string, ifInfo driverapi.InterfaceInfo, epOptions map[string]interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) DeleteEndpoint(nid, eid string) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) EndpointOperInfo(nid, eid string) (map[string]interface{}, error) {
	return nil, types.NotImplementedErrorf("not implemented")
}

// Join method is invoked when a Sandbox is attached to an endpoint.
func (d *driver) Join(nid, eid string, sboxKey string, jinfo driverapi.JoinInfo, options map[string]interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

// Leave method is invoked when a Sandbox detaches from an endpoint.
func (d *driver) Leave(nid, eid string) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) Type() string {
	return networkType
}

// DiscoverNew is a notification for a new discovery event, such as a new node joining a cluster
func (d *driver) DiscoverNew(dType discoverapi.DiscoveryType, data interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

// DiscoverDelete is a notification for a discovery delete event, such as a node leaving a cluster
func (d *driver) DiscoverDelete(dType discoverapi.DiscoveryType, data interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) ProgramExternalConnectivity(nid, eid string, options map[string]interface{}) error {
	return types.NotImplementedErrorf("not implemented")
}

func (d *driver) RevokeExternalConnectivity(nid, eid string) error {
	return types.NotImplementedErrorf("not implemented")
}
<|endoftext|>


<commit_before>// Copyright 2015 tsuru-admin authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/cmd/cmdtest"
	"github.com/tsuru/tsuru/iaas"
	"gopkg.in/check.v1"
)

func (s *S) TestMachineListInfo(c *check.C) {
	expected := &cmd.Info{
		Name:    "machine-list",
		Usage:   "machine-list",
		Desc:    "List all machines created using a IaaS.",
		MinArgs: 0,
	}
	c.Assert((&machineList{}).Info(), check.DeepEquals, expected)
}

func (s *S) TestMachineListRun(c *check.C) {
	var stdout, stderr bytes.Buffer
	context := cmd.Context{
		Stdout: &stdout,
		Stderr: &stderr,
	}
	m1 := iaas.Machine{Id: "id1", Address: "addr1", Iaas: "iaas1", CreationParams: map[string]string{
		"param1": "value1",
	}}
	m2 := iaas.Machine{Id: "id2", Address: "addr2", Iaas: "iaas2", CreationParams: map[string]string{
		"param1": "value1",
		"param2": "value2",
	}}
	data, err := json.Marshal([]iaas.Machine{m1, m2})
	c.Assert(err, check.IsNil)
	expected := `+-----+-------+---------+-----------------+
| Id  | IaaS  | Address | Creation Params |
+-----+-------+---------+-----------------+
| id1 | iaas1 | addr1   | param1=value1   |
+-----+-------+---------+-----------------+
| id2 | iaas2 | addr2   | param1=value1   |
|     |       |         | param2=value2   |
+-----+-------+---------+-----------------+
`
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: string(data), Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			return req.URL.Path == "/iaas/machines" && req.Method == "GET"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, manager)
	command := machineList{}
	err = command.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(stdout.String(), check.Equals, expected)
}

func (s *S) TestMachineDestroyInfo(c *check.C) {
	expected := &cmd.Info{
		Name:    "machine-destroy",
		Usage:   "machine-destroy <machine id>",
		Desc:    "Destroy an existing machine created using a IaaS.",
		MinArgs: 1,
	}
	c.Assert((&machineDestroy{}).Info(), check.DeepEquals, expected)
}

func (s *S) TestMachineDestroyRun(c *check.C) {
	var stdout, stderr bytes.Buffer
	context := cmd.Context{
		Stdout: &stdout,
		Stderr: &stderr,
		Args:   []string{"myid1"},
	}
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: "", Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			return req.URL.Path == "/iaas/machines/myid1" && req.Method == "DELETE"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, manager)
	command := machineDestroy{}
	err := command.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(stdout.String(), check.Equals, "Machine successfully destroyed.\n")
}

func (s *S) TestTemplateListInfo(c *check.C) {
	expected := &cmd.Info{
		Name:    "machine-template-list",
		Usage:   "machine-template-list",
		Desc:    "List all machine templates.",
		MinArgs: 0,
	}
	c.Assert((&templateList{}).Info(), check.DeepEquals, expected)
}

func (s *S) TestTemplateListRun(c *check.C) {
	var stdout, stderr bytes.Buffer
	context := cmd.Context{
		Stdout: &stdout,
		Stderr: &stderr,
	}
	tpl1 := iaas.Template{Name: "machine1", IaaSName: "ec2", Data: iaas.TemplateDataList{
		{Name: "region", Value: "us-east-1"},
		{Name: "type", Value: "m1.small"},
	}}
	tpl2 := iaas.Template{Name: "tpl1", IaaSName: "ec2", Data: iaas.TemplateDataList{
		{Name: "region", Value: "xxxx"},
		{Name: "type", Value: "l1.large"},
	}}
	data, err := json.Marshal([]iaas.Template{tpl1, tpl2})
	c.Assert(err, check.IsNil)
	expected := `+----------+------+------------------+
| Name     | IaaS | Params           |
+----------+------+------------------+
| machine1 | ec2  | region=us-east-1 |
|          |      | type=m1.small    |
+----------+------+------------------+
| tpl1     | ec2  | region=xxxx      |
|          |      | type=l1.large    |
+----------+------+------------------+
`
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: string(data), Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			return req.URL.Path == "/iaas/templates" && req.Method == "GET"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, manager)
	command := templateList{}
	err = command.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(stdout.String(), check.Equals, expected)
}

func (s *S) TestTemplateAddCmdInfo(c *check.C) {
	expected := cmd.Info{
		Name:    "machine-template-add",
		Usage:   "machine-template-add <name> <iaas> <param>=<value>...",
		Desc:    "Add a new machine template.",
		MinArgs: 3,
	}
	cmd := templateAdd{}
	c.Assert(cmd.Info(), check.DeepEquals, &expected)
}

func (s *S) TestTemplateAddCmdRun(c *check.C) {
	var buf bytes.Buffer
	context := cmd.Context{Args: []string{"my-tpl", "ec2", "zone=xyz", "image=ami-something"}, Stdout: &buf}
	expectedBody := `{"Name":"my-tpl","IaaSName":"ec2",` +
		`"Data":[{"Name":"zone","Value":"xyz"},{"Name":"image","Value":"ami-something"}]}`
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: "", Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			body, _ := ioutil.ReadAll(req.Body)
			c.Assert(string(body), check.DeepEquals, expectedBody)
			return req.URL.Path == "/iaas/templates" && req.Method == "POST"
		},
	}
	manager := cmd.Manager{}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, &manager)
	cmd := templateAdd{}
	err := cmd.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(buf.String(), check.Equals, "Template successfully added.\n")
}

func (s *S) TestTemplateRemoveCmdInfo(c *check.C) {
	expected := cmd.Info{
		Name:    "machine-template-remove",
		Usage:   "machine-template-remove <name>",
		Desc:    "Remove an existing machine template.",
		MinArgs: 1,
	}
	cmd := templateRemove{}
	c.Assert(cmd.Info(), check.DeepEquals, &expected)
}

func (s *S) TestTemplateRemoveCmdRun(c *check.C) {
	var buf bytes.Buffer
	context := cmd.Context{Args: []string{"my-tpl"}, Stdout: &buf, Stderr: &buf}
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: "", Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			return req.URL.Path == "/iaas/templates/my-tpl" && req.Method == "DELETE"
		},
	}
	manager := cmd.Manager{}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, &manager)
	cmd := templateRemove{}
	err := cmd.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(buf.String(), check.Equals, "Template successfully removed.\n")
}

func (s *S) TestTemplateUpdateCmdRun(c *check.C) {
	var buf bytes.Buffer
	context := cmd.Context{Args: []string{"my-tpl", "zone=", "image=ami-something"}, Stdout: &buf}
	expectedBody := `{"Name":"my-tpl","IaaSName":"",` +
		`"Data":[{"Name":"zone","Value":""},{"Name":"image","Value":"ami-something"}]}`
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: "", Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			body, _ := ioutil.ReadAll(req.Body)
			c.Assert(string(body), check.DeepEquals, expectedBody)
			return req.URL.Path == "/iaas/templates/my-tpl" && req.Method == "PUT"
		},
	}
	manager := cmd.Manager{}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, &manager)
	cmd := templateUpdate{}
	err := cmd.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(buf.String(), check.Equals, "Template successfully updated.\n")
}
<commit_msg>iaas: remove tests that only copy implementation<commit_after>// Copyright 2015 tsuru-admin authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/cmd/cmdtest"
	"github.com/tsuru/tsuru/iaas"
	"gopkg.in/check.v1"
)

func (s *S) TestMachineListRun(c *check.C) {
	var stdout, stderr bytes.Buffer
	context := cmd.Context{
		Stdout: &stdout,
		Stderr: &stderr,
	}
	m1 := iaas.Machine{Id: "id1", Address: "addr1", Iaas: "iaas1", CreationParams: map[string]string{
		"param1": "value1",
	}}
	m2 := iaas.Machine{Id: "id2", Address: "addr2", Iaas: "iaas2", CreationParams: map[string]string{
		"param1": "value1",
		"param2": "value2",
	}}
	data, err := json.Marshal([]iaas.Machine{m1, m2})
	c.Assert(err, check.IsNil)
	expected := `+-----+-------+---------+-----------------+
| Id  | IaaS  | Address | Creation Params |
+-----+-------+---------+-----------------+
| id1 | iaas1 | addr1   | param1=value1   |
+-----+-------+---------+-----------------+
| id2 | iaas2 | addr2   | param1=value1   |
|     |       |         | param2=value2   |
+-----+-------+---------+-----------------+
`
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: string(data), Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			return req.URL.Path == "/iaas/machines" && req.Method == "GET"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, manager)
	command := machineList{}
	err = command.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(stdout.String(), check.Equals, expected)
}

func (s *S) TestMachineDestroyRun(c *check.C) {
	var stdout, stderr bytes.Buffer
	context := cmd.Context{
		Stdout: &stdout,
		Stderr: &stderr,
		Args:   []string{"myid1"},
	}
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: "", Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			return req.URL.Path == "/iaas/machines/myid1" && req.Method == "DELETE"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, manager)
	command := machineDestroy{}
	err := command.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(stdout.String(), check.Equals, "Machine successfully destroyed.\n")
}

func (s *S) TestTemplateListRun(c *check.C) {
	var stdout, stderr bytes.Buffer
	context := cmd.Context{
		Stdout: &stdout,
		Stderr: &stderr,
	}
	tpl1 := iaas.Template{Name: "machine1", IaaSName: "ec2", Data: iaas.TemplateDataList{
		{Name: "region", Value: "us-east-1"},
		{Name: "type", Value: "m1.small"},
	}}
	tpl2 := iaas.Template{Name: "tpl1", IaaSName: "ec2", Data: iaas.TemplateDataList{
		{Name: "region", Value: "xxxx"},
		{Name: "type", Value: "l1.large"},
	}}
	data, err := json.Marshal([]iaas.Template{tpl1, tpl2})
	c.Assert(err, check.IsNil)
	expected := `+----------+------+------------------+
| Name     | IaaS | Params           |
+----------+------+------------------+
| machine1 | ec2  | region=us-east-1 |
|          |      | type=m1.small    |
+----------+------+------------------+
| tpl1     | ec2  | region=xxxx      |
|          |      | type=l1.large    |
+----------+------+------------------+
`
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: string(data), Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			return req.URL.Path == "/iaas/templates" && req.Method == "GET"
		},
	}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, manager)
	command := templateList{}
	err = command.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(stdout.String(), check.Equals, expected)
}

func (s *S) TestTemplateAddCmdRun(c *check.C) {
	var buf bytes.Buffer
	context := cmd.Context{Args: []string{"my-tpl", "ec2", "zone=xyz", "image=ami-something"}, Stdout: &buf}
	expectedBody := `{"Name":"my-tpl","IaaSName":"ec2",` +
		`"Data":[{"Name":"zone","Value":"xyz"},{"Name":"image","Value":"ami-something"}]}`
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: "", Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			body, _ := ioutil.ReadAll(req.Body)
			c.Assert(string(body), check.DeepEquals, expectedBody)
			return req.URL.Path == "/iaas/templates" && req.Method == "POST"
		},
	}
	manager := cmd.Manager{}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, &manager)
	cmd := templateAdd{}
	err := cmd.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(buf.String(), check.Equals, "Template successfully added.\n")
}

func (s *S) TestTemplateRemoveCmdRun(c *check.C) {
	var buf bytes.Buffer
	context := cmd.Context{Args: []string{"my-tpl"}, Stdout: &buf, Stderr: &buf}
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: "", Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			return req.URL.Path == "/iaas/templates/my-tpl" && req.Method == "DELETE"
		},
	}
	manager := cmd.Manager{}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, &manager)
	cmd := templateRemove{}
	err := cmd.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(buf.String(), check.Equals, "Template successfully removed.\n")
}

func (s *S) TestTemplateUpdateCmdRun(c *check.C) {
	var buf bytes.Buffer
	context := cmd.Context{Args: []string{"my-tpl", "zone=", "image=ami-something"}, Stdout: &buf}
	expectedBody := `{"Name":"my-tpl","IaaSName":"",` +
		`"Data":[{"Name":"zone","Value":""},{"Name":"image","Value":"ami-something"}]}`
	trans := &cmdtest.ConditionalTransport{
		Transport: cmdtest.Transport{Message: "", Status: http.StatusOK},
		CondFunc: func(req *http.Request) bool {
			body, _ := ioutil.ReadAll(req.Body)
			c.Assert(string(body), check.DeepEquals, expectedBody)
			return req.URL.Path == "/iaas/templates/my-tpl" && req.Method == "PUT"
		},
	}
	manager := cmd.Manager{}
	client := cmd.NewClient(&http.Client{Transport: trans}, nil, &manager)
	cmd := templateUpdate{}
	err := cmd.Run(&context, client)
	c.Assert(err, check.IsNil)
	c.Assert(buf.String(), check.Equals, "Template successfully updated.\n")
}
<|endoftext|>


<commit_before>package iffy

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"

	"testing"
	"text/template"
)

type Tester struct {
	t      *testing.T
	r      http.Handler
	Calls  []*Call
	values Values
}

type Headers map[string]string

type Call struct {
	Name       string
	Method     string
	QueryStr   string
	Body       string
	headers    Headers
	respObject interface{}
	checkers   []Checker
}

func (c *Call) ResponseObject(respObject interface{}) *Call {
	c.respObject = respObject
	return c
}

func (c *Call) Headers(h Headers) *Call {
	c.headers = h
	return c
}

func (c *Call) Checkers(ch ...Checker) *Call {
	c.checkers = ch
	return c
}

type Checker func(r *http.Response, body string, respObject interface{}) error

// Tester

func NewTester(t *testing.T, r http.Handler, calls ...*Call) *Tester {
	return &Tester{
		t:      t,
		r:      r,
		values: make(Values),
	}
}

func (t *Tester) AddCall(name, method, querystr, body string) *Call {
	c := &Call{
		Name:     name,
		Method:   method,
		QueryStr: querystr,
		Body:     body,
	}
	t.Calls = append(t.Calls, c)
	return c
}

func (t *Tester) Run() {
	for _, c := range t.Calls {
		var body io.Reader
		if c.Body != "" {
			body = bytes.NewBuffer([]byte(t.applyTemplate(c.Body)))
		}
		req, err := http.NewRequest(c.Method, t.applyTemplate(c.QueryStr), body)
		if err != nil {
			t.t.Error(err)
			continue
		}
		if c.Body != "" {
			req.Header.Set("content-type", "application/json")
		}
		if c.headers != nil {
			for k, v := range c.headers {
				req.Header.Set(t.applyTemplate(k), t.applyTemplate(v))
			}
		}
		w := httptest.NewRecorder()
		t.r.ServeHTTP(w, req)
		resp := w.Result()
		var respBody string
		if resp.Body != nil {
			rb, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				t.t.Error(err)
			}
			respBody = string(rb)
			resp.Body.Close()
			if c.respObject != nil {
				err = json.Unmarshal(rb, c.respObject)
				if err != nil {
					t.t.Error(err)
				}
			}
			var retJson map[string]interface{}
			err = json.Unmarshal(rb, &retJson)
			if err == nil {
				t.values[c.Name] = retJson
			}
		}
		for _, checker := range c.checkers {
			err = checker(resp, respBody, c.respObject)
			if err != nil {
				t.t.Errorf("%s: %s", c.Name, err)
			}
		}
	}
}

func (t *Tester) applyTemplate(s string) string {
	b, err := t.values.Apply(s)
	if err != nil {
		t.t.Error(err)
		return ""
	}
	return string(b)
}

type Values map[string]interface{}

func (v Values) Apply(templateStr string) ([]byte, error) {

	var funcMap = template.FuncMap{
		"field": v.fieldTmpl,
		"json":  v.jsonFieldTmpl,
	}

	tmpl, err := template.New("tmpl").Funcs(funcMap).Parse(templateStr)
	if err != nil {
		return nil, err
	}

	b := new(bytes.Buffer)

	err = tmpl.Execute(b, v)
	if err != nil {
		return nil, err
	}

	return b.Bytes(), nil
}

// templating funcs

func (v Values) fieldTmpl(key ...string) (interface{}, error) {
	var i interface{}

	i = map[string]interface{}(v)
	var ok bool

	for _, k := range key {
		switch i.(type) {
		case map[string]interface{}:
			i, ok = i.(map[string]interface{})[k]
			if !ok {
				i = "<no value>"
			}
		case map[string]string:
			i, ok = i.(map[string]string)[k]
			if !ok {
				i = "<no value>"
			}
		default:
			return nil, fmt.Errorf("cannot dereference %T", i)
		}
	}
	return i, nil
}

func (v Values) jsonFieldTmpl(key ...string) (interface{}, error) {
	i, err := v.fieldTmpl(key...)
	if err != nil {
		return nil, err
	}
	marshalled, err := json.Marshal(i)
	if err != nil {
		return nil, err
	}
	return string(marshalled), nil
}

// BUILT IN CHECKERS

func ExpectStatus(st int) Checker {
	return func(r *http.Response, body string, respObject interface{}) error {
		if r.StatusCode != st {
			return fmt.Errorf("Bad status code: expected %d, got %d", st, r.StatusCode)
		}
		return nil
	}
}

func ExpectJSONFields(fields ...string) Checker {
	return func(r *http.Response, body string, respObject interface{}) error {
		m := map[string]interface{}{}
		err := json.Unmarshal([]byte(body), &m)
		if err != nil {
			return err
		}
		for _, f := range fields {
			if _, ok := m[f]; !ok {
				return fmt.Errorf("Missing expected field '%s'", f)
			}
		}
		return nil
	}
}

func ExpectListLength(length int) Checker {
	return func(r *http.Response, body string, respObject interface{}) error {
		l := []interface{}{}
		err := json.Unmarshal([]byte(body), &l)
		if err != nil {
			return err
		}
		if len(l) != length {
			return fmt.Errorf("Expected a list of length %d, got %d", length, len(l))
		}
		return nil
	}
}

func ExpectListNonEmpty(r *http.Response, body string, respObject interface{}) error {
	l := []interface{}{}
	err := json.Unmarshal([]byte(body), &l)
	if err != nil {
		return err
	}
	if len(l) == 0 {
		return errors.New("Expected a non empty list")
	}
	return nil
}

func ExpectJSONBranch(nodes ...string) Checker {
	return func(r *http.Response, body string, respObject interface{}) error {
		m := map[string]interface{}{}
		err := json.Unmarshal([]byte(body), &m)
		if err != nil {
			return err
		}
		for i, n := range nodes {
			v, ok := m[n]
			if !ok {
				return fmt.Errorf("Missing node '%s'", n)
			}
			if child, ok := v.(map[string]interface{}); ok {
				m = child
			} else if i == len(nodes)-2 {
				// last child is not an object anymore
				// and there's only one more node to check
				// test last child against last provided node
				lastNode := nodes[i+1]
				if v != lastNode {
					return fmt.Errorf("Wrong value: expected '%v', got '%v'", lastNode, v)
				}
				return nil
			}
		}
		return nil
	}
}
<commit_msg>support boolean and number comparison<commit_after>package iffy

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"

	"testing"
	"text/template"
)

type Tester struct {
	t      *testing.T
	r      http.Handler
	Calls  []*Call
	values Values
}

type Headers map[string]string

type Call struct {
	Name       string
	Method     string
	QueryStr   string
	Body       string
	headers    Headers
	respObject interface{}
	checkers   []Checker
}

func (c *Call) ResponseObject(respObject interface{}) *Call {
	c.respObject = respObject
	return c
}

func (c *Call) Headers(h Headers) *Call {
	c.headers = h
	return c
}

func (c *Call) Checkers(ch ...Checker) *Call {
	c.checkers = ch
	return c
}

type Checker func(r *http.Response, body string, respObject interface{}) error

// Tester

func NewTester(t *testing.T, r http.Handler, calls ...*Call) *Tester {
	return &Tester{
		t:      t,
		r:      r,
		values: make(Values),
	}
}

func (t *Tester) AddCall(name, method, querystr, body string) *Call {
	c := &Call{
		Name:     name,
		Method:   method,
		QueryStr: querystr,
		Body:     body,
	}
	t.Calls = append(t.Calls, c)
	return c
}

func (t *Tester) Run() {
	for _, c := range t.Calls {
		var body io.Reader
		if c.Body != "" {
			body = bytes.NewBuffer([]byte(t.applyTemplate(c.Body)))
		}
		req, err := http.NewRequest(c.Method, t.applyTemplate(c.QueryStr), body)
		if err != nil {
			t.t.Error(err)
			continue
		}
		if c.Body != "" {
			req.Header.Set("content-type", "application/json")
		}
		if c.headers != nil {
			for k, v := range c.headers {
				req.Header.Set(t.applyTemplate(k), t.applyTemplate(v))
			}
		}
		w := httptest.NewRecorder()
		t.r.ServeHTTP(w, req)
		resp := w.Result()
		var respBody string
		if resp.Body != nil {
			rb, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				t.t.Error(err)
			}
			respBody = string(rb)
			resp.Body.Close()
			if c.respObject != nil {
				err = json.Unmarshal(rb, c.respObject)
				if err != nil {
					t.t.Error(err)
				}
			}
			var retJson map[string]interface{}
			err = json.Unmarshal(rb, &retJson)
			if err == nil {
				t.values[c.Name] = retJson
			}
		}
		for _, checker := range c.checkers {
			err = checker(resp, respBody, c.respObject)
			if err != nil {
				t.t.Errorf("%s: %s", c.Name, err)
			}
		}
	}
}

func (t *Tester) applyTemplate(s string) string {
	b, err := t.values.Apply(s)
	if err != nil {
		t.t.Error(err)
		return ""
	}
	return string(b)
}

type Values map[string]interface{}

func (v Values) Apply(templateStr string) ([]byte, error) {

	var funcMap = template.FuncMap{
		"field": v.fieldTmpl,
		"json":  v.jsonFieldTmpl,
	}

	tmpl, err := template.New("tmpl").Funcs(funcMap).Parse(templateStr)
	if err != nil {
		return nil, err
	}

	b := new(bytes.Buffer)

	err = tmpl.Execute(b, v)
	if err != nil {
		return nil, err
	}

	return b.Bytes(), nil
}

// templating funcs

func (v Values) fieldTmpl(key ...string) (interface{}, error) {
	var i interface{}

	i = map[string]interface{}(v)
	var ok bool

	for _, k := range key {
		switch i.(type) {
		case map[string]interface{}:
			i, ok = i.(map[string]interface{})[k]
			if !ok {
				i = "<no value>"
			}
		case map[string]string:
			i, ok = i.(map[string]string)[k]
			if !ok {
				i = "<no value>"
			}
		default:
			return nil, fmt.Errorf("cannot dereference %T", i)
		}
	}
	return i, nil
}

func (v Values) jsonFieldTmpl(key ...string) (interface{}, error) {
	i, err := v.fieldTmpl(key...)
	if err != nil {
		return nil, err
	}
	marshalled, err := json.Marshal(i)
	if err != nil {
		return nil, err
	}
	return string(marshalled), nil
}

// BUILT IN CHECKERS

func ExpectStatus(st int) Checker {
	return func(r *http.Response, body string, respObject interface{}) error {
		if r.StatusCode != st {
			return fmt.Errorf("Bad status code: expected %d, got %d", st, r.StatusCode)
		}
		return nil
	}
}

func ExpectJSONFields(fields ...string) Checker {
	return func(r *http.Response, body string, respObject interface{}) error {
		m := map[string]interface{}{}
		err := json.Unmarshal([]byte(body), &m)
		if err != nil {
			return err
		}
		for _, f := range fields {
			if _, ok := m[f]; !ok {
				return fmt.Errorf("Missing expected field '%s'", f)
			}
		}
		return nil
	}
}

func ExpectListLength(length int) Checker {
	return func(r *http.Response, body string, respObject interface{}) error {
		l := []interface{}{}
		err := json.Unmarshal([]byte(body), &l)
		if err != nil {
			return err
		}
		if len(l) != length {
			return fmt.Errorf("Expected a list of length %d, got %d", length, len(l))
		}
		return nil
	}
}

func ExpectListNonEmpty(r *http.Response, body string, respObject interface{}) error {
	l := []interface{}{}
	err := json.Unmarshal([]byte(body), &l)
	if err != nil {
		return err
	}
	if len(l) == 0 {
		return errors.New("Expected a non empty list")
	}
	return nil
}

func ExpectJSONBranch(nodes ...string) Checker {
	return func(r *http.Response, body string, respObject interface{}) error {
		m := map[string]interface{}{}
		err := json.Unmarshal([]byte(body), &m)
		if err != nil {
			return err
		}
		for i, n := range nodes {
			v, ok := m[n]
			if !ok {
				return fmt.Errorf("Missing node '%s'", n)
			}
			if child, ok := v.(map[string]interface{}); ok {
				m = child
			} else if i == len(nodes)-2 {
				// last child is not an object anymore
				// and there's only one more node to check
				// test last child against last provided node
				lastNode := nodes[i+1]
				if fmt.Sprintf("%v", v) != lastNode {
					return fmt.Errorf("Wrong value: expected '%v', got '%v'", lastNode, v)
				}
				return nil
			}
		}
		return nil
	}
}
<|endoftext|>


<commit_before>package peer_protocol

import (
	"encoding/binary"
	"io"
)

type Integer uint32

func (i *Integer) Read(r io.Reader) error {
	return binary.Read(r, binary.BigEndian, i)
}

// It's perfectly fine to cast these to an int. TODO: Or is it?
func (i Integer) Int() int {
	return int(i)
}

func (i Integer) Uint64() uint64 {
	return uint64(i)
}
<commit_msg>Add peer_protocol.Integer.Uint32<commit_after>package peer_protocol

import (
	"encoding/binary"
	"io"
)

type Integer uint32

func (i *Integer) Read(r io.Reader) error {
	return binary.Read(r, binary.BigEndian, i)
}

// It's perfectly fine to cast these to an int. TODO: Or is it?
func (i Integer) Int() int {
	return int(i)
}

func (i Integer) Uint64() uint64 {
	return uint64(i)
}

func (i Integer) Uint32() uint32 {
	return uint32(i)
}
<|endoftext|>


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

package namespace

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"sync"
	"time"

	"istio.io/istio/pkg/test/framework/components/environment/kube"
	"istio.io/istio/pkg/test/framework/components/istio"
	"istio.io/istio/pkg/test/framework/resource"
	k "istio.io/istio/pkg/test/kube"
	"istio.io/istio/pkg/test/scopes"
)

var (
	idctr int64
	rnd   = rand.New(rand.NewSource(time.Now().UnixNano()))
	mu    sync.Mutex
)

// kubeNamespace represents a Kubernetes namespace. It is tracked as a resource.
type kubeNamespace struct {
	id   resource.ID
	name string
	a    *k.Accessor
	ctx  resource.Context
}

func (n *kubeNamespace) Dump() {
	scopes.CI.Errorf("=== Dumping Namespace %s State...", n.name)

	d, err := n.ctx.CreateTmpDirectory(n.name + "-state")
	if err != nil {
		scopes.CI.Errorf("Unable to create directory for dumping %s contents: %v", n.name, err)
		return
	}

	pods, err := n.a.GetPods(n.name)
	if err != nil {
		scopes.CI.Errorf("Unable to get pods from the namespace: %v", err)
		return
	}

	for _, pod := range pods {
		for _, container := range pod.Spec.Containers {
			l, err := n.a.Logs(pod.Namespace, pod.Name, container.Name, false /* previousLog */)
			if err != nil {
				scopes.CI.Errorf("Unable to get logs for pod/container: %s/%s/%s", pod.Namespace, pod.Name, container.Name)
				continue
			}

			fname := path.Join(d, fmt.Sprintf("%s-%s.log", pod.Name, container.Name))
			if err = ioutil.WriteFile(fname, []byte(l), os.ModePerm); err != nil {
				scopes.CI.Errorf("Unable to write logs for pod/container: %s/%s/%s", pod.Namespace, pod.Name, container.Name)
			}
		}
	}
}

var _ Instance = &kubeNamespace{}
var _ io.Closer = &kubeNamespace{}
var _ resource.Resource = &kubeNamespace{}
var _ resource.Dumper = &kubeNamespace{}

func (n *kubeNamespace) Name() string {
	return n.name
}

func (n *kubeNamespace) ID() resource.ID {
	return n.id
}

// Close implements io.Closer
func (n *kubeNamespace) Close() (err error) {
	if n.name != "" {
		scopes.Framework.Debugf("%s deleting namespace", n.id)
		ns := n.name
		n.name = ""
		err = n.a.DeleteNamespace(ns)
	}

	scopes.Framework.Debugf("%s close complete (err:%v)", n.id, err)
	return
}

func claimKube(ctx resource.Context, name string) (Instance, error) {
	env := ctx.Environment().(*kube.Environment)
	cfg, err := istio.DefaultConfig(ctx)
	if err != nil {
		return nil, err
	}

	if !env.Accessor.NamespaceExists(name) {
		nsConfig := Config{
			Inject:                  true,
			CustomInjectorNamespace: cfg.CustomSidecarInjectorNamespace,
		}
		nsLabels := createNamespaceLabels(&nsConfig)
		if err := env.CreateNamespaceWithLabels(name, "istio-test", nsLabels); err != nil {
			return nil, err
		}

	}
	return &kubeNamespace{name: name}, nil
}

// NewNamespace allocates a new testing namespace.
func newKube(ctx resource.Context, nsConfig *Config) (Instance, error) {
	mu.Lock()
	idctr++
	nsid := idctr
	r := rnd.Intn(99999)
	mu.Unlock()

	env := ctx.Environment().(*kube.Environment)
	ns := fmt.Sprintf("%s-%d-%d", nsConfig.Prefix, nsid, r)

	nsLabels := createNamespaceLabels(nsConfig)
	if err := env.CreateNamespaceWithLabels(ns, "istio-test", nsLabels); err != nil {
		return nil, err
	}

	n := &kubeNamespace{name: ns, a: env.Accessor, ctx: ctx}
	id := ctx.TrackResource(n)
	n.id = id

	return n, nil
}

// createNamespaceLabels will take a namespace config and generate the proper k8s labels
func createNamespaceLabels(cfg *Config) map[string]string {
	l := make(map[string]string)
	if cfg.Inject {
		l["istio-injection"] = "enabled"
		if cfg.CustomInjectorNamespace != "" {
			l["istio-env"] = cfg.CustomInjectorNamespace
		}
	}

	// bring over supplied labels
	for k, v := range cfg.Labels {
		l[k] = v
	}
	return l
}
<commit_msg>Capture init container logs for dumped pods (#20573)<commit_after>//  Copyright 2019 Istio Authors
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

package namespace

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"sync"
	"time"

	"istio.io/istio/pkg/test/framework/components/environment/kube"
	"istio.io/istio/pkg/test/framework/components/istio"
	"istio.io/istio/pkg/test/framework/resource"
	k "istio.io/istio/pkg/test/kube"
	"istio.io/istio/pkg/test/scopes"
)

var (
	idctr int64
	rnd   = rand.New(rand.NewSource(time.Now().UnixNano()))
	mu    sync.Mutex
)

// kubeNamespace represents a Kubernetes namespace. It is tracked as a resource.
type kubeNamespace struct {
	id   resource.ID
	name string
	a    *k.Accessor
	ctx  resource.Context
}

func (n *kubeNamespace) Dump() {
	scopes.CI.Errorf("=== Dumping Namespace %s State...", n.name)

	d, err := n.ctx.CreateTmpDirectory(n.name + "-state")
	if err != nil {
		scopes.CI.Errorf("Unable to create directory for dumping %s contents: %v", n.name, err)
		return
	}

	pods, err := n.a.GetPods(n.name)
	if err != nil {
		scopes.CI.Errorf("Unable to get pods from the namespace: %v", err)
		return
	}

	for _, pod := range pods {
		containers := append(pod.Spec.Containers, pod.Spec.InitContainers...)
		for _, container := range containers {
			l, err := n.a.Logs(pod.Namespace, pod.Name, container.Name, false /* previousLog */)
			if err != nil {
				scopes.CI.Errorf("Unable to get logs for pod/container: %s/%s/%s", pod.Namespace, pod.Name, container.Name)
				continue
			}

			fname := path.Join(d, fmt.Sprintf("%s-%s.log", pod.Name, container.Name))
			if err = ioutil.WriteFile(fname, []byte(l), os.ModePerm); err != nil {
				scopes.CI.Errorf("Unable to write logs for pod/container: %s/%s/%s", pod.Namespace, pod.Name, container.Name)
			}
		}
	}
}

var _ Instance = &kubeNamespace{}
var _ io.Closer = &kubeNamespace{}
var _ resource.Resource = &kubeNamespace{}
var _ resource.Dumper = &kubeNamespace{}

func (n *kubeNamespace) Name() string {
	return n.name
}

func (n *kubeNamespace) ID() resource.ID {
	return n.id
}

// Close implements io.Closer
func (n *kubeNamespace) Close() (err error) {
	if n.name != "" {
		scopes.Framework.Debugf("%s deleting namespace", n.id)
		ns := n.name
		n.name = ""
		err = n.a.DeleteNamespace(ns)
	}

	scopes.Framework.Debugf("%s close complete (err:%v)", n.id, err)
	return
}

func claimKube(ctx resource.Context, name string) (Instance, error) {
	env := ctx.Environment().(*kube.Environment)
	cfg, err := istio.DefaultConfig(ctx)
	if err != nil {
		return nil, err
	}

	if !env.Accessor.NamespaceExists(name) {
		nsConfig := Config{
			Inject:                  true,
			CustomInjectorNamespace: cfg.CustomSidecarInjectorNamespace,
		}
		nsLabels := createNamespaceLabels(&nsConfig)
		if err := env.CreateNamespaceWithLabels(name, "istio-test", nsLabels); err != nil {
			return nil, err
		}

	}
	return &kubeNamespace{name: name}, nil
}

// NewNamespace allocates a new testing namespace.
func newKube(ctx resource.Context, nsConfig *Config) (Instance, error) {
	mu.Lock()
	idctr++
	nsid := idctr
	r := rnd.Intn(99999)
	mu.Unlock()

	env := ctx.Environment().(*kube.Environment)
	ns := fmt.Sprintf("%s-%d-%d", nsConfig.Prefix, nsid, r)

	nsLabels := createNamespaceLabels(nsConfig)
	if err := env.CreateNamespaceWithLabels(ns, "istio-test", nsLabels); err != nil {
		return nil, err
	}

	n := &kubeNamespace{name: ns, a: env.Accessor, ctx: ctx}
	id := ctx.TrackResource(n)
	n.id = id

	return n, nil
}

// createNamespaceLabels will take a namespace config and generate the proper k8s labels
func createNamespaceLabels(cfg *Config) map[string]string {
	l := make(map[string]string)
	if cfg.Inject {
		l["istio-injection"] = "enabled"
		if cfg.CustomInjectorNamespace != "" {
			l["istio-env"] = cfg.CustomInjectorNamespace
		}
	}

	// bring over supplied labels
	for k, v := range cfg.Labels {
		l[k] = v
	}
	return l
}
<|endoftext|>