<commit_before>// Copyright 2018 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this currentFile 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 graph

import (
	"encoding/json"
	"reflect"
	"testing"
	"time"

	"istio.io/tools/isotope/convert/pkg/graph/script"
	"istio.io/tools/isotope/convert/pkg/graph/svc"

	"istio.io/tools/isotope/convert/pkg/graph/svctype"
)

func TestServiceGraph_UnmarshalJSON(t *testing.T) {
	tests := []struct {
		input []byte
		graph ServiceGraph
		err   error
	}{
		{jsonWithOneService, graphWithOneService, nil},
		{jsonWithDefaultsAndManyServices, graphWithDefaultsAndManyServices, nil},
		{
			jsonWithRequestToUndefinedService,
			ServiceGraph{},
			ErrRequestToUndefinedService{"b"},
		},
		{
			jsonWithNestedConcurrentCommand,
			ServiceGraph{},
			ErrNestedConcurrentCommand,
		},
	}

	for _, test := range tests {
		test := test
		t.Run("", func(t *testing.T) {
			t.Parallel()

			var graph ServiceGraph
			err := json.Unmarshal(test.input, &graph)
			if err == nil {
				if !reflect.DeepEqual(test.graph, graph) {
					t.Errorf("expected %v; actual %v", test.graph, graph)
				}
			} else {
				if test.err != err {
					t.Errorf("expected %v; actual %v", test.err, err)
				}
			}
		})
	}
}

var (
	jsonWithOneService = []byte(`
		{
			"services": [{"name": "a"}]
		}
	`)
	graphWithOneService = ServiceGraph{[]svc.Service{
		{
			Name:        "a",
			Type:        svctype.ServiceHTTP,
			NumReplicas: 1,
		},
	}}
	jsonWithDefaultsAndManyServices = []byte(`
		{
			"defaults": {
				"errorRate": 0.1,
				"numReplicas": 2,
				"requestSize": 516,
				"responseSize": 128,
				"script": [
					{ "sleep": "100ms" }
				]
			},
			"services": [
				{
					"name": "a",
					"numReplicas": 5
				},
				{
					"name": "b",
					"script": [
						{
							"call": {
								"service": "a",
								"size": "1KiB"
							}
						},
						{ "sleep": "10ms" }
					]
				},
				{
					"name": "c",
					"type": "grpc",
					"numReplicas": 1,
					"errorRate": "20%",
					"responseSize": "1K",
					"script": [
						[
							{ "call": "a" },
							{ "call": "b" }
						],
						{ "sleep": "10ms" }
					]
				}
			]
		}
	`)
	graphWithDefaultsAndManyServices = ServiceGraph{[]svc.Service{
		{
			Name:         "a",
			Type:         svctype.ServiceHTTP,
			NumReplicas:  5,
			ErrorRate:    0.1,
			ResponseSize: 128,
			Script: script.Script([]script.Command{
				script.SleepCommand(100 * time.Millisecond),
			}),
		},
		{
			Name:         "b",
			Type:         svctype.ServiceHTTP,
			NumReplicas:  2,
			ErrorRate:    0.1,
			ResponseSize: 128,
			Script: script.Script([]script.Command{
				script.RequestCommand{ServiceName: "a", Size: 1024},
				script.SleepCommand(10 * time.Millisecond),
			}),
		},
		{
			Name:         "c",
			Type:         svctype.ServiceGRPC,
			NumReplicas:  1,
			ErrorRate:    0.2,
			ResponseSize: 1024,
			Script: script.Script([]script.Command{
				script.ConcurrentCommand{
					script.RequestCommand{ServiceName: "a", Size: 516},
					script.RequestCommand{ServiceName: "b", Size: 516},
				},
				script.SleepCommand(10 * time.Millisecond),
			}),
		},
	}}
	jsonWithRequestToUndefinedService = []byte(`
		{
			"services": [
				{
					"name": "a",
					"script": [{ "call": "b"}]
				}
			]
		}
	`)
	jsonWithNestedConcurrentCommand = []byte(`
		{
			"services": [
				{
					"name": "a"
				},
				{
					"name": "b",
					"script": [
						[
							[{ "call": "a" }, { "call": "a" }],
							{ "sleep": "10ms" }
						]
					]
				}
			]
		}
	`)
)
<commit_msg>remove reduntant type conversion (#1582)<commit_after>// Copyright 2018 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this currentFile 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 graph

import (
	"encoding/json"
	"reflect"
	"testing"
	"time"

	"istio.io/tools/isotope/convert/pkg/graph/script"
	"istio.io/tools/isotope/convert/pkg/graph/svc"

	"istio.io/tools/isotope/convert/pkg/graph/svctype"
)

func TestServiceGraph_UnmarshalJSON(t *testing.T) {
	tests := []struct {
		input []byte
		graph ServiceGraph
		err   error
	}{
		{jsonWithOneService, graphWithOneService, nil},
		{jsonWithDefaultsAndManyServices, graphWithDefaultsAndManyServices, nil},
		{
			jsonWithRequestToUndefinedService,
			ServiceGraph{},
			ErrRequestToUndefinedService{"b"},
		},
		{
			jsonWithNestedConcurrentCommand,
			ServiceGraph{},
			ErrNestedConcurrentCommand,
		},
	}

	for _, test := range tests {
		test := test
		t.Run("", func(t *testing.T) {
			t.Parallel()

			var graph ServiceGraph
			err := json.Unmarshal(test.input, &graph)
			if err == nil {
				if !reflect.DeepEqual(test.graph, graph) {
					t.Errorf("expected %v; actual %v", test.graph, graph)
				}
			} else {
				if test.err != err {
					t.Errorf("expected %v; actual %v", test.err, err)
				}
			}
		})
	}
}

var (
	jsonWithOneService = []byte(`
		{
			"services": [{"name": "a"}]
		}
	`)
	graphWithOneService = ServiceGraph{[]svc.Service{
		{
			Name:        "a",
			Type:        svctype.ServiceHTTP,
			NumReplicas: 1,
		},
	}}
	jsonWithDefaultsAndManyServices = []byte(`
		{
			"defaults": {
				"errorRate": 0.1,
				"numReplicas": 2,
				"requestSize": 516,
				"responseSize": 128,
				"script": [
					{ "sleep": "100ms" }
				]
			},
			"services": [
				{
					"name": "a",
					"numReplicas": 5
				},
				{
					"name": "b",
					"script": [
						{
							"call": {
								"service": "a",
								"size": "1KiB"
							}
						},
						{ "sleep": "10ms" }
					]
				},
				{
					"name": "c",
					"type": "grpc",
					"numReplicas": 1,
					"errorRate": "20%",
					"responseSize": "1K",
					"script": [
						[
							{ "call": "a" },
							{ "call": "b" }
						],
						{ "sleep": "10ms" }
					]
				}
			]
		}
	`)
	graphWithDefaultsAndManyServices = ServiceGraph{[]svc.Service{
		{
			Name:         "a",
			Type:         svctype.ServiceHTTP,
			NumReplicas:  5,
			ErrorRate:    0.1,
			ResponseSize: 128,
			Script: script.Script([]script.Command{
				100 * time.Millisecond,
			}),
		},
		{
			Name:         "b",
			Type:         svctype.ServiceHTTP,
			NumReplicas:  2,
			ErrorRate:    0.1,
			ResponseSize: 128,
			Script: script.Script([]script.Command{
				script.RequestCommand{ServiceName: "a", Size: 1024},
				10 * time.Millisecond,
			}),
		},
		{
			Name:         "c",
			Type:         svctype.ServiceGRPC,
			NumReplicas:  1,
			ErrorRate:    0.2,
			ResponseSize: 1024,
			Script: script.Script([]script.Command{
				script.ConcurrentCommand{
					script.RequestCommand{ServiceName: "a", Size: 516},
					script.RequestCommand{ServiceName: "b", Size: 516},
				},
				10 * time.Millisecond,
			}),
		},
	}}
	jsonWithRequestToUndefinedService = []byte(`
		{
			"services": [
				{
					"name": "a",
					"script": [{ "call": "b"}]
				}
			]
		}
	`)
	jsonWithNestedConcurrentCommand = []byte(`
		{
			"services": [
				{
					"name": "a"
				},
				{
					"name": "b",
					"script": [
						[
							[{ "call": "a" }, { "call": "a" }],
							{ "sleep": "10ms" }
						]
					]
				}
			]
		}
	`)
)
<|endoftext|>


<commit_before>// Copyright 2009 The Go 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 template

import (
	"fmt";
	"io";
	"os";
	"reflect";
	"template";
	"testing";
)

type Test struct {
	in, out string
}

type T struct {
	item string;
	value string;
}

type S struct {
	header string;
	integer int;
	data []T;
	pdata []*T;
	empty []*T;
	null []*T;
}

var t1 = T{ "ItemNumber1", "ValueNumber1" }
var t2 = T{ "ItemNumber2", "ValueNumber2" }

func uppercase(v reflect.Value) string {
	s := reflect.Indirect(v).(reflect.StringValue).Get();
	t := "";
	for i := 0; i < len(s); i++ {
		c := s[i];
		if 'a' <= c && c <= 'z' {
			c = c + 'A' - 'a'
		}
		t += string(c);
	}
	return t;
}

func plus1(v reflect.Value) string {
	i := reflect.Indirect(v).(reflect.IntValue).Get();
	return fmt.Sprint(i + 1);
}

var formatters = FormatterMap {
	"uppercase" : uppercase,
	"+1" : plus1,
}

var tests = []*Test {
	// Simple
	&Test{ "", "" },
	&Test{ "abc\ndef\n", "abc\ndef\n" },
	&Test{ " {.meta-left}   \n", "{" },
	&Test{ " {.meta-right}   \n", "}" },
	&Test{ " {.space}   \n", " " },
	&Test{ "     {#comment}   \n", "" },

	// Section
	&Test{
		"{.section data }\n"
		"some text for the section\n"
		"{.end}\n",

		"some text for the section\n"
	},
	&Test{
		"{.section data }\n"
		"{header}={integer}\n"
		"{.end}\n",

		"Header=77\n"
	},
	&Test{
		"{.section pdata }\n"
		"{header}={integer}\n"
		"{.end}\n",

		"Header=77\n"
	},
	&Test{
		"{.section pdata }\n"
		"data present\n"
		"{.or}\n"
		"data not present\n"
		"{.end}\n",

		"data present\n"
	},
	&Test{
		"{.section empty }\n"
		"data present\n"
		"{.or}\n"
		"data not present\n"
		"{.end}\n",

		"data not present\n"
	},
	&Test{
		"{.section null }\n"
		"data present\n"
		"{.or}\n"
		"data not present\n"
		"{.end}\n",

		"data not present\n"
	},
	&Test{
		"{.section pdata }\n"
		"{header}={integer}\n"
		"{.section @ }\n"
		"{header}={integer}\n"
		"{.end}\n"
		"{.end}\n",

		"Header=77\n"
		"Header=77\n"
	},

	// Repeated
	&Test{
		"{.section pdata }\n"
		"{.repeated section @ }\n"
		"{item}={value}\n"
		"{.end}\n"
		"{.end}\n",

		"ItemNumber1=ValueNumber1\n"
		"ItemNumber2=ValueNumber2\n"
	},

	// Formatters
	&Test{
		"{.section pdata }\n"
		"{header|uppercase}={integer|+1}\n"
		"{header|html}={integer|str}\n"
		"{.end}\n",

		"HEADER=78\n"
		"Header=77\n"
	},
}

func TestAll(t *testing.T) {
	s := new(S);
	// initialized by hand for clarity.
	s.header = "Header";
	s.integer = 77;
	s.data = []T{ t1, t2 };
	s.pdata = []*T{ &t1, &t2 };
	s.empty = []*T{ };
	s.null = nil;

	var buf io.ByteBuffer;
	for i, test := range tests {
		buf.Reset();
		err := Execute(test.in, s, formatters, &buf);
		if err != nil {
			t.Error("unexpected error:", err)
		}
		if string(buf.Data()) != test.out {
			t.Errorf("for %q: expected %q got %q", test.in, test.out, string(buf.Data()));
		}
	}
}

func TestBadDriverType(t *testing.T) {
	err := Execute("hi", "hello", nil, os.Stdout);
	if err == nil {
		t.Error("failed to detect string as driver type")
	}
	var s S;
}
<commit_msg>template bug<commit_after>// Copyright 2009 The Go 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 template

import (
	"fmt";
	"io";
	"os";
	"reflect";
	"template";
	"testing";
)

type Test struct {
	in, out string
}

type T struct {
	item string;
	value string;
}

type S struct {
	header string;
	integer int;
	data []T;
	pdata []*T;
	empty []*T;
	null []*T;
}

var t1 = T{ "ItemNumber1", "ValueNumber1" }
var t2 = T{ "ItemNumber2", "ValueNumber2" }

func uppercase(v interface{}) string {
	s := v.(string);
	t := "";
	for i := 0; i < len(s); i++ {
		c := s[i];
		if 'a' <= c && c <= 'z' {
			c = c + 'A' - 'a'
		}
		t += string(c);
	}
	return t;
}

func plus1(v interface{}) string {
	i := v.(int);
	return fmt.Sprint(i + 1);
}

func writer(f func(interface{}) string) (func(io.Write, interface{}, string)) {
	return func(w io.Write, v interface{}, format string) {
		io.WriteString(w, f(v));
	}
}

var formatters = FormatterMap {
	"uppercase" : writer(uppercase),
	"+1" : writer(plus1),
}

var tests = []*Test {
	// Simple
	&Test{ "", "" },
	&Test{ "abc\ndef\n", "abc\ndef\n" },
	&Test{ " {.meta-left}   \n", "{" },
	&Test{ " {.meta-right}   \n", "}" },
	&Test{ " {.space}   \n", " " },
	&Test{ "     {#comment}   \n", "" },

	// Section
	&Test{
		"{.section data }\n"
		"some text for the section\n"
		"{.end}\n",

		"some text for the section\n"
	},
	&Test{
		"{.section data }\n"
		"{header}={integer}\n"
		"{.end}\n",

		"Header=77\n"
	},
	&Test{
		"{.section pdata }\n"
		"{header}={integer}\n"
		"{.end}\n",

		"Header=77\n"
	},
	&Test{
		"{.section pdata }\n"
		"data present\n"
		"{.or}\n"
		"data not present\n"
		"{.end}\n",

		"data present\n"
	},
	&Test{
		"{.section empty }\n"
		"data present\n"
		"{.or}\n"
		"data not present\n"
		"{.end}\n",

		"data not present\n"
	},
	&Test{
		"{.section null }\n"
		"data present\n"
		"{.or}\n"
		"data not present\n"
		"{.end}\n",

		"data not present\n"
	},
	&Test{
		"{.section pdata }\n"
		"{header}={integer}\n"
		"{.section @ }\n"
		"{header}={integer}\n"
		"{.end}\n"
		"{.end}\n",

		"Header=77\n"
		"Header=77\n"
	},

	// Repeated
	&Test{
		"{.section pdata }\n"
		"{.repeated section @ }\n"
		"{item}={value}\n"
		"{.end}\n"
		"{.end}\n",

		"ItemNumber1=ValueNumber1\n"
		"ItemNumber2=ValueNumber2\n"
	},

	// Formatters
	&Test{
		"{.section pdata }\n"
		"{header|uppercase}={integer|+1}\n"
		"{header|html}={integer|str}\n"
		"{.end}\n",

		"HEADER=78\n"
		"Header=77\n"
	},

	// Bugs
//	&Test{
//		"{.section data}{.end} {header}\n",
//
//		" 77\n"
//	},
}

func TestAll(t *testing.T) {
	s := new(S);
	// initialized by hand for clarity.
	s.header = "Header";
	s.integer = 77;
	s.data = []T{ t1, t2 };
	s.pdata = []*T{ &t1, &t2 };
	s.empty = []*T{ };
	s.null = nil;

	var buf io.ByteBuffer;
	for i, test := range tests {
		buf.Reset();
		err := Execute(test.in, s, formatters, &buf);
		if err != nil {
			t.Error("unexpected error:", err)
		}
		if string(buf.Data()) != test.out {
			t.Errorf("for %q: expected %q got %q", test.in, test.out, string(buf.Data()));
		}
	}
}

func TestBadDriverType(t *testing.T) {
	err := Execute("hi", "hello", nil, os.Stdout);
	if err == nil {
		t.Error("failed to detect string as driver type")
	}
	var s S;
}
<|endoftext|>


<commit_before>//Package aphcollection contains collection functions for string
//slices
package aphcollection

// Index returns the index of the first instance of s in slice a, or -1 if s is
// not present in a
func Index(a []string, s string) int {
	for i, v := range a {
		if v == s {
			return i
		}
	}
	return -1
}

// Contains reports whether s is present in slice a
func Contains(a []string, s string) bool {
	return Index(a, s) >= 0
}

// Map applies the given function to each element of a, returning slice of
// results
func Map(a []string, fn func(string) string) []string {
	sl := make([]string, len(a))
	for i, v := range a {
		sl[i] = fn(v)
	}
	return sl
}
<commit_msg>added check for empty slice<commit_after>//Package aphcollection contains collection functions for string
//slices
package aphcollection

// Index returns the index of the first instance of s in slice a, or -1 if s is
// not present in a
func Index(a []string, s string) int {
	if len(a) == 0 {
		return -1
	}
	for i, v := range a {
		if v == s {
			return i
		}
	}
	return -1
}

// Contains reports whether s is present in slice a
func Contains(a []string, s string) bool {
	if len(a) == 0 {
		return false
	}
	return Index(a, s) >= 0
}

// Map applies the given function to each element of a, returning slice of
// results
func Map(a []string, fn func(string) string) []string {
	if len(a) == 0 {
		return a
	}
	sl := make([]string, len(a))
	for i, v := range a {
		sl[i] = fn(v)
	}
	return sl
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/mux"
	"github.com/mozilla-services/pushgo/client"
)

func newTestHandler(t *testing.T) (*Handler, *Application) {

	tlogger, _ := NewLogger(&TestLogger{DEBUG, t})

	mx := &TestMetrics{}
	mx.Init(nil, nil)
	store := &NoStore{logger: tlogger, maxChannels: 10}
	count := int32(0)
	pping := &NoopPing{}
	app := &Application{
		hostname:           "test",
		host:               "test",
		clientMinPing:      10 * time.Second,
		clientHelloTimeout: 10 * time.Second,
		clientMux:          new(sync.RWMutex),
		pushLongPongs:      true,
		tokenKey:           []byte(""),
		metrics:            mx,
		clients:            make(map[string]*Client),
		clientCount:        &count,
		store:              store,
		propping:           pping,
	}
	app.SetLogger(tlogger)
	server := &Serv{}
	server.Init(app, server.ConfigStruct())
	app.SetServer(server)
	locator := &NoLocator{logger: tlogger}
	router := NewBroadcastRouter()
	router.Init(app, router.ConfigStruct())
	router.SetLocator(locator)
	app.SetRouter(router)

	handler := &Handler{
		app:        app,
		logger:     tlogger,
		store:      store,
		router:     router,
		metrics:    mx,
		tokenKey:   app.TokenKey(),
		maxDataLen: 140,
		propping:   pping,
	}
	return handler, app
}

func Test_UpdateHandler(t *testing.T) {
	var err error
	uaid := "deadbeef000000000000000000000000"
	chid := "decafbad000000000000000000000000"
	data := "This is a test of the emergency broadcasting system."

	handler, app := newTestHandler(t)
	noPush := &PushWS{
		Socket: nil,
		Born:   time.Now(),
	}
	noPush.SetUAID(uaid)

	worker := &NoWorker{Socket: noPush,
		Logger: app.Logger(),
	}

	app.AddClient(uaid, &Client{
		Worker(worker),
		noPush,
		uaid})
	resp := httptest.NewRecorder()
	// don't bother with encryption right now.
	key, _ := app.Store().IDsToKey(uaid, chid)
	req, err := http.NewRequest("PUT",
		fmt.Sprintf("http://test/update/%s", key),
		nil)
	if req == nil {
		t.Fatal("Update put returned nil")
	}
	if err != nil {
		t.Fatal(err)
	}
	req.Form = make(url.Values)
	req.Form.Add("version", "1")
	req.Form.Add("data", data)
	tmux := mux.NewRouter()

	// Yay! Actually try the test!
	tmux.HandleFunc("/update/{key}", handler.UpdateHandler)
	tmux.ServeHTTP(resp, req)
	if resp.Body.String() != "{}" {
		t.Error("Unexpected response from server")
	}
	rep := FlushData{}
	if err = json.Unmarshal(worker.Outbuffer, &rep); err != nil {
		t.Errorf("Could not read output buffer %s", err.Error())
	}
	if rep.Data != data {
		t.Error("Returned data does not match expected value")
	}
	if rep.Version != 1 {
		t.Error("Returned version does not match expected value")
	}
}

func endpointIds(uri *url.URL) (deviceId, channelId string, ok bool) {
	if !uri.IsAbs() {
		ok = false
		return
	}
	pathPrefix := "/update/"
	i := strings.Index(uri.Path, pathPrefix)
	if i < 0 {
		ok = false
		return
	}
	key := strings.SplitN(uri.Path[i+len(pathPrefix):], ".", 2)
	if len(key) < 2 {
		ok = false
		return
	}
	return key[0], key[1], true
}

func TestBadKey(t *testing.T) {
	origin, err := Server.Origin()
	if err != nil {
		t.Fatalf("Error initializing test server: %#v", err)
	}
	conn, deviceId, err := client.Dial(origin)
	if err != nil {
		t.Fatalf("Error dialing origin: %#v", err)
	}
	defer conn.Close()
	channelId, endpoint, err := conn.Subscribe()
	if err != nil {
		t.Fatalf("Error subscribing to channel: %#v", err)
	}
	uri, err := url.ParseRequestURI(endpoint)
	if err != nil {
		t.Fatalf("Error parsing push endpoint %#v: %#v", endpoint, err)
	}
	keyDevice, keyChannel, ok := endpointIds(uri)
	if !ok {
		t.Errorf("Incomplete push endpoint: %#v", endpoint)
	}
	if keyDevice != deviceId {
		t.Errorf("Mismatched device IDs: got %#v; want %#v", keyDevice, deviceId)
	}
	if keyChannel != channelId {
		t.Errorf("Mismatched channel IDs: got %#v; want %#v", keyChannel, channelId)
	}
	newURI, _ := uri.Parse(fmt.Sprintf("/update/%s", keyDevice))
	err = client.Notify(newURI.String(), 1)
	clientErr, ok := err.(client.Error)
	if !ok {
		t.Errorf("Type assertion failed for endpoint error: %#v", err)
	} else if clientErr.Status() != 404 {
		t.Errorf("Unexpected endpoint status: got %#v; want 404", clientErr.Status())
	}
}

func TestMissingKey(t *testing.T) {
	origin, err := Server.Origin()
	if err != nil {
		t.Fatalf("Error initializing test server: %#v", err)
	}
	conn, deviceId, err := client.Dial(origin)
	if err != nil {
		t.Fatalf("Error dialing origin: %#v", err)
	}
	defer conn.Close()
	channelId, endpoint, err := conn.Subscribe()
	if err != nil {
		t.Fatalf("Error subscribing to channel: %#v", err)
	}
	uri, err := url.ParseRequestURI(endpoint)
	if err != nil {
		t.Fatalf("Error parsing push endpoint %#v: %#v", endpoint, err)
	}
	keyDevice, keyChannel, ok := endpointIds(uri)
	if !ok {
		t.Errorf("Incomplete push endpoint: %#v", endpoint)
	}
	if keyDevice != deviceId {
		t.Errorf("Mismatched device IDs: got %#v; want %#v", keyDevice, deviceId)
	}
	if keyChannel != channelId {
		t.Errorf("Mismatched channel IDs: got %#v; want %#v", keyChannel, channelId)
	}
	newURI, _ := uri.Parse(fmt.Sprintf("/update/%s.", keyDevice))
	err = client.Notify(newURI.String(), 1)
	clientErr, ok := err.(client.Error)
	if !ok {
		t.Errorf("Type assertion failed for endpoint error: %#v", err)
	} else if clientErr.Status() != 404 {
		t.Errorf("Unexpected endpoint status: got %#v; want 404", clientErr.Status())
	}
}
<commit_msg>Add test for empty Content-Type header<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/mux"
	"github.com/mozilla-services/pushgo/client"
)

func newTestHandler(t *testing.T) (*Handler, *Application) {

	tlogger, _ := NewLogger(&TestLogger{DEBUG, t})

	mx := &TestMetrics{}
	mx.Init(nil, nil)
	store := &NoStore{logger: tlogger, maxChannels: 10}
	count := int32(0)
	pping := &NoopPing{}
	app := &Application{
		hostname:           "test",
		host:               "test",
		clientMinPing:      10 * time.Second,
		clientHelloTimeout: 10 * time.Second,
		clientMux:          new(sync.RWMutex),
		pushLongPongs:      true,
		tokenKey:           []byte(""),
		metrics:            mx,
		clients:            make(map[string]*Client),
		clientCount:        &count,
		store:              store,
		propping:           pping,
	}
	app.SetLogger(tlogger)
	server := &Serv{}
	server.Init(app, server.ConfigStruct())
	app.SetServer(server)
	locator := &NoLocator{logger: tlogger}
	router := NewBroadcastRouter()
	router.Init(app, router.ConfigStruct())
	router.SetLocator(locator)
	app.SetRouter(router)

	handler := &Handler{
		app:        app,
		logger:     tlogger,
		store:      store,
		router:     router,
		metrics:    mx,
		tokenKey:   app.TokenKey(),
		maxDataLen: 140,
		propping:   pping,
	}
	return handler, app
}

func Test_UpdateHandler(t *testing.T) {
	var err error
	uaid := "deadbeef000000000000000000000000"
	chid := "decafbad000000000000000000000000"
	data := "This is a test of the emergency broadcasting system."

	handler, app := newTestHandler(t)
	noPush := &PushWS{
		Socket: nil,
		Born:   time.Now(),
	}
	noPush.SetUAID(uaid)

	worker := &NoWorker{Socket: noPush,
		Logger: app.Logger(),
	}

	app.AddClient(uaid, &Client{
		Worker(worker),
		noPush,
		uaid})
	resp := httptest.NewRecorder()
	// don't bother with encryption right now.
	key, _ := app.Store().IDsToKey(uaid, chid)
	req, err := http.NewRequest("PUT",
		fmt.Sprintf("http://test/update/%s", key),
		nil)
	if req == nil {
		t.Fatal("Update put returned nil")
	}
	if err != nil {
		t.Fatal(err)
	}
	req.Form = make(url.Values)
	req.Form.Add("version", "1")
	req.Form.Add("data", data)
	tmux := mux.NewRouter()

	// Yay! Actually try the test!
	tmux.HandleFunc("/update/{key}", handler.UpdateHandler)
	tmux.ServeHTTP(resp, req)
	if resp.Body.String() != "{}" {
		t.Error("Unexpected response from server")
	}
	rep := FlushData{}
	if err = json.Unmarshal(worker.Outbuffer, &rep); err != nil {
		t.Errorf("Could not read output buffer %s", err.Error())
	}
	if rep.Data != data {
		t.Error("Returned data does not match expected value")
	}
	if rep.Version != 1 {
		t.Error("Returned version does not match expected value")
	}

	//retry without a Content-Type header
	resp = httptest.NewRecorder()
	req.Header.Set("Content-Type", "")
	tmux.ServeHTTP(resp, req)
	if resp.Body.String() != "{}" {
		t.Error("Unexpected response from server")
	}
	rep = FlushData{}
	if err = json.Unmarshal(worker.Outbuffer, &rep); err != nil {
		t.Errorf("Could not read output buffer %s", err.Error())
	}
	if rep.Data != data {
		t.Error("Returned data does not match expected value")
	}
	if rep.Version != 1 {
		t.Error("Returned version does not match expected value")
	}

}

func endpointIds(uri *url.URL) (deviceId, channelId string, ok bool) {
	if !uri.IsAbs() {
		ok = false
		return
	}
	pathPrefix := "/update/"
	i := strings.Index(uri.Path, pathPrefix)
	if i < 0 {
		ok = false
		return
	}
	key := strings.SplitN(uri.Path[i+len(pathPrefix):], ".", 2)
	if len(key) < 2 {
		ok = false
		return
	}
	return key[0], key[1], true
}

func TestBadKey(t *testing.T) {
	origin, err := Server.Origin()
	if err != nil {
		t.Fatalf("Error initializing test server: %#v", err)
	}
	conn, deviceId, err := client.Dial(origin)
	if err != nil {
		t.Fatalf("Error dialing origin: %#v", err)
	}
	defer conn.Close()
	channelId, endpoint, err := conn.Subscribe()
	if err != nil {
		t.Fatalf("Error subscribing to channel: %#v", err)
	}
	uri, err := url.ParseRequestURI(endpoint)
	if err != nil {
		t.Fatalf("Error parsing push endpoint %#v: %#v", endpoint, err)
	}
	keyDevice, keyChannel, ok := endpointIds(uri)
	if !ok {
		t.Errorf("Incomplete push endpoint: %#v", endpoint)
	}
	if keyDevice != deviceId {
		t.Errorf("Mismatched device IDs: got %#v; want %#v", keyDevice, deviceId)
	}
	if keyChannel != channelId {
		t.Errorf("Mismatched channel IDs: got %#v; want %#v", keyChannel, channelId)
	}
	newURI, _ := uri.Parse(fmt.Sprintf("/update/%s", keyDevice))
	err = client.Notify(newURI.String(), 1)
	clientErr, ok := err.(client.Error)
	if !ok {
		t.Errorf("Type assertion failed for endpoint error: %#v", err)
	} else if clientErr.Status() != 404 {
		t.Errorf("Unexpected endpoint status: got %#v; want 404", clientErr.Status())
	}
}

func TestMissingKey(t *testing.T) {
	origin, err := Server.Origin()
	if err != nil {
		t.Fatalf("Error initializing test server: %#v", err)
	}
	conn, deviceId, err := client.Dial(origin)
	if err != nil {
		t.Fatalf("Error dialing origin: %#v", err)
	}
	defer conn.Close()
	channelId, endpoint, err := conn.Subscribe()
	if err != nil {
		t.Fatalf("Error subscribing to channel: %#v", err)
	}
	uri, err := url.ParseRequestURI(endpoint)
	if err != nil {
		t.Fatalf("Error parsing push endpoint %#v: %#v", endpoint, err)
	}
	keyDevice, keyChannel, ok := endpointIds(uri)
	if !ok {
		t.Errorf("Incomplete push endpoint: %#v", endpoint)
	}
	if keyDevice != deviceId {
		t.Errorf("Mismatched device IDs: got %#v; want %#v", keyDevice, deviceId)
	}
	if keyChannel != channelId {
		t.Errorf("Mismatched channel IDs: got %#v; want %#v", keyChannel, channelId)
	}
	newURI, _ := uri.Parse(fmt.Sprintf("/update/%s.", keyDevice))
	err = client.Notify(newURI.String(), 1)
	clientErr, ok := err.(client.Error)
	if !ok {
		t.Errorf("Type assertion failed for endpoint error: %#v", err)
	} else if clientErr.Status() != 404 {
		t.Errorf("Unexpected endpoint status: got %#v; want 404", clientErr.Status())
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/qiniu/qshell/v2/iqshell"
	"github.com/qiniu/qshell/v2/iqshell/common/alert"
	"github.com/qiniu/qshell/v2/iqshell/common/data"
	"github.com/qiniu/qshell/v2/iqshell/common/export"
	"github.com/qiniu/qshell/v2/iqshell/common/flow"
	"github.com/qiniu/qshell/v2/iqshell/common/log"
	"github.com/qiniu/qshell/v2/iqshell/common/utils"
	"github.com/qiniu/qshell/v2/iqshell/common/workspace"
	"github.com/qiniu/qshell/v2/iqshell/storage/object"
	"github.com/qiniu/qshell/v2/iqshell/storage/object/batch"
	"path/filepath"
	"time"
)

type FetchInfo struct {
	QiniuBucket   string
	Host          string
	BatchInfo     batch.Info
	AwsBucketInfo ListBucketInfo
}

func (info *FetchInfo) Check() *data.CodeError {
	// check AWS bucket
	if info.AwsBucketInfo.Bucket == "" {
		return alert.CannotEmptyError("AWS bucket", "")
	}

	// check AWS region
	if info.AwsBucketInfo.Region == "" {
		return alert.CannotEmptyError("AWS region", "")
	}

	// check AWS region
	if info.QiniuBucket == "" {
		return alert.CannotEmptyError("Qiniu bucket", "")
	}

	if info.AwsBucketInfo.Id == "" || info.AwsBucketInfo.SecretKey == "" {
		return alert.CannotEmptyError("AWS ID and SecretKey", "")
	}

	if info.BatchInfo.WorkerCount <= 0 || info.BatchInfo.WorkerCount >= 1000 {
		info.BatchInfo.WorkerCount = 20
	}

	return nil
}

func Fetch(cfg *iqshell.Config, info FetchInfo) {
	cfg.JobPathBuilder = func(cmdPath string) string {
		jobId := utils.Md5Hex(fmt.Sprintf("%s:%s:%s:%s", cfg.CmdCfg.CmdId, info.AwsBucketInfo.Region, info.AwsBucketInfo.Bucket, info.QiniuBucket))
		return filepath.Join(cmdPath, jobId)
	}
	if shouldContinue := iqshell.CheckAndLoad(cfg, iqshell.CheckAndLoadInfo{
		Checker: &info,
	}); !shouldContinue {
		return
	}

	exporter, err := export.NewFileExport(export.FileExporterConfig{
		SuccessExportFilePath:   info.BatchInfo.SuccessExportFilePath,
		FailExportFilePath:      info.BatchInfo.FailExportFilePath,
		OverwriteExportFilePath: info.BatchInfo.OverwriteExportFilePath,
	})

	if err != nil {
		log.ErrorF("get export error:%v", err)
		return
	}

	fetchInfoChan := make(chan flow.Work, info.BatchInfo.WorkerCount)
	// 生产者
	go func() {
		if e := listBucket(info.AwsBucketInfo, func(svc *s3.S3, obj *s3.Object) {
			req, _ := svc.GetObjectRequest(&s3.GetObjectInput{
				Bucket: aws.String(info.AwsBucketInfo.Bucket),
				Key:    obj.Key,
			})
			if downloadUrl, e := req.Presign(5 * 3600 * time.Second); e == nil {
				fetchInfoChan <- &object.FetchApiInfo{
					Bucket:  info.QiniuBucket,
					Key:     *obj.Key,
					FromUrl: downloadUrl,
				}
				log.DebugF("get object:%s\t%d\t%s\t%s\n%s", *obj.Key, *obj.Size, *obj.ETag, *obj.LastModified, downloadUrl)
			} else {
				log.ErrorF("fetch([%s:%s]) create download url error: %v", info.AwsBucketInfo.Bucket, *obj.Key, e)
			}
		}); e != nil {
			log.Error(e)
		}
		close(fetchInfoChan)
	}()

	var overseer flow.Overseer
	if info.BatchInfo.EnableRecord {
		dbPath := filepath.Join(workspace.GetJobDir(), ".recorder")
		log.DebugF("aws batch fetch recorder:%s", dbPath)
		if overseer, err = flow.NewDBRecordOverseer(dbPath, func() *flow.WorkRecord {
			return &flow.WorkRecord{
				WorkInfo: &flow.WorkInfo{
					Data: "",
					Work: nil,
				},
				Result: &object.FetchResult{},
				Err:    nil,
			}
		}); err != nil {
			log.ErrorF("aws batch fetch create overseer error:%v", err)
			return
		}
	} else {
		log.Debug("aws batch fetch recorder:Not Enable")
	}

	metric := &batch.Metric{}
	metric.Start()
	flow.New(info.BatchInfo.Info).
		WorkProviderWithChan(fetchInfoChan).
		WorkerProvider(flow.NewWorkerProvider(func() (flow.Worker, *data.CodeError) {
			return flow.NewSimpleWorker(func(workInfo *flow.WorkInfo) (flow.Result, *data.CodeError) {
				in := workInfo.Work.(*object.FetchApiInfo)
				return object.Fetch(*in)
			}), nil
		})).
		FlowWillStartFunc(func(flow *flow.Flow) (err *data.CodeError) {
			metric.AddTotalCount(flow.WorkProvider.WorkTotalCount())
			return nil
		}).
		SetOverseer(overseer).
		ShouldRedo(func(workInfo *flow.WorkInfo, workRecord *flow.WorkRecord) (shouldRedo bool, cause *data.CodeError) {
			if workRecord.Err == nil {
				return false, nil
			}

			if !info.BatchInfo.RecordRedoWhileError {
				return false, workRecord.Err
			}

			result, _ := workRecord.Result.(*object.FetchResult)
			if result == nil {
				return true, data.NewEmptyError().AppendDesc("no result found")
			}
			if !result.IsValid() {
				return true, data.NewEmptyError().AppendDesc("result is invalid")
			}
			return false, nil
		}).
		OnWorkSkip(func(work *flow.WorkInfo, result flow.Result, err *data.CodeError) {
			metric.AddCurrentCount(1)
			metric.PrintProgress("Batching:" + work.Data)

			operationResult, _ := result.(*object.FetchResult)
			if err != nil && err.Code == data.ErrorCodeAlreadyDone {
				if operationResult != nil && operationResult.IsValid() {
					metric.AddSuccessCount(1)
					log.DebugF("Skip line:%s because have done and success", work.Data)
				} else {
					metric.AddFailureCount(1)
					log.DebugF("Skip line:%s because have done and failure, %v", work.Data, err)
				}
			} else {
				metric.AddSkippedCount(1)
				exporter.Fail().ExportF("%s%s%v", work.Data, flow.ErrorSeparate, err)
				log.DebugF("Skip line:%s because:%v", work.Data, err)
			}

		}).
		OnWorkSuccess(func(workInfo *flow.WorkInfo, result flow.Result) {
			metric.AddCurrentCount(1)
			metric.AddSuccessCount(1)
			metric.PrintProgress("Batching:" + workInfo.Data)

			in, _ := workInfo.Work.(*object.FetchApiInfo)
			exporter.Success().ExportF("%s\t%s", in.FromUrl, in.Bucket)
			log.InfoF("AWS Fetch Success, '%s' => [%s:%s]", in.FromUrl, in.Bucket, in.Key)
		}).
		OnWorkFail(func(workInfo *flow.WorkInfo, err *data.CodeError) {
			metric.AddCurrentCount(1)
			metric.AddFailureCount(1)
			metric.PrintProgress("AWS Batching:" + workInfo.Data)

			exporter.Fail().ExportF("%s%s%v", workInfo.Data, flow.ErrorSeparate, err)
			if in, ok := workInfo.Work.(*object.FetchApiInfo); ok {
				log.ErrorF("AWS Fetch Failed, '%s' => [%s:%s], Error: %v", in.FromUrl, in.Bucket, in.Key, err)
			} else {
				log.ErrorF("AWS Fetch Failed, %s, Error: %s", workInfo.Data, err)
			}
		}).Build().Start()

	metric.End()
	if metric.TotalCount <= 0 {
		metric.TotalCount = metric.SuccessCount + metric.FailureCount + metric.SkippedCount
	}

	// 输出结果
	resultPath := filepath.Join(workspace.GetJobDir(), ".result")
	if e := utils.MarshalToFile(resultPath, metric); e != nil {
		log.ErrorF("save aws batch fetch result to path:%s error:%v", resultPath, e)
	} else {
		log.DebugF("save aws batch fetch result to path:%s", resultPath)
	}

	log.Info("------------- AWS Batch Result --------------")
	log.InfoF("%20s%10d", "Total:", metric.TotalCount)
	log.InfoF("%20s%10d", "Success:", metric.SuccessCount)
	log.InfoF("%20s%10d", "Failure:", metric.FailureCount)
	log.InfoF("%20s%10d", "Skipped:", metric.SkippedCount)
	log.InfoF("%20s%10ds", "Duration:", metric.Duration)
	log.InfoF("--------------------------------------------")
}
<commit_msg>gofmt<commit_after>package aws

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/qiniu/qshell/v2/iqshell"
	"github.com/qiniu/qshell/v2/iqshell/common/alert"
	"github.com/qiniu/qshell/v2/iqshell/common/data"
	"github.com/qiniu/qshell/v2/iqshell/common/export"
	"github.com/qiniu/qshell/v2/iqshell/common/flow"
	"github.com/qiniu/qshell/v2/iqshell/common/log"
	"github.com/qiniu/qshell/v2/iqshell/common/utils"
	"github.com/qiniu/qshell/v2/iqshell/common/workspace"
	"github.com/qiniu/qshell/v2/iqshell/storage/object"
	"github.com/qiniu/qshell/v2/iqshell/storage/object/batch"
	"path/filepath"
	"time"
)

type FetchInfo struct {
	QiniuBucket   string
	Host          string
	BatchInfo     batch.Info
	AwsBucketInfo ListBucketInfo
}

func (info *FetchInfo) Check() *data.CodeError {
	// check AWS bucket
	if info.AwsBucketInfo.Bucket == "" {
		return alert.CannotEmptyError("AWS bucket", "")
	}

	// check AWS region
	if info.AwsBucketInfo.Region == "" {
		return alert.CannotEmptyError("AWS region", "")
	}

	// check AWS region
	if info.QiniuBucket == "" {
		return alert.CannotEmptyError("Qiniu bucket", "")
	}

	if info.AwsBucketInfo.Id == "" || info.AwsBucketInfo.SecretKey == "" {
		return alert.CannotEmptyError("AWS ID and SecretKey", "")
	}

	if info.BatchInfo.WorkerCount <= 0 || info.BatchInfo.WorkerCount >= 1000 {
		info.BatchInfo.WorkerCount = 20
	}

	return nil
}

func Fetch(cfg *iqshell.Config, info FetchInfo) {
	cfg.JobPathBuilder = func(cmdPath string) string {
		jobId := utils.Md5Hex(fmt.Sprintf("%s:%s:%s:%s", cfg.CmdCfg.CmdId, info.AwsBucketInfo.Region, info.AwsBucketInfo.Bucket, info.QiniuBucket))
		return filepath.Join(cmdPath, jobId)
	}
	if shouldContinue := iqshell.CheckAndLoad(cfg, iqshell.CheckAndLoadInfo{
		Checker: &info,
	}); !shouldContinue {
		return
	}

	exporter, err := export.NewFileExport(export.FileExporterConfig{
		SuccessExportFilePath:   info.BatchInfo.SuccessExportFilePath,
		FailExportFilePath:      info.BatchInfo.FailExportFilePath,
		OverwriteExportFilePath: info.BatchInfo.OverwriteExportFilePath,
	})

	if err != nil {
		log.ErrorF("get export error:%v", err)
		return
	}

	fetchInfoChan := make(chan flow.Work, info.BatchInfo.WorkerCount)
	// 生产者
	go func() {
		if e := listBucket(info.AwsBucketInfo, func(svc *s3.S3, obj *s3.Object) {
			req, _ := svc.GetObjectRequest(&s3.GetObjectInput{
				Bucket: aws.String(info.AwsBucketInfo.Bucket),
				Key:    obj.Key,
			})
			if downloadUrl, e := req.Presign(5 * 3600 * time.Second); e == nil {
				fetchInfoChan <- &object.FetchApiInfo{
					Bucket:  info.QiniuBucket,
					Key:     *obj.Key,
					FromUrl: downloadUrl,
				}
				log.DebugF("get object:%s\t%d\t%s\t%s\n%s", *obj.Key, *obj.Size, *obj.ETag, *obj.LastModified, downloadUrl)
			} else {
				log.ErrorF("fetch([%s:%s]) create download url error: %v", info.AwsBucketInfo.Bucket, *obj.Key, e)
			}
		}); e != nil {
			log.Error(e)
		}
		close(fetchInfoChan)
	}()

	var overseer flow.Overseer
	if info.BatchInfo.EnableRecord {
		dbPath := filepath.Join(workspace.GetJobDir(), ".recorder")
		log.DebugF("aws batch fetch recorder:%s", dbPath)
		if overseer, err = flow.NewDBRecordOverseer(dbPath, func() *flow.WorkRecord {
			return &flow.WorkRecord{
				WorkInfo: &flow.WorkInfo{
					Data: "",
					Work: nil,
				},
				Result: &object.FetchResult{},
				Err:    nil,
			}
		}); err != nil {
			log.ErrorF("aws batch fetch create overseer error:%v", err)
			return
		}
	} else {
		log.Debug("aws batch fetch recorder:Not Enable")
	}

	metric := &batch.Metric{}
	metric.Start()
	flow.New(info.BatchInfo.Info).
		WorkProviderWithChan(fetchInfoChan).
		WorkerProvider(flow.NewWorkerProvider(func() (flow.Worker, *data.CodeError) {
			return flow.NewSimpleWorker(func(workInfo *flow.WorkInfo) (flow.Result, *data.CodeError) {
				in := workInfo.Work.(*object.FetchApiInfo)
				return object.Fetch(*in)
			}), nil
		})).
		FlowWillStartFunc(func(flow *flow.Flow) (err *data.CodeError) {
			metric.AddTotalCount(flow.WorkProvider.WorkTotalCount())
			return nil
		}).
		SetOverseer(overseer).
		ShouldRedo(func(workInfo *flow.WorkInfo, workRecord *flow.WorkRecord) (shouldRedo bool, cause *data.CodeError) {
			if workRecord.Err == nil {
				return false, nil
			}

			if !info.BatchInfo.RecordRedoWhileError {
				return false, workRecord.Err
			}

			result, _ := workRecord.Result.(*object.FetchResult)
			if result == nil {
				return true, data.NewEmptyError().AppendDesc("no result found")
			}
			if !result.IsValid() {
				return true, data.NewEmptyError().AppendDesc("result is invalid")
			}
			return false, nil
		}).
		OnWorkSkip(func(work *flow.WorkInfo, result flow.Result, err *data.CodeError) {
			metric.AddCurrentCount(1)
			metric.PrintProgress("Batching:" + work.Data)

			operationResult, _ := result.(*object.FetchResult)
			if err != nil && err.Code == data.ErrorCodeAlreadyDone {
				if operationResult != nil && operationResult.IsValid() {
					metric.AddSuccessCount(1)
					log.DebugF("Skip line:%s because have done and success", work.Data)
				} else {
					metric.AddFailureCount(1)
					log.DebugF("Skip line:%s because have done and failure, %v", work.Data, err)
				}
			} else {
				metric.AddSkippedCount(1)
				exporter.Fail().ExportF("%s%s%v", work.Data, flow.ErrorSeparate, err)
				log.DebugF("Skip line:%s because:%v", work.Data, err)
			}

		}).
		OnWorkSuccess(func(workInfo *flow.WorkInfo, result flow.Result) {
			metric.AddCurrentCount(1)
			metric.AddSuccessCount(1)
			metric.PrintProgress("Batching:" + workInfo.Data)

			in, _ := workInfo.Work.(*object.FetchApiInfo)
			exporter.Success().ExportF("%s\t%s", in.FromUrl, in.Bucket)
			log.InfoF("AWS Fetch Success, '%s' => [%s:%s]", in.FromUrl, in.Bucket, in.Key)
		}).
		OnWorkFail(func(workInfo *flow.WorkInfo, err *data.CodeError) {
			metric.AddCurrentCount(1)
			metric.AddFailureCount(1)
			metric.PrintProgress("AWS Batching:" + workInfo.Data)

			exporter.Fail().ExportF("%s%s%v", workInfo.Data, flow.ErrorSeparate, err)
			if in, ok := workInfo.Work.(*object.FetchApiInfo); ok {
				log.ErrorF("AWS Fetch Failed, '%s' => [%s:%s], Error: %v", in.FromUrl, in.Bucket, in.Key, err)
			} else {
				log.ErrorF("AWS Fetch Failed, %s, Error: %s", workInfo.Data, err)
			}
		}).Build().Start()

	metric.End()
	if metric.TotalCount <= 0 {
		metric.TotalCount = metric.SuccessCount + metric.FailureCount + metric.SkippedCount
	}

	// 输出结果
	resultPath := filepath.Join(workspace.GetJobDir(), ".result")
	if e := utils.MarshalToFile(resultPath, metric); e != nil {
		log.ErrorF("save aws batch fetch result to path:%s error:%v", resultPath, e)
	} else {
		log.DebugF("save aws batch fetch result to path:%s", resultPath)
	}

	log.Info("------------- AWS Batch Result --------------")
	log.InfoF("%20s%10d", "Total:", metric.TotalCount)
	log.InfoF("%20s%10d", "Success:", metric.SuccessCount)
	log.InfoF("%20s%10d", "Failure:", metric.FailureCount)
	log.InfoF("%20s%10d", "Skipped:", metric.SkippedCount)
	log.InfoF("%20s%10ds", "Duration:", metric.Duration)
	log.InfoF("--------------------------------------------")
}
<|endoftext|>


<commit_before>package s3

// Tests for the PFS' S3 emulation API. Note that, in calls to
// `tu.UniqueString`, all lowercase characters are used, unlike in other
// tests. This is in order to generate repo names that are also valid bucket
// names. Otherwise minio complains that the bucket name is not valid.

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"

	minio "github.com/minio/minio-go"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/server/pfs/server"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	tu "github.com/pachyderm/pachyderm/src/server/pkg/testutil"
)

func serve(t *testing.T, pc *client.APIClient) (*http.Server, uint16) {
	port := tu.UniquePort()
	srv := Server(pc, port)

	go func() {
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			t.Fatalf("http server returned an error: %v", err)
		}
	}()

	// Wait for the server to start
	require.NoError(t, backoff.Retry(func() error {
		c := &http.Client{}
		res, err := c.Get(fmt.Sprintf("http://127.0.0.1:%d/_ping", port))
		if err != nil {
			return err
		} else if res.StatusCode != 200 {
			return fmt.Errorf("Unexpected status code: %d", res.StatusCode)
		}
		return nil
	}, backoff.NewTestingBackOff()))
	return srv, port
}

func getObject(c *minio.Client, repo, branch, file string) (string, error) {
	obj, err := c.GetObject(repo, fmt.Sprintf("%s/%s", branch, file))
	if err != nil {
		return "", err
	}
	bytes, err := ioutil.ReadAll(obj)
	if err != nil {
		return "", err
	}
	if err = obj.Close(); err != nil {
		return "", err
	}
	return string(bytes), nil
}

func TestGetFile(t *testing.T) {
	repo := tu.UniqueString("testgetfile")
	pc := server.GetPachClient(t)
	require.NoError(t, pc.CreateRepo(repo))
	_, err := pc.PutFile(repo, "master", "file", strings.NewReader("content"))
	require.NoError(t, err)

	srv, port := serve(t, pc)
	c, err := minio.New(fmt.Sprintf("127.0.0.1:%d", port), "id", "secret", false)
	require.NoError(t, err)

	fetchedContent, err := getObject(c, repo, "master", "file")
	require.NoError(t, err)
	require.Equal(t, "content", fetchedContent)

	require.NoError(t, srv.Close())
}

func TestGetFileInBranch(t *testing.T) {
	repo := tu.UniqueString("testgetfileinbranch")
	pc := server.GetPachClient(t)
	require.NoError(t, pc.CreateRepo(repo))
	require.NoError(t, pc.CreateBranch(repo, "branch", "", nil))
	_, err := pc.PutFile(repo, "branch", "file", strings.NewReader("content"))
	require.NoError(t, err)

	srv, port := serve(t, pc)
	c, err := minio.New(fmt.Sprintf("127.0.0.1:%d", port), "id", "secret", false)
	require.NoError(t, err)

	fetchedContent, err := getObject(c, repo, "branch", "file")
	require.NoError(t, err)
	require.Equal(t, "content", fetchedContent)

	require.NoError(t, srv.Close())
}

func TestNonExistingBranch(t *testing.T) {
	repo := tu.UniqueString("testnonexistingbranch")
	pc := server.GetPachClient(t)
	require.NoError(t, pc.CreateRepo(repo))

	srv, port := serve(t, pc)
	c, err := minio.New(fmt.Sprintf("127.0.0.1:%d", port), "id", "secret", false)
	require.NoError(t, err)

	_, err = getObject(c, repo, "branch", "file")
	require.YesError(t, err)
	require.Equal(t, err.Error(), "The specified key does not exist.")
	require.NoError(t, srv.Close())
}

func TestNonExistingRepo(t *testing.T) {
	repo := tu.UniqueString("testnonexistingrepo")
	pc := server.GetPachClient(t)

	srv, port := serve(t, pc)
	c, err := minio.New(fmt.Sprintf("127.0.0.1:%d", port), "id", "secret", false)
	require.NoError(t, err)

	_, err = getObject(c, repo, "master", "file")
	require.YesError(t, err)
	require.Equal(t, err.Error(), "The specified bucket does not exist.")
	require.NoError(t, srv.Close())
}
<commit_msg>Use defer to ensure the object is closed<commit_after>package s3

// Tests for the PFS' S3 emulation API. Note that, in calls to
// `tu.UniqueString`, all lowercase characters are used, unlike in other
// tests. This is in order to generate repo names that are also valid bucket
// names. Otherwise minio complains that the bucket name is not valid.

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"

	minio "github.com/minio/minio-go"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/server/pfs/server"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	tu "github.com/pachyderm/pachyderm/src/server/pkg/testutil"
)

func serve(t *testing.T, pc *client.APIClient) (*http.Server, uint16) {
	port := tu.UniquePort()
	srv := Server(pc, port)

	go func() {
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			t.Fatalf("http server returned an error: %v", err)
		}
	}()

	// Wait for the server to start
	require.NoError(t, backoff.Retry(func() error {
		c := &http.Client{}
		res, err := c.Get(fmt.Sprintf("http://127.0.0.1:%d/_ping", port))
		if err != nil {
			return err
		} else if res.StatusCode != 200 {
			return fmt.Errorf("Unexpected status code: %d", res.StatusCode)
		}
		return nil
	}, backoff.NewTestingBackOff()))
	return srv, port
}

func getObject(c *minio.Client, repo, branch, file string) (string, error) {
	obj, err := c.GetObject(repo, fmt.Sprintf("%s/%s", branch, file))
	if err != nil {
		return "", err
	}
	defer func() { err = obj.Close() }()
	bytes, err := ioutil.ReadAll(obj)
	if err != nil {
		return "", err
	}
	return string(bytes), err
}

func TestGetFile(t *testing.T) {
	repo := tu.UniqueString("testgetfile")
	pc := server.GetPachClient(t)
	require.NoError(t, pc.CreateRepo(repo))
	_, err := pc.PutFile(repo, "master", "file", strings.NewReader("content"))
	require.NoError(t, err)

	srv, port := serve(t, pc)
	c, err := minio.New(fmt.Sprintf("127.0.0.1:%d", port), "id", "secret", false)
	require.NoError(t, err)

	fetchedContent, err := getObject(c, repo, "master", "file")
	require.NoError(t, err)
	require.Equal(t, "content", fetchedContent)

	require.NoError(t, srv.Close())
}

func TestGetFileInBranch(t *testing.T) {
	repo := tu.UniqueString("testgetfileinbranch")
	pc := server.GetPachClient(t)
	require.NoError(t, pc.CreateRepo(repo))
	require.NoError(t, pc.CreateBranch(repo, "branch", "", nil))
	_, err := pc.PutFile(repo, "branch", "file", strings.NewReader("content"))
	require.NoError(t, err)

	srv, port := serve(t, pc)
	c, err := minio.New(fmt.Sprintf("127.0.0.1:%d", port), "id", "secret", false)
	require.NoError(t, err)

	fetchedContent, err := getObject(c, repo, "branch", "file")
	require.NoError(t, err)
	require.Equal(t, "content", fetchedContent)

	require.NoError(t, srv.Close())
}

func TestNonExistingBranch(t *testing.T) {
	repo := tu.UniqueString("testnonexistingbranch")
	pc := server.GetPachClient(t)
	require.NoError(t, pc.CreateRepo(repo))

	srv, port := serve(t, pc)
	c, err := minio.New(fmt.Sprintf("127.0.0.1:%d", port), "id", "secret", false)
	require.NoError(t, err)

	_, err = getObject(c, repo, "branch", "file")
	require.YesError(t, err)
	require.Equal(t, err.Error(), "The specified key does not exist.")
	require.NoError(t, srv.Close())
}

func TestNonExistingRepo(t *testing.T) {
	repo := tu.UniqueString("testnonexistingrepo")
	pc := server.GetPachClient(t)

	srv, port := serve(t, pc)
	c, err := minio.New(fmt.Sprintf("127.0.0.1:%d", port), "id", "secret", false)
	require.NoError(t, err)

	_, err = getObject(c, repo, "master", "file")
	require.YesError(t, err)
	require.Equal(t, err.Error(), "The specified bucket does not exist.")
	require.NoError(t, srv.Close())
}
<|endoftext|>


<commit_before>package handlers

//go:generate make

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

type SecurityType int64

const (
	Currency SecurityType = 1
	Stock                 = 2
)

func GetSecurityType(typestring string) SecurityType {
	if strings.EqualFold(typestring, "currency") {
		return Currency
	} else if strings.EqualFold(typestring, "stock") {
		return Stock
	} else {
		return 0
	}
}

type Security struct {
	SecurityId  int64
	UserId      int64
	Name        string
	Description string
	Symbol      string
	// Number of decimal digits (to the right of the decimal point) this
	// security is precise to
	Precision int
	Type      SecurityType
	// AlternateId is CUSIP for Type=Stock, ISO4217 for Type=Currency
	AlternateId string
}

type SecurityList struct {
	Securities *[]*Security `json:"securities"`
}

func (s *Security) Read(json_str string) error {
	dec := json.NewDecoder(strings.NewReader(json_str))
	return dec.Decode(s)
}

func (s *Security) Write(w http.ResponseWriter) error {
	enc := json.NewEncoder(w)
	return enc.Encode(s)
}

func (sl *SecurityList) Read(json_str string) error {
	dec := json.NewDecoder(strings.NewReader(json_str))
	return dec.Decode(sl)
}

func (sl *SecurityList) Write(w http.ResponseWriter) error {
	enc := json.NewEncoder(w)
	return enc.Encode(sl)
}

func SearchSecurityTemplates(search string, _type SecurityType, limit int64) []*Security {
	upperSearch := strings.ToUpper(search)
	var results []*Security
	for i, security := range SecurityTemplates {
		if strings.Contains(strings.ToUpper(security.Name), upperSearch) ||
			strings.Contains(strings.ToUpper(security.Description), upperSearch) ||
			strings.Contains(strings.ToUpper(security.Symbol), upperSearch) {
			if _type == 0 || _type == security.Type {
				results = append(results, &SecurityTemplates[i])
				if limit != -1 && int64(len(results)) >= limit {
					break
				}
			}
		}
	}
	return results
}

func FindSecurityTemplate(name string, _type SecurityType) *Security {
	for _, security := range SecurityTemplates {
		if name == security.Name && _type == security.Type {
			return &security
		}
	}
	return nil
}

func FindCurrencyTemplate(iso4217 int64) *Security {
	iso4217string := strconv.FormatInt(iso4217, 10)
	for _, security := range SecurityTemplates {
		if security.Type == Currency && security.AlternateId == iso4217string {
			return &security
		}
	}
	return nil
}

func GetSecurity(tx *Tx, securityid int64, userid int64) (*Security, error) {
	var s Security

	err := tx.SelectOne(&s, "SELECT * from securities where UserId=? AND SecurityId=?", userid, securityid)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func GetSecurities(tx *Tx, userid int64) (*[]*Security, error) {
	var securities []*Security

	_, err := tx.Select(&securities, "SELECT * from securities where UserId=?", userid)
	if err != nil {
		return nil, err
	}
	return &securities, nil
}

func InsertSecurity(tx *Tx, s *Security) error {
	err := tx.Insert(s)
	if err != nil {
		return err
	}
	return nil
}

func UpdateSecurity(tx *Tx, s *Security) (err error) {
	user, err := GetUser(tx, s.UserId)
	if err != nil {
		return
	} else if user.DefaultCurrency == s.SecurityId && s.Type != Currency {
		return errors.New("Cannot change security which is user's default currency to be non-currency")
	}

	count, err := tx.Update(s)
	if err != nil {
		return
	}
	if count > 1 {
		return fmt.Errorf("Updated %d securities (expected 1)", count)
	}

	return nil
}

type SecurityInUseError struct {
	message string
}

func (e SecurityInUseError) Error() string {
	return e.message
}

func DeleteSecurity(tx *Tx, s *Security) error {
	// First, ensure no accounts are using this security
	accounts, err := tx.SelectInt("SELECT count(*) from accounts where UserId=? and SecurityId=?", s.UserId, s.SecurityId)

	if accounts != 0 {
		return SecurityInUseError{"One or more accounts still use this security"}
	}

	user, err := GetUser(tx, s.UserId)
	if err != nil {
		return err
	} else if user.DefaultCurrency == s.SecurityId {
		return SecurityInUseError{"Cannot delete security which is user's default currency"}
	}

	// Remove all prices involving this security (either of this security, or
	// using it as a currency)
	_, err = tx.Exec("DELETE FROM prices WHERE SecurityId=? OR CurrencyId=?", s.SecurityId, s.SecurityId)
	if err != nil {
		return err
	}

	count, err := tx.Delete(s)
	if err != nil {
		return err
	}
	if count != 1 {
		return errors.New("Deleted more than one security")
	}

	return nil
}

func ImportGetCreateSecurity(tx *Tx, userid int64, security *Security) (*Security, error) {
	security.UserId = userid
	if len(security.AlternateId) == 0 {
		// Always create a new local security if we can't match on the AlternateId
		err := InsertSecurity(tx, security)
		if err != nil {
			return nil, err
		}
		return security, nil
	}

	var securities []*Security

	_, err := tx.Select(&securities, "SELECT * from securities where UserId=? AND Type=? AND AlternateId=? AND Precision=?", userid, security.Type, security.AlternateId, security.Precision)
	if err != nil {
		return nil, err
	}

	// First try to find a case insensitive match on the name or symbol
	upperName := strings.ToUpper(security.Name)
	upperSymbol := strings.ToUpper(security.Symbol)
	for _, s := range securities {
		if (len(s.Name) > 0 && strings.ToUpper(s.Name) == upperName) ||
			(len(s.Symbol) > 0 && strings.ToUpper(s.Symbol) == upperSymbol) {
			return s, nil
		}
	}
	//		if strings.Contains(strings.ToUpper(security.Name), upperSearch) ||

	// Try to find a partial string match on the name or symbol
	for _, s := range securities {
		sUpperName := strings.ToUpper(s.Name)
		sUpperSymbol := strings.ToUpper(s.Symbol)
		if (len(upperName) > 0 && len(s.Name) > 0 && (strings.Contains(upperName, sUpperName) || strings.Contains(sUpperName, upperName))) ||
			(len(upperSymbol) > 0 && len(s.Symbol) > 0 && (strings.Contains(upperSymbol, sUpperSymbol) || strings.Contains(sUpperSymbol, upperSymbol))) {
			return s, nil
		}
	}

	// Give up and return the first security in the list
	if len(securities) > 0 {
		return securities[0], nil
	}

	// If there wasn't even one security in the list, make a new one
	err = InsertSecurity(tx, security)
	if err != nil {
		return nil, err
	}

	return security, nil
}

func SecurityHandler(r *http.Request, context *Context) ResponseWriterWriter {
	user, err := GetUserFromSession(context.Tx, r)
	if err != nil {
		return NewError(1 /*Not Signed In*/)
	}

	if r.Method == "POST" {
		if !context.LastLevel() {
			securityid, err := context.NextID()
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}
			if context.NextLevel() != "prices" {
				return NewError(3 /*Invalid Request*/)
			}
			return PriceHandler(r, context, user, securityid)
		}

		var security Security
		if err := ReadJSON(r, &security); err != nil {
			return NewError(3 /*Invalid Request*/)
		}
		security.SecurityId = -1
		security.UserId = user.UserId

		err = InsertSecurity(context.Tx, &security)
		if err != nil {
			log.Print(err)
			return NewError(999 /*Internal Error*/)
		}

		return ResponseWrapper{201, &security}
	} else if r.Method == "GET" {
		if context.LastLevel() {
			//Return all securities
			var sl SecurityList

			securities, err := GetSecurities(context.Tx, user.UserId)
			if err != nil {
				log.Print(err)
				return NewError(999 /*Internal Error*/)
			}

			sl.Securities = securities
			return &sl
		} else {
			securityid, err := context.NextID()
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}

			if !context.LastLevel() {
				if context.NextLevel() != "prices" {
					return NewError(3 /*Invalid Request*/)
				}
				return PriceHandler(r, context, user, securityid)
			}

			security, err := GetSecurity(context.Tx, securityid, user.UserId)
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}

			return security
		}
	} else {
		securityid, err := context.NextID()
		if err != nil {
			return NewError(3 /*Invalid Request*/)
		}
		if !context.LastLevel() {
			if context.NextLevel() != "prices" {
				return NewError(3 /*Invalid Request*/)
			}
			return PriceHandler(r, context, user, securityid)
		}

		if r.Method == "PUT" {
			var security Security
			if err := ReadJSON(r, &security); err != nil || security.SecurityId != securityid {
				return NewError(3 /*Invalid Request*/)
			}
			security.UserId = user.UserId

			err = UpdateSecurity(context.Tx, &security)
			if err != nil {
				log.Print(err)
				return NewError(999 /*Internal Error*/)
			}

			return &security
		} else if r.Method == "DELETE" {
			security, err := GetSecurity(context.Tx, securityid, user.UserId)
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}

			err = DeleteSecurity(context.Tx, security)
			if _, ok := err.(SecurityInUseError); ok {
				return NewError(7 /*In Use Error*/)
			} else if err != nil {
				log.Print(err)
				return NewError(999 /*Internal Error*/)
			}

			return SuccessWriter{}
		}
	}
	return NewError(3 /*Invalid Request*/)
}

func SecurityTemplateHandler(r *http.Request, context *Context) ResponseWriterWriter {
	if r.Method == "GET" {
		var sl SecurityList

		query, _ := url.ParseQuery(r.URL.RawQuery)

		var limit int64 = -1
		search := query.Get("search")

		var _type SecurityType = 0
		typestring := query.Get("type")
		if len(typestring) > 0 {
			_type = GetSecurityType(typestring)
			if _type == 0 {
				return NewError(3 /*Invalid Request*/)
			}
		}

		limitstring := query.Get("limit")
		if limitstring != "" {
			limitint, err := strconv.ParseInt(limitstring, 10, 0)
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}
			limit = limitint
		}

		securities := SearchSecurityTemplates(search, _type, limit)

		sl.Securities = &securities
		return &sl
	} else {
		return NewError(3 /*Invalid Request*/)
	}
}
<commit_msg>securities: Don't use 'precision', a MySQL reserved word, in DB<commit_after>package handlers

//go:generate make

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

type SecurityType int64

const (
	Currency SecurityType = 1
	Stock                 = 2
)

func GetSecurityType(typestring string) SecurityType {
	if strings.EqualFold(typestring, "currency") {
		return Currency
	} else if strings.EqualFold(typestring, "stock") {
		return Stock
	} else {
		return 0
	}
}

type Security struct {
	SecurityId  int64
	UserId      int64
	Name        string
	Description string
	Symbol      string
	// Number of decimal digits (to the right of the decimal point) this
	// security is precise to
	Precision int `db:"Preciseness"`
	Type      SecurityType
	// AlternateId is CUSIP for Type=Stock, ISO4217 for Type=Currency
	AlternateId string
}

type SecurityList struct {
	Securities *[]*Security `json:"securities"`
}

func (s *Security) Read(json_str string) error {
	dec := json.NewDecoder(strings.NewReader(json_str))
	return dec.Decode(s)
}

func (s *Security) Write(w http.ResponseWriter) error {
	enc := json.NewEncoder(w)
	return enc.Encode(s)
}

func (sl *SecurityList) Read(json_str string) error {
	dec := json.NewDecoder(strings.NewReader(json_str))
	return dec.Decode(sl)
}

func (sl *SecurityList) Write(w http.ResponseWriter) error {
	enc := json.NewEncoder(w)
	return enc.Encode(sl)
}

func SearchSecurityTemplates(search string, _type SecurityType, limit int64) []*Security {
	upperSearch := strings.ToUpper(search)
	var results []*Security
	for i, security := range SecurityTemplates {
		if strings.Contains(strings.ToUpper(security.Name), upperSearch) ||
			strings.Contains(strings.ToUpper(security.Description), upperSearch) ||
			strings.Contains(strings.ToUpper(security.Symbol), upperSearch) {
			if _type == 0 || _type == security.Type {
				results = append(results, &SecurityTemplates[i])
				if limit != -1 && int64(len(results)) >= limit {
					break
				}
			}
		}
	}
	return results
}

func FindSecurityTemplate(name string, _type SecurityType) *Security {
	for _, security := range SecurityTemplates {
		if name == security.Name && _type == security.Type {
			return &security
		}
	}
	return nil
}

func FindCurrencyTemplate(iso4217 int64) *Security {
	iso4217string := strconv.FormatInt(iso4217, 10)
	for _, security := range SecurityTemplates {
		if security.Type == Currency && security.AlternateId == iso4217string {
			return &security
		}
	}
	return nil
}

func GetSecurity(tx *Tx, securityid int64, userid int64) (*Security, error) {
	var s Security

	err := tx.SelectOne(&s, "SELECT * from securities where UserId=? AND SecurityId=?", userid, securityid)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func GetSecurities(tx *Tx, userid int64) (*[]*Security, error) {
	var securities []*Security

	_, err := tx.Select(&securities, "SELECT * from securities where UserId=?", userid)
	if err != nil {
		return nil, err
	}
	return &securities, nil
}

func InsertSecurity(tx *Tx, s *Security) error {
	err := tx.Insert(s)
	if err != nil {
		return err
	}
	return nil
}

func UpdateSecurity(tx *Tx, s *Security) (err error) {
	user, err := GetUser(tx, s.UserId)
	if err != nil {
		return
	} else if user.DefaultCurrency == s.SecurityId && s.Type != Currency {
		return errors.New("Cannot change security which is user's default currency to be non-currency")
	}

	count, err := tx.Update(s)
	if err != nil {
		return
	}
	if count > 1 {
		return fmt.Errorf("Updated %d securities (expected 1)", count)
	}

	return nil
}

type SecurityInUseError struct {
	message string
}

func (e SecurityInUseError) Error() string {
	return e.message
}

func DeleteSecurity(tx *Tx, s *Security) error {
	// First, ensure no accounts are using this security
	accounts, err := tx.SelectInt("SELECT count(*) from accounts where UserId=? and SecurityId=?", s.UserId, s.SecurityId)

	if accounts != 0 {
		return SecurityInUseError{"One or more accounts still use this security"}
	}

	user, err := GetUser(tx, s.UserId)
	if err != nil {
		return err
	} else if user.DefaultCurrency == s.SecurityId {
		return SecurityInUseError{"Cannot delete security which is user's default currency"}
	}

	// Remove all prices involving this security (either of this security, or
	// using it as a currency)
	_, err = tx.Exec("DELETE FROM prices WHERE SecurityId=? OR CurrencyId=?", s.SecurityId, s.SecurityId)
	if err != nil {
		return err
	}

	count, err := tx.Delete(s)
	if err != nil {
		return err
	}
	if count != 1 {
		return errors.New("Deleted more than one security")
	}

	return nil
}

func ImportGetCreateSecurity(tx *Tx, userid int64, security *Security) (*Security, error) {
	security.UserId = userid
	if len(security.AlternateId) == 0 {
		// Always create a new local security if we can't match on the AlternateId
		err := InsertSecurity(tx, security)
		if err != nil {
			return nil, err
		}
		return security, nil
	}

	var securities []*Security

	_, err := tx.Select(&securities, "SELECT * from securities where UserId=? AND Type=? AND AlternateId=? AND Preciseness=?", userid, security.Type, security.AlternateId, security.Precision)
	if err != nil {
		return nil, err
	}

	// First try to find a case insensitive match on the name or symbol
	upperName := strings.ToUpper(security.Name)
	upperSymbol := strings.ToUpper(security.Symbol)
	for _, s := range securities {
		if (len(s.Name) > 0 && strings.ToUpper(s.Name) == upperName) ||
			(len(s.Symbol) > 0 && strings.ToUpper(s.Symbol) == upperSymbol) {
			return s, nil
		}
	}
	//		if strings.Contains(strings.ToUpper(security.Name), upperSearch) ||

	// Try to find a partial string match on the name or symbol
	for _, s := range securities {
		sUpperName := strings.ToUpper(s.Name)
		sUpperSymbol := strings.ToUpper(s.Symbol)
		if (len(upperName) > 0 && len(s.Name) > 0 && (strings.Contains(upperName, sUpperName) || strings.Contains(sUpperName, upperName))) ||
			(len(upperSymbol) > 0 && len(s.Symbol) > 0 && (strings.Contains(upperSymbol, sUpperSymbol) || strings.Contains(sUpperSymbol, upperSymbol))) {
			return s, nil
		}
	}

	// Give up and return the first security in the list
	if len(securities) > 0 {
		return securities[0], nil
	}

	// If there wasn't even one security in the list, make a new one
	err = InsertSecurity(tx, security)
	if err != nil {
		return nil, err
	}

	return security, nil
}

func SecurityHandler(r *http.Request, context *Context) ResponseWriterWriter {
	user, err := GetUserFromSession(context.Tx, r)
	if err != nil {
		return NewError(1 /*Not Signed In*/)
	}

	if r.Method == "POST" {
		if !context.LastLevel() {
			securityid, err := context.NextID()
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}
			if context.NextLevel() != "prices" {
				return NewError(3 /*Invalid Request*/)
			}
			return PriceHandler(r, context, user, securityid)
		}

		var security Security
		if err := ReadJSON(r, &security); err != nil {
			return NewError(3 /*Invalid Request*/)
		}
		security.SecurityId = -1
		security.UserId = user.UserId

		err = InsertSecurity(context.Tx, &security)
		if err != nil {
			log.Print(err)
			return NewError(999 /*Internal Error*/)
		}

		return ResponseWrapper{201, &security}
	} else if r.Method == "GET" {
		if context.LastLevel() {
			//Return all securities
			var sl SecurityList

			securities, err := GetSecurities(context.Tx, user.UserId)
			if err != nil {
				log.Print(err)
				return NewError(999 /*Internal Error*/)
			}

			sl.Securities = securities
			return &sl
		} else {
			securityid, err := context.NextID()
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}

			if !context.LastLevel() {
				if context.NextLevel() != "prices" {
					return NewError(3 /*Invalid Request*/)
				}
				return PriceHandler(r, context, user, securityid)
			}

			security, err := GetSecurity(context.Tx, securityid, user.UserId)
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}

			return security
		}
	} else {
		securityid, err := context.NextID()
		if err != nil {
			return NewError(3 /*Invalid Request*/)
		}
		if !context.LastLevel() {
			if context.NextLevel() != "prices" {
				return NewError(3 /*Invalid Request*/)
			}
			return PriceHandler(r, context, user, securityid)
		}

		if r.Method == "PUT" {
			var security Security
			if err := ReadJSON(r, &security); err != nil || security.SecurityId != securityid {
				return NewError(3 /*Invalid Request*/)
			}
			security.UserId = user.UserId

			err = UpdateSecurity(context.Tx, &security)
			if err != nil {
				log.Print(err)
				return NewError(999 /*Internal Error*/)
			}

			return &security
		} else if r.Method == "DELETE" {
			security, err := GetSecurity(context.Tx, securityid, user.UserId)
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}

			err = DeleteSecurity(context.Tx, security)
			if _, ok := err.(SecurityInUseError); ok {
				return NewError(7 /*In Use Error*/)
			} else if err != nil {
				log.Print(err)
				return NewError(999 /*Internal Error*/)
			}

			return SuccessWriter{}
		}
	}
	return NewError(3 /*Invalid Request*/)
}

func SecurityTemplateHandler(r *http.Request, context *Context) ResponseWriterWriter {
	if r.Method == "GET" {
		var sl SecurityList

		query, _ := url.ParseQuery(r.URL.RawQuery)

		var limit int64 = -1
		search := query.Get("search")

		var _type SecurityType = 0
		typestring := query.Get("type")
		if len(typestring) > 0 {
			_type = GetSecurityType(typestring)
			if _type == 0 {
				return NewError(3 /*Invalid Request*/)
			}
		}

		limitstring := query.Get("limit")
		if limitstring != "" {
			limitint, err := strconv.ParseInt(limitstring, 10, 0)
			if err != nil {
				return NewError(3 /*Invalid Request*/)
			}
			limit = limitint
		}

		securities := SearchSecurityTemplates(search, _type, limit)

		sl.Securities = &securities
		return &sl
	} else {
		return NewError(3 /*Invalid Request*/)
	}
}
<|endoftext|>