<commit_before>package bundler

// Bundling test results are stored in snapshot files, located in the
// "snapshots" directory. This allows test results to be updated easily without
// manually rewriting all of the expected values. To update the tests run
// "UPDATE_SNAPSHOTS=1 make test" and commit the updated values. Make sure to
// inspect the diff to ensure the expected values are valid.

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"testing"

	"github.com/evanw/esbuild/internal/compat"
	"github.com/evanw/esbuild/internal/config"
	"github.com/evanw/esbuild/internal/fs"
	"github.com/evanw/esbuild/internal/logging"
	"github.com/evanw/esbuild/internal/resolver"
	"github.com/kylelemons/godebug/diff"
)

func es(version int) compat.Feature {
	return compat.UnsupportedFeatures(map[compat.Engine][]int{
		compat.ES: {version},
	})
}

func assertEqual(t *testing.T, a interface{}, b interface{}) {
	if a != b {
		stringA := fmt.Sprintf("%v", a)
		stringB := fmt.Sprintf("%v", b)
		if strings.Contains(stringA, "\n") {
			t.Fatal(diff.Diff(stringB, stringA))
		} else {
			t.Fatalf("%s != %s", a, b)
		}
	}
}

func assertLog(t *testing.T, msgs []logging.Msg, expected string) {
	text := ""
	for _, msg := range msgs {
		text += msg.String(logging.StderrOptions{}, logging.TerminalInfo{})
	}
	assertEqual(t, text, expected)
}

func hasErrors(msgs []logging.Msg) bool {
	for _, msg := range msgs {
		if msg.Kind == logging.Error {
			return true
		}
	}
	return false
}

type bundled struct {
	files              map[string]string
	entryPaths         []string
	expectedScanLog    string
	expectedCompileLog string
	options            config.Options
}

type suite struct {
	name               string
	path               string
	mutex              sync.Mutex
	expectedSnapshots  map[string]string
	generatedSnapshots map[string]string
}

func (s *suite) expectBundled(t *testing.T, args bundled) {
	testName := t.Name()
	t.Run("", func(t *testing.T) {
		fs := fs.MockFS(args.files)
		args.options.ExtensionOrder = []string{".tsx", ".ts", ".jsx", ".js", ".json"}
		if args.options.AbsOutputFile != "" {
			args.options.AbsOutputDir = path.Dir(args.options.AbsOutputFile)
		}
		log := logging.NewDeferLog()
		resolver := resolver.NewResolver(fs, log, args.options)
		bundle := ScanBundle(log, fs, resolver, args.entryPaths, args.options)
		msgs := log.Done()
		assertLog(t, msgs, args.expectedScanLog)

		// Stop now if there were any errors during the scan
		if hasErrors(msgs) {
			return
		}

		log = logging.NewDeferLog()
		args.options.OmitRuntimeForTests = true
		results := bundle.Compile(log, args.options)
		msgs = log.Done()
		assertLog(t, msgs, args.expectedCompileLog)

		// Stop now if there were any errors during the compile
		if hasErrors(msgs) {
			return
		}

		// Don't include source maps in results since they are just noise. Source
		// map validity is tested separately in a test that uses Mozilla's source
		// map parsing library.
		generated := ""
		for _, result := range results {
			if !strings.HasSuffix(result.AbsPath, ".map") {
				if generated != "" {
					generated += "\n"
				}
				generated += fmt.Sprintf("---------- %s ----------\n%s", result.AbsPath, string(result.Contents))
			}
		}
		s.compareSnapshot(t, testName, generated)
	})
}

const snapshotsDir = "snapshots"
const snapshotSplitter = "\n================================================================================\n"

var globalTestMutex sync.Mutex
var globalSuites map[*suite]bool
var globalUpdateSnapshots bool

func (s *suite) compareSnapshot(t *testing.T, testName string, generated string) {
	// Initialize the test suite during the first test
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if s.path == "" {
		s.path = snapshotsDir + "/snapshots_" + s.name + ".txt"
		s.generatedSnapshots = make(map[string]string)
		s.expectedSnapshots = make(map[string]string)
		if contents, err := ioutil.ReadFile(s.path); err == nil {
			for _, part := range strings.Split(string(contents), snapshotSplitter) {
				if newline := strings.IndexByte(part, '\n'); newline != -1 {
					key := part[:newline]
					value := part[newline+1:]
					s.expectedSnapshots[key] = value
				} else {
					s.expectedSnapshots[part] = ""
				}
			}
		}
		globalTestMutex.Lock()
		defer globalTestMutex.Unlock()
		if globalSuites == nil {
			globalSuites = make(map[*suite]bool)
		}
		globalSuites[s] = true
		_, globalUpdateSnapshots = os.LookupEnv("UPDATE_SNAPSHOTS")
	}

	// Check against the stored snapshot if present
	s.generatedSnapshots[testName] = generated
	if !globalUpdateSnapshots {
		if expected, ok := s.expectedSnapshots[testName]; ok {
			assertEqual(t, generated, expected)
		} else {
			t.Fatalf("No snapshot saved for %s\n%s", testName, generated)
		}
	}
}

func (s *suite) updateSnapshots() {
	os.Mkdir(snapshotsDir, 0755)
	keys := make([]string, 0, len(s.generatedSnapshots))
	for key := range s.generatedSnapshots {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	contents := ""
	for i, key := range keys {
		if i > 0 {
			contents += snapshotSplitter
		}
		contents += fmt.Sprintf("%s\n%s", key, s.generatedSnapshots[key])
	}
	if err := ioutil.WriteFile(s.path, []byte(contents), 0644); err != nil {
		panic(err)
	}
}

func (s *suite) validateSnapshots() bool {
	isValid := true
	for key := range s.expectedSnapshots {
		if _, ok := s.generatedSnapshots[key]; !ok {
			if isValid {
				fmt.Printf("%s\n", s.path)
			}
			fmt.Printf("    No test found for snapshot %s\n", key)
			isValid = false
		}
	}
	return isValid
}

func TestMain(m *testing.M) {
	code := m.Run()
	if globalSuites != nil {
		if globalUpdateSnapshots {
			for s := range globalSuites {
				s.updateSnapshots()
			}
		} else {
			for s := range globalSuites {
				if !s.validateSnapshots() {
					code = 1
				}
			}
		}
	}
	os.Exit(code)
}
<commit_msg>fix windows-only ci issues due to crlf<commit_after>package bundler

// Bundling test results are stored in snapshot files, located in the
// "snapshots" directory. This allows test results to be updated easily without
// manually rewriting all of the expected values. To update the tests run
// "UPDATE_SNAPSHOTS=1 make test" and commit the updated values. Make sure to
// inspect the diff to ensure the expected values are valid.

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"testing"

	"github.com/evanw/esbuild/internal/compat"
	"github.com/evanw/esbuild/internal/config"
	"github.com/evanw/esbuild/internal/fs"
	"github.com/evanw/esbuild/internal/logging"
	"github.com/evanw/esbuild/internal/resolver"
	"github.com/kylelemons/godebug/diff"
)

func es(version int) compat.Feature {
	return compat.UnsupportedFeatures(map[compat.Engine][]int{
		compat.ES: {version},
	})
}

func assertEqual(t *testing.T, a interface{}, b interface{}) {
	if a != b {
		stringA := fmt.Sprintf("%v", a)
		stringB := fmt.Sprintf("%v", b)
		if strings.Contains(stringA, "\n") {
			t.Fatal(diff.Diff(stringB, stringA))
		} else {
			t.Fatalf("%s != %s", a, b)
		}
	}
}

func assertLog(t *testing.T, msgs []logging.Msg, expected string) {
	text := ""
	for _, msg := range msgs {
		text += msg.String(logging.StderrOptions{}, logging.TerminalInfo{})
	}
	assertEqual(t, text, expected)
}

func hasErrors(msgs []logging.Msg) bool {
	for _, msg := range msgs {
		if msg.Kind == logging.Error {
			return true
		}
	}
	return false
}

type bundled struct {
	files              map[string]string
	entryPaths         []string
	expectedScanLog    string
	expectedCompileLog string
	options            config.Options
}

type suite struct {
	name               string
	path               string
	mutex              sync.Mutex
	expectedSnapshots  map[string]string
	generatedSnapshots map[string]string
}

func (s *suite) expectBundled(t *testing.T, args bundled) {
	testName := t.Name()
	t.Run("", func(t *testing.T) {
		fs := fs.MockFS(args.files)
		args.options.ExtensionOrder = []string{".tsx", ".ts", ".jsx", ".js", ".json"}
		if args.options.AbsOutputFile != "" {
			args.options.AbsOutputDir = path.Dir(args.options.AbsOutputFile)
		}
		log := logging.NewDeferLog()
		resolver := resolver.NewResolver(fs, log, args.options)
		bundle := ScanBundle(log, fs, resolver, args.entryPaths, args.options)
		msgs := log.Done()
		assertLog(t, msgs, args.expectedScanLog)

		// Stop now if there were any errors during the scan
		if hasErrors(msgs) {
			return
		}

		log = logging.NewDeferLog()
		args.options.OmitRuntimeForTests = true
		results := bundle.Compile(log, args.options)
		msgs = log.Done()
		assertLog(t, msgs, args.expectedCompileLog)

		// Stop now if there were any errors during the compile
		if hasErrors(msgs) {
			return
		}

		// Don't include source maps in results since they are just noise. Source
		// map validity is tested separately in a test that uses Mozilla's source
		// map parsing library.
		generated := ""
		for _, result := range results {
			if !strings.HasSuffix(result.AbsPath, ".map") {
				if generated != "" {
					generated += "\n"
				}
				generated += fmt.Sprintf("---------- %s ----------\n%s", result.AbsPath, string(result.Contents))
			}
		}
		s.compareSnapshot(t, testName, generated)
	})
}

const snapshotsDir = "snapshots"
const snapshotSplitter = "\n================================================================================\n"

var globalTestMutex sync.Mutex
var globalSuites map[*suite]bool
var globalUpdateSnapshots bool

func (s *suite) compareSnapshot(t *testing.T, testName string, generated string) {
	// Initialize the test suite during the first test
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if s.path == "" {
		s.path = snapshotsDir + "/snapshots_" + s.name + ".txt"
		s.generatedSnapshots = make(map[string]string)
		s.expectedSnapshots = make(map[string]string)
		if contents, err := ioutil.ReadFile(s.path); err == nil {
			// Replacing CRLF with LF is necessary to fix tests in GitHub actions,
			// which for some reason check out the source code in CLRF mode
			for _, part := range strings.Split(strings.ReplaceAll(string(contents), "\r\n", "\n"), snapshotSplitter) {
				if newline := strings.IndexByte(part, '\n'); newline != -1 {
					key := part[:newline]
					value := part[newline+1:]
					s.expectedSnapshots[key] = value
				} else {
					s.expectedSnapshots[part] = ""
				}
			}
		}
		globalTestMutex.Lock()
		defer globalTestMutex.Unlock()
		if globalSuites == nil {
			globalSuites = make(map[*suite]bool)
		}
		globalSuites[s] = true
		_, globalUpdateSnapshots = os.LookupEnv("UPDATE_SNAPSHOTS")
	}

	// Check against the stored snapshot if present
	s.generatedSnapshots[testName] = generated
	if !globalUpdateSnapshots {
		if expected, ok := s.expectedSnapshots[testName]; ok {
			assertEqual(t, generated, expected)
		} else {
			t.Fatalf("No snapshot saved for %s\n%s", testName, generated)
		}
	}
}

func (s *suite) updateSnapshots() {
	os.Mkdir(snapshotsDir, 0755)
	keys := make([]string, 0, len(s.generatedSnapshots))
	for key := range s.generatedSnapshots {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	contents := ""
	for i, key := range keys {
		if i > 0 {
			contents += snapshotSplitter
		}
		contents += fmt.Sprintf("%s\n%s", key, s.generatedSnapshots[key])
	}
	if err := ioutil.WriteFile(s.path, []byte(contents), 0644); err != nil {
		panic(err)
	}
}

func (s *suite) validateSnapshots() bool {
	isValid := true
	for key := range s.expectedSnapshots {
		if _, ok := s.generatedSnapshots[key]; !ok {
			if isValid {
				fmt.Printf("%s\n", s.path)
			}
			fmt.Printf("    No test found for snapshot %s\n", key)
			isValid = false
		}
	}
	return isValid
}

func TestMain(m *testing.M) {
	code := m.Run()
	if globalSuites != nil {
		if globalUpdateSnapshots {
			for s := range globalSuites {
				s.updateSnapshots()
			}
		} else {
			for s := range globalSuites {
				if !s.validateSnapshots() {
					code = 1
				}
			}
		}
	}
	os.Exit(code)
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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
//
//     https://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 recon

import (
	"context"
	"fmt"
	"sort"
	"strings"

	pb "github.com/datacommonsorg/mixer/internal/proto"
	"github.com/datacommonsorg/mixer/internal/store"
	"github.com/datacommonsorg/mixer/internal/store/bigtable"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
)

var (
	// This is a preferred list.  The props ranked higher are preferred over
	// those ranked lower for resolving.
	//
	// NOTE: Needs to be kept in sync with PLACE_RESOLVABLE_AND_ASSIGNABLE_IDS in
	// https://github.com/datacommonsorg/import repo.
	rankedIDProps = []string{
		"dcid",
		"geoId",
		"isoCode",
		"nutsCode",
		"wikidataId",
		"geoNamesId",
		"istatId",
		"austrianMunicipalityKey",
		"indianCensusAreaCode2011",
		"indianCensusAreaCode2001",
		"lgdCode",
		"udiseCode",
		"fips52AlphaCode",
	}
)

// Get the value of a given property, assuming single value.
func getPropVal(node *pb.McfGraph_PropertyValues, prop string) string {
	values, ok := (node.GetPvs())[prop]
	if !ok {
		return ""
	}
	typedValues := values.GetTypedValues()
	if len(typedValues) == 0 {
		return ""
	}
	return typedValues[0].GetValue()
}

// ResolveEntities implements API for ReconServer.ResolveEntities.
func ResolveEntities(
	ctx context.Context, in *pb.ResolveEntitiesRequest, store *store.Store,
) (
	*pb.ResolveEntitiesResponse, error) {
	idKeyToSourceIDs := map[string][]string{}
	sourceIDs := map[string]struct{}{}
	idKeys := []string{}

	// Collect to-be-resolved IDs to rowList and idKeyToSourceID.
	for _, entity := range in.GetEntities() {
		sourceID := entity.GetSourceId()

		// Try to resolve all the supported IDs
		// For the resolved ones, only rely on the one ranked higher.
		switch t := entity.GraphRepresentation.(type) {
		case *pb.EntitySubGraph_SubGraph:
			node, ok := (entity.GetSubGraph().GetNodes())[sourceID]
			if !ok {
				continue
			}
			for _, idProp := range rankedIDProps {
				idVal := getPropVal(node, idProp)
				if idVal == "" {
					continue
				}
				idKey := fmt.Sprintf("%s^%s", idProp, idVal)
				idKeys = append(idKeys, idKey)
				idKeyToSourceIDs[idKey] = append(idKeyToSourceIDs[idKey], sourceID)
			}
		case *pb.EntitySubGraph_EntityIds:
			idStore := map[string]string{} // Map: ID prop -> ID val.
			for _, id := range entity.GetEntityIds().GetIds() {
				idStore[id.GetProp()] = id.GetVal()
			}
			for _, idProp := range rankedIDProps {
				idVal, ok := idStore[idProp]
				if !ok {
					continue
				}
				idKey := fmt.Sprintf("%s^%s", idProp, idVal)
				idKeys = append(idKeys, idKey)
				idKeyToSourceIDs[idKey] = append(idKeyToSourceIDs[idKey], sourceID)
			}
		default:
			return nil, fmt.Errorf("Entity.GraphRepresentation has unexpected type %T", t)
		}
		sourceIDs[sourceID] = struct{}{}
	}

	// Read ReconIdMap cache.
	btDataList, err := bigtable.Read(
		ctx,
		store.BtGroup,
		bigtable.BtReconIDMapPrefix,
		[][]string{idKeys},
		func(jsonRaw []byte) (interface{}, error) {
			var reconEntities pb.ReconEntities
			if err := proto.Unmarshal(jsonRaw, &reconEntities); err != nil {
				return nil, err
			}
			return &reconEntities, nil
		},
	)
	if err != nil {
		return nil, err
	}

	// Source ID -> ID Prop -> ReconEntities.
	reconEntityStore := map[string]map[string]*pb.ReconEntities{}

	// Group resolving cache result by source ID.
	for _, btData := range btDataList {
		if len(btData) == 0 {
			continue
		}
		for _, row := range btData {
			idKey := fmt.Sprintf("%s^%s", row.Parts[0], row.Parts[1])
			reconEntities := row.Data
			if reconEntities == nil {
				continue
			}
			sourceIDs, ok := idKeyToSourceIDs[idKey]
			if !ok {
				continue
			}
			parts := strings.Split(idKey, "^")
			if len(parts) != 2 {
				return nil, status.Errorf(codes.Internal, "Invalid id key %s", idKey)
			}
			idProp := parts[0]

			for _, sourceID := range sourceIDs {
				if _, ok := reconEntityStore[sourceID]; !ok {
					reconEntityStore[sourceID] = map[string]*pb.ReconEntities{}
				}
				if re := reconEntities.(*pb.ReconEntities); len(re.GetEntities()) > 0 {
					reconEntityStore[sourceID][idProp] = re
				}
			}
		}
		// Only process data from one preferred import group.
		// TODO: merge entities from differente import groups.
		break
	}

	// Assemble response.
	res := &pb.ResolveEntitiesResponse{}
	for sourceID, idProp2ReconEntities := range reconEntityStore {
		var reconEntities *pb.ReconEntities
		for _, idProp := range rankedIDProps {
			if val, ok := idProp2ReconEntities[idProp]; ok {
				reconEntities = val
				break
			}
		}
		if reconEntities == nil {
			continue
		}

		// If it is resolved to multiple DC entities, each resolved entity has an equal probability.
		probability := float64(1.0 / len(reconEntities.GetEntities()))

		resolvedEntity := &pb.ResolveEntitiesResponse_ResolvedEntity{
			SourceId: sourceID,
		}

		for _, entity := range reconEntities.GetEntities() {
			resolvedID := &pb.ResolveEntitiesResponse_ResolvedId{
				Probability: probability,
			}
			for _, id := range entity.GetIds() {
				resolvedID.Ids = append(resolvedID.Ids,
					&pb.IdWithProperty{
						Prop: id.GetProp(),
						Val:  id.GetVal(),
					})
			}
			resolvedEntity.ResolvedIds = append(resolvedEntity.ResolvedIds, resolvedID)
		}

		res.ResolvedEntities = append(res.ResolvedEntities, resolvedEntity)
	}

	// Add entities that are not resolved as empty result.
	for sourceID := range sourceIDs {
		if _, ok := reconEntityStore[sourceID]; ok { // Resolved.
			continue
		}
		res.ResolvedEntities = append(res.ResolvedEntities,
			&pb.ResolveEntitiesResponse_ResolvedEntity{
				SourceId: sourceID,
			})
	}

	// Sort to make the result deterministic.
	sort.Slice(res.ResolvedEntities, func(i, j int) bool {
		return res.ResolvedEntities[i].GetSourceId() > res.ResolvedEntities[j].GetSourceId()
	})

	return res, nil
}
<commit_msg>Add support for countryAlpha3Code and countryNumericCode in resolve.go (#879)<commit_after>// Copyright 2022 Google LLC
//
// 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
//
//     https://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 recon

import (
	"context"
	"fmt"
	"sort"
	"strings"

	pb "github.com/datacommonsorg/mixer/internal/proto"
	"github.com/datacommonsorg/mixer/internal/store"
	"github.com/datacommonsorg/mixer/internal/store/bigtable"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
)

var (
	// This is a preferred list.  The props ranked higher are preferred over
	// those ranked lower for resolving.
	//
	// NOTE: Needs to be kept in sync with PLACE_RESOLVABLE_AND_ASSIGNABLE_IDS in
	// https://github.com/datacommonsorg/import repo.
	rankedIDProps = []string{
		"dcid",
		"geoId",
		"isoCode",
		"nutsCode",
		"wikidataId",
		"geoNamesId",
		"istatId",
		"austrianMunicipalityKey",
		"indianCensusAreaCode2011",
		"indianCensusAreaCode2001",
		"lgdCode",
		"udiseCode",
		"fips52AlphaCode",
		"countryAlpha3Code",
		"countryNumericCode",
	}
)

// Get the value of a given property, assuming single value.
func getPropVal(node *pb.McfGraph_PropertyValues, prop string) string {
	values, ok := (node.GetPvs())[prop]
	if !ok {
		return ""
	}
	typedValues := values.GetTypedValues()
	if len(typedValues) == 0 {
		return ""
	}
	return typedValues[0].GetValue()
}

// ResolveEntities implements API for ReconServer.ResolveEntities.
func ResolveEntities(
	ctx context.Context, in *pb.ResolveEntitiesRequest, store *store.Store,
) (
	*pb.ResolveEntitiesResponse, error) {
	idKeyToSourceIDs := map[string][]string{}
	sourceIDs := map[string]struct{}{}
	idKeys := []string{}

	// Collect to-be-resolved IDs to rowList and idKeyToSourceID.
	for _, entity := range in.GetEntities() {
		sourceID := entity.GetSourceId()

		// Try to resolve all the supported IDs
		// For the resolved ones, only rely on the one ranked higher.
		switch t := entity.GraphRepresentation.(type) {
		case *pb.EntitySubGraph_SubGraph:
			node, ok := (entity.GetSubGraph().GetNodes())[sourceID]
			if !ok {
				continue
			}
			for _, idProp := range rankedIDProps {
				idVal := getPropVal(node, idProp)
				if idVal == "" {
					continue
				}
				idKey := fmt.Sprintf("%s^%s", idProp, idVal)
				idKeys = append(idKeys, idKey)
				idKeyToSourceIDs[idKey] = append(idKeyToSourceIDs[idKey], sourceID)
			}
		case *pb.EntitySubGraph_EntityIds:
			idStore := map[string]string{} // Map: ID prop -> ID val.
			for _, id := range entity.GetEntityIds().GetIds() {
				idStore[id.GetProp()] = id.GetVal()
			}
			for _, idProp := range rankedIDProps {
				idVal, ok := idStore[idProp]
				if !ok {
					continue
				}
				idKey := fmt.Sprintf("%s^%s", idProp, idVal)
				idKeys = append(idKeys, idKey)
				idKeyToSourceIDs[idKey] = append(idKeyToSourceIDs[idKey], sourceID)
			}
		default:
			return nil, fmt.Errorf("Entity.GraphRepresentation has unexpected type %T", t)
		}
		sourceIDs[sourceID] = struct{}{}
	}

	// Read ReconIdMap cache.
	btDataList, err := bigtable.Read(
		ctx,
		store.BtGroup,
		bigtable.BtReconIDMapPrefix,
		[][]string{idKeys},
		func(jsonRaw []byte) (interface{}, error) {
			var reconEntities pb.ReconEntities
			if err := proto.Unmarshal(jsonRaw, &reconEntities); err != nil {
				return nil, err
			}
			return &reconEntities, nil
		},
	)
	if err != nil {
		return nil, err
	}

	// Source ID -> ID Prop -> ReconEntities.
	reconEntityStore := map[string]map[string]*pb.ReconEntities{}

	// Group resolving cache result by source ID.
	for _, btData := range btDataList {
		if len(btData) == 0 {
			continue
		}
		for _, row := range btData {
			idKey := fmt.Sprintf("%s^%s", row.Parts[0], row.Parts[1])
			reconEntities := row.Data
			if reconEntities == nil {
				continue
			}
			sourceIDs, ok := idKeyToSourceIDs[idKey]
			if !ok {
				continue
			}
			parts := strings.Split(idKey, "^")
			if len(parts) != 2 {
				return nil, status.Errorf(codes.Internal, "Invalid id key %s", idKey)
			}
			idProp := parts[0]

			for _, sourceID := range sourceIDs {
				if _, ok := reconEntityStore[sourceID]; !ok {
					reconEntityStore[sourceID] = map[string]*pb.ReconEntities{}
				}
				if re := reconEntities.(*pb.ReconEntities); len(re.GetEntities()) > 0 {
					reconEntityStore[sourceID][idProp] = re
				}
			}
		}
		// Only process data from one preferred import group.
		// TODO: merge entities from differente import groups.
		break
	}

	// Assemble response.
	res := &pb.ResolveEntitiesResponse{}
	for sourceID, idProp2ReconEntities := range reconEntityStore {
		var reconEntities *pb.ReconEntities
		for _, idProp := range rankedIDProps {
			if val, ok := idProp2ReconEntities[idProp]; ok {
				reconEntities = val
				break
			}
		}
		if reconEntities == nil {
			continue
		}

		// If it is resolved to multiple DC entities, each resolved entity has an equal probability.
		probability := float64(1.0 / len(reconEntities.GetEntities()))

		resolvedEntity := &pb.ResolveEntitiesResponse_ResolvedEntity{
			SourceId: sourceID,
		}

		for _, entity := range reconEntities.GetEntities() {
			resolvedID := &pb.ResolveEntitiesResponse_ResolvedId{
				Probability: probability,
			}
			for _, id := range entity.GetIds() {
				resolvedID.Ids = append(resolvedID.Ids,
					&pb.IdWithProperty{
						Prop: id.GetProp(),
						Val:  id.GetVal(),
					})
			}
			resolvedEntity.ResolvedIds = append(resolvedEntity.ResolvedIds, resolvedID)
		}

		res.ResolvedEntities = append(res.ResolvedEntities, resolvedEntity)
	}

	// Add entities that are not resolved as empty result.
	for sourceID := range sourceIDs {
		if _, ok := reconEntityStore[sourceID]; ok { // Resolved.
			continue
		}
		res.ResolvedEntities = append(res.ResolvedEntities,
			&pb.ResolveEntitiesResponse_ResolvedEntity{
				SourceId: sourceID,
			})
	}

	// Sort to make the result deterministic.
	sort.Slice(res.ResolvedEntities, func(i, j int) bool {
		return res.ResolvedEntities[i].GetSourceId() > res.ResolvedEntities[j].GetSourceId()
	})

	return res, nil
}
<|endoftext|>


<commit_before>// Copyright 2012 Arne Roomann-Kurrik
//
// 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 fauxfile

import (
	"errors"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

func GetPathError(path string, message string) error {
	return &os.PathError{
		Path: path,
		Err:  errors.New(message),
	}
}

type MockFilesystem struct {
	cwd  *MockFileInfo
	root *MockFileInfo
}

func NewMockFilesystem() *MockFilesystem {
	root := &MockFileInfo{
		name:       "/",
		filesystem: nil,
		mode:       os.ModeDir | 0755,
		modified:   time.Now(),
		data:       nil,
		parent:     nil,
		children:   map[string]*MockFileInfo{},
	}
	mf := &MockFilesystem{
		cwd:  root,
		root: root,
	}
	root.filesystem = mf
	return mf
}

func (mf *MockFilesystem) getpath(path string) string {
	if filepath.IsAbs(path) {
		return filepath.Clean(path)
	}
	return filepath.Join(mf.cwd.path(), path)
}

func (mf *MockFilesystem) resolve(path string) (*MockFileInfo, error) {
	path = mf.getpath(path)
	parts := strings.Split(path, string(filepath.Separator))
	ptr := mf.root
	for _, part := range parts {
		if part == "" {
			continue
		}
		if child := ptr.Child(part); child != nil {
			ptr = ptr.Child(part)
		} else {
			return nil, GetPathError(path, "Path does not exist")
		}
	}
	return ptr, nil
}

func (mf *MockFilesystem) exists(path string) bool {
	_, err := mf.resolve(path)
	return err == nil
}

func (mf *MockFilesystem) Chdir(dir string) error {
	fi, err := mf.resolve(dir)
	if err == nil {
		if fi.IsDir() {
			mf.cwd = fi
			return nil
		} else {
			return GetPathError(dir, "Path is not a directory")
		}
	}
	return err
}

func (mf *MockFilesystem) Mkdir(name string, perm os.FileMode) error {
	path := mf.getpath(name)
	parentpath, dirname := filepath.Split(path)
	fi, err := mf.resolve(parentpath)
	if err != nil {
		return err
	}
	if child := fi.Child(dirname); child != nil {
		return GetPathError(path, "Path already exists")
	}
	fi.children[dirname] = &MockFileInfo{
		name:       dirname,
		filesystem: mf,
		mode:       perm | os.ModeDir,
		modified:   time.Now(),
		data:       nil,
		parent:     fi,
		children:   map[string]*MockFileInfo{},
	}
	fi.modified = time.Now() // Update directory's timestamp.
	return nil
}

func (mf *MockFilesystem) MkdirAll(path string, perm os.FileMode) error {
	path = mf.getpath(path)
	parts := strings.Split(path, string(filepath.Separator))
	base := "/"
	for _, part := range parts {
		if part == "" {
			continue
		}
		base = filepath.Join(base, part)
		if err := mf.Mkdir(base, perm); err != nil {
			return err
		}
	}
	return nil
}

func (mf *MockFilesystem) Remove(name string) error {
	fi, err := mf.resolve(name)
	if err != nil {
		return err
	}
	if len(fi.Children()) > 0 {
		return GetPathError(name, "Directory contains children")
	}
	delete(fi.Parent().Children(), fi.name)
	fi.Parent().modified = time.Now()
	return nil
}

func (mf *MockFilesystem) RemoveAll(path string) error {
	fi, err := mf.resolve(path)
	if err != nil {
		return err
	}
	delete(fi.Parent().Children(), fi.name)
	fi.Parent().modified = time.Now()
	return nil
}

func (mf *MockFilesystem) Rename(oldname string, newname string) error {
	return errors.New("Not implemented")
}

func (mf *MockFilesystem) Create(name string) (file File, err error) {
	path := mf.getpath(name)
	dir, filename := filepath.Split(path)
	fi, err := mf.resolve(dir)
	if err != nil {
		return nil, err
	}
	fi.children[filename] = &MockFileInfo{
		name:       filename,
		filesystem: mf,
		mode:       0666,
		modified:   time.Now(),
		data:       nil,
		parent:     fi,
		children:   nil,
	}
	fi.modified = time.Now()
	f := &MockFile{
		filesystem: mf,
		closed: false,
		path: path,
	}
	return f, nil
}

func (mf *MockFilesystem) Open(name string) (file File, err error) {
	_, err = mf.resolve(name)
	if err != nil {
		return nil, err
	}
	f := &MockFile{
		filesystem: mf,
		closed: false,
		path: name,
	}
	return f, nil
}

func (mf *MockFilesystem) OpenFile(name string, flag int, perm os.FileMode) (file File, err error) {
	return nil, errors.New("Not implemented")
}

type MockFile struct {
	path       string
	closed     bool
	filesystem *MockFilesystem
}

func (mf *MockFile) stat() (mfi *MockFileInfo, err error) {
	return mf.filesystem.resolve(mf.path)
}

func (mf *MockFile) Chdir() error {
	return mf.filesystem.Chdir(filepath.Dir(mf.path))
}

func (mf *MockFile) Chmod(mode os.FileMode) error {
	var (
		mfi *MockFileInfo
		err error
	)
	if mfi, err = mf.stat(); err != nil {
		return err
	}
	mfi.mode = mode
	return nil
}

func (mf *MockFile) Close() error {
	mf.closed = true
	return nil
}

func (mf *MockFile) Name() string {
	_, filename := filepath.Split(mf.path)
	return filename
}

func (mf *MockFile) Read(b []byte) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) ReadAt(b []byte, off int64) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) Readdir(n int) (fi []os.FileInfo, err error) {
	// TODO: Enable returning additional elements in subsequent calls.
	var (
		mfi *MockFileInfo
	)
	fi = make([]os.FileInfo, 0)
	if mfi, err = mf.stat(); err != nil {
		return nil, err
	}
	limit := len(mfi.children)
	if n > 0 {
		limit = n
	}
	if len(mfi.children) < limit {
		err = io.EOF
	}
	if len(mfi.children) == 0 {
		return
	}
	i := 0
	for _, child := range mfi.children {
		if i == limit {
			break
		}
		fi = append(fi, child)
		i++
	}
	return
}

func (mf *MockFile) Readdirnames(n int) (names []string, err error) {
	fi, err := mf.Readdir(n)
	names = make([]string, len(fi))
	for i, f := range fi {
		names[i] = f.Name()
	}
	return names, err
}

func (mf *MockFile) Seek(offset int64, whence int) (ret int64, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) Stat() (fi os.FileInfo, err error) {
	return mf.stat()
}

func (mf *MockFile) Sync() (err error) {
	return nil
}

func (mf *MockFile) Truncate(size int64) error {
	return errors.New("Not implemented")
}

func (mf *MockFile) Write(b []byte) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) WriteAt(b []byte, off int64) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) WriteString(s string) (ret int, err error) {
	return 0, errors.New("Not implemented")
}

type MockFileInfo struct {
	data       *[]byte
	name       string
	filesystem *MockFilesystem
	mode       os.FileMode
	modified   time.Time
	parent     *MockFileInfo
	children   map[string]*MockFileInfo
}

// Not needed for interface.
func (mfi *MockFileInfo) path() string {
	ptr := mfi
	filepath := mfi.name
	for ptr != mfi.filesystem.root {
		ptr = ptr.parent
		filepath = path.Join(ptr.name, filepath)
	}
	return filepath
}

func (mfi *MockFileInfo) Parent() *MockFileInfo {
	return mfi.parent
}

func (mfi *MockFileInfo) Children() map[string]*MockFileInfo {
	return mfi.children
}

func (mfi *MockFileInfo) Child(name string) *MockFileInfo {
	return mfi.children[name]
}

func (mfi *MockFileInfo) Name() string {
	return mfi.name
}

func (mfi *MockFileInfo) Size() int64 {
	return int64(len(*mfi.data))
}

func (mfi *MockFileInfo) Mode() os.FileMode {
	return mfi.mode
}

func (mfi *MockFileInfo) ModTime() time.Time {
	return mfi.modified
}

func (mfi *MockFileInfo) IsDir() bool {
	return mfi.mode.IsDir()
}

func (mfi *MockFileInfo) Sys() interface{} {
	return nil
}
<commit_msg>Implements non-seeking file methods<commit_after>// Copyright 2012 Arne Roomann-Kurrik
//
// 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 fauxfile

import (
	"errors"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"bytes"
	"time"
)

func GetPathError(path string, message string) error {
	return &os.PathError{
		Path: path,
		Err:  errors.New(message),
	}
}

type MockFilesystem struct {
	cwd  *MockFileInfo
	root *MockFileInfo
}

func NewMockFilesystem() *MockFilesystem {
	root := &MockFileInfo{
		name:       "/",
		filesystem: nil,
		mode:       os.ModeDir | 0755,
		modified:   time.Now(),
		data:       nil,
		parent:     nil,
		children:   map[string]*MockFileInfo{},
	}
	mf := &MockFilesystem{
		cwd:  root,
		root: root,
	}
	root.filesystem = mf
	return mf
}

func (mf *MockFilesystem) getpath(path string) string {
	if filepath.IsAbs(path) {
		return filepath.Clean(path)
	}
	return filepath.Join(mf.cwd.path(), path)
}

func (mf *MockFilesystem) resolve(path string) (*MockFileInfo, error) {
	path = mf.getpath(path)
	parts := strings.Split(path, string(filepath.Separator))
	ptr := mf.root
	for _, part := range parts {
		if part == "" {
			continue
		}
		if child := ptr.Child(part); child != nil {
			ptr = ptr.Child(part)
		} else {
			return nil, GetPathError(path, "Path does not exist")
		}
	}
	return ptr, nil
}

func (mf *MockFilesystem) exists(path string) bool {
	_, err := mf.resolve(path)
	return err == nil
}

func (mf *MockFilesystem) Chdir(dir string) error {
	fi, err := mf.resolve(dir)
	if err == nil {
		if fi.IsDir() {
			mf.cwd = fi
			return nil
		} else {
			return GetPathError(dir, "Path is not a directory")
		}
	}
	return err
}

func (mf *MockFilesystem) Mkdir(name string, perm os.FileMode) error {
	path := mf.getpath(name)
	parentpath, dirname := filepath.Split(path)
	fi, err := mf.resolve(parentpath)
	if err != nil {
		return err
	}
	if child := fi.Child(dirname); child != nil {
		return GetPathError(path, "Path already exists")
	}
	fi.children[dirname] = &MockFileInfo{
		name:       dirname,
		filesystem: mf,
		mode:       perm | os.ModeDir,
		modified:   time.Now(),
		data:       new(bytes.Buffer),
		parent:     fi,
		children:   map[string]*MockFileInfo{},
	}
	fi.modified = time.Now() // Update directory's timestamp.
	return nil
}

func (mf *MockFilesystem) MkdirAll(path string, perm os.FileMode) error {
	path = mf.getpath(path)
	parts := strings.Split(path, string(filepath.Separator))
	base := "/"
	for _, part := range parts {
		if part == "" {
			continue
		}
		base = filepath.Join(base, part)
		if err := mf.Mkdir(base, perm); err != nil {
			return err
		}
	}
	return nil
}

func (mf *MockFilesystem) Remove(name string) error {
	fi, err := mf.resolve(name)
	if err != nil {
		return err
	}
	if len(fi.Children()) > 0 {
		return GetPathError(name, "Directory contains children")
	}
	delete(fi.Parent().Children(), fi.name)
	fi.Parent().modified = time.Now()
	return nil
}

func (mf *MockFilesystem) RemoveAll(path string) error {
	fi, err := mf.resolve(path)
	if err != nil {
		return err
	}
	delete(fi.Parent().Children(), fi.name)
	fi.Parent().modified = time.Now()
	return nil
}

func (mf *MockFilesystem) Rename(oldname string, newname string) error {
	return errors.New("Not implemented")
}

func (mf *MockFilesystem) Create(name string) (file File, err error) {
	path := mf.getpath(name)
	dir, filename := filepath.Split(path)
	fi, err := mf.resolve(dir)
	if err != nil {
		return nil, err
	}
	fi.children[filename] = &MockFileInfo{
		name:       filename,
		filesystem: mf,
		mode:       0666,
		modified:   time.Now(),
		data:       new(bytes.Buffer),
		parent:     fi,
		children:   nil,
	}
	fi.modified = time.Now()
	f := &MockFile{
		filesystem: mf,
		closed: false,
		path: path,
	}
	return f, nil
}

func (mf *MockFilesystem) Open(name string) (file File, err error) {
	_, err = mf.resolve(name)
	if err != nil {
		return nil, err
	}
	f := &MockFile{
		filesystem: mf,
		closed: false,
		path: name,
	}
	return f, nil
}

func (mf *MockFilesystem) OpenFile(name string, flag int, perm os.FileMode) (file File, err error) {
	return nil, errors.New("Not implemented")
}

type MockFile struct {
	path       string
	closed     bool
	filesystem *MockFilesystem
}

func (mf *MockFile) stat() (mfi *MockFileInfo, err error) {
	return mf.filesystem.resolve(mf.path)
}

func (mf *MockFile) Chdir() error {
	return mf.filesystem.Chdir(filepath.Dir(mf.path))
}

func (mf *MockFile) Chmod(mode os.FileMode) error {
	var (
		mfi *MockFileInfo
		err error
	)
	if mfi, err = mf.stat(); err != nil {
		return err
	}
	mfi.mode = mode
	return nil
}

func (mf *MockFile) Close() error {
	mf.closed = true
	return nil
}

func (mf *MockFile) Name() string {
	_, filename := filepath.Split(mf.path)
	return filename
}

func (mf *MockFile) Read(b []byte) (n int, err error) {
	var mfi *MockFileInfo
	if mfi, err = mf.stat(); err != nil {
		return 0, err
	}
	return mfi.data.Read(b)
}

func (mf *MockFile) ReadAt(b []byte, off int64) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) Readdir(n int) (fi []os.FileInfo, err error) {
	// TODO: Enable returning additional elements in subsequent calls.
	var (
		mfi *MockFileInfo
	)
	fi = make([]os.FileInfo, 0)
	if mfi, err = mf.stat(); err != nil {
		return nil, err
	}
	limit := len(mfi.children)
	if n > 0 {
		limit = n
	}
	if len(mfi.children) < limit {
		err = io.EOF
	}
	if len(mfi.children) == 0 {
		return
	}
	i := 0
	for _, child := range mfi.children {
		if i == limit {
			break
		}
		fi = append(fi, child)
		i++
	}
	return
}

func (mf *MockFile) Readdirnames(n int) (names []string, err error) {
	fi, err := mf.Readdir(n)
	names = make([]string, len(fi))
	for i, f := range fi {
		names[i] = f.Name()
	}
	return names, err
}

func (mf *MockFile) Seek(offset int64, whence int) (ret int64, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) Stat() (fi os.FileInfo, err error) {
	return mf.stat()
}

func (mf *MockFile) Sync() (err error) {
	return nil
}

func (mf *MockFile) Truncate(size int64) error {
	var mfi *MockFileInfo
	if mfi, err = mf.stat(); err != nil {
		return 0, err
	}
	return mfi.data.Truncate(size)
}

func (mf *MockFile) Write(b []byte) (n int, err error) {
	var mfi *MockFileInfo
	if mfi, err = mf.stat(); err != nil {
		return 0, err
	}
	return mfi.data.Write(b)
}

func (mf *MockFile) WriteAt(b []byte, off int64) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) WriteString(s string) (ret int, err error) {
	var mfi *MockFileInfo
	if mfi, err = mf.stat(); err != nil {
		return 0, err
	}
	return mfi.data.Write(s)
}

type MockFileInfo struct {
	data       *bytes.Buffer
	name       string
	filesystem *MockFilesystem
	mode       os.FileMode
	modified   time.Time
	parent     *MockFileInfo
	children   map[string]*MockFileInfo
}

func (mfi *MockFileInfo) path() string {
	ptr := mfi
	filepath := mfi.name
	for ptr != mfi.filesystem.root {
		ptr = ptr.parent
		filepath = path.Join(ptr.name, filepath)
	}
	return filepath
}

func (mfi *MockFileInfo) Parent() *MockFileInfo {
	return mfi.parent
}

func (mfi *MockFileInfo) Children() map[string]*MockFileInfo {
	return mfi.children
}

func (mfi *MockFileInfo) Child(name string) *MockFileInfo {
	return mfi.children[name]
}

func (mfi *MockFileInfo) Name() string {
	return mfi.name
}

func (mfi *MockFileInfo) Size() int64 {
	return int64(mfi.data.Len())
}

func (mfi *MockFileInfo) Mode() os.FileMode {
	return mfi.mode
}

func (mfi *MockFileInfo) ModTime() time.Time {
	return mfi.modified
}

func (mfi *MockFileInfo) IsDir() bool {
	return mfi.mode.IsDir()
}

func (mfi *MockFileInfo) Sys() interface{} {
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012 Arne Roomann-Kurrik
//
// 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 fauxfile

import (
	"errors"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

func GetPathError(path string, message string) error {
	return &os.PathError{
		Path: path,
		Err:  errors.New(message),
	}
}

type MockFilesystem struct {
	cwd  *MockFileInfo
	root *MockFileInfo
}

func NewMockFilesystem() *MockFilesystem {
	root := &MockFileInfo{
		name:       "/",
		filesystem: nil,
		mode:       os.ModeDir | 0755,
		modified:   time.Now(),
		data:       nil,
		parent:     nil,
		children:   map[string]*MockFileInfo{},
	}
	mf := &MockFilesystem{
		cwd:  root,
		root: root,
	}
	root.filesystem = mf
	return mf
}

func (mf *MockFilesystem) getpath(path string) string {
	if filepath.IsAbs(path) {
		return filepath.Clean(path)
	}
	return filepath.Join(mf.cwd.path(), path)
}

func (mf *MockFilesystem) resolve(path string) (*MockFileInfo, error) {
	path = mf.getpath(path)
	parts := strings.Split(path, string(filepath.Separator))
	ptr := mf.root
	for _, part := range parts {
		if part == "" {
			continue
		}
		if child := ptr.Child(part); child != nil {
			ptr = ptr.Child(part)
		} else {
			return nil, GetPathError(path, "Path does not exist")
		}
	}
	return ptr, nil
}

func (mf *MockFilesystem) exists(path string) bool {
	_, err := mf.resolve(path)
	return err == nil
}

func (mf *MockFilesystem) Chdir(dir string) error {
	fi, err := mf.resolve(dir)
	if err == nil {
		if fi.IsDir() {
			mf.cwd = fi
			return nil
		} else {
			return GetPathError(dir, "Path is not a directory")
		}
	}
	return err
}

func (mf *MockFilesystem) Mkdir(name string, perm os.FileMode) error {
	path := mf.getpath(name)
	parentpath, dirname := filepath.Split(path)
	fi, err := mf.resolve(parentpath)
	if err != nil {
		return err
	}
	if child := fi.Child(dirname); child != nil {
		return GetPathError(path, "Path already exists")
	}
	fi.children[dirname] = &MockFileInfo{
		name:       dirname,
		filesystem: mf,
		mode:       perm | os.ModeDir,
		modified:   time.Now(),
		data:       nil,
		parent:     fi,
		children:   map[string]*MockFileInfo{},
	}
	fi.modified = time.Now() // Update directory's timestamp.
	return nil
}

func (mf *MockFilesystem) MkdirAll(path string, perm os.FileMode) error {
	path = mf.getpath(path)
	parts := strings.Split(path, string(filepath.Separator))
	base := "/"
	for _, part := range parts {
		if part == "" {
			continue
		}
		base = filepath.Join(base, part)
		if err := mf.Mkdir(base, perm); err != nil {
			return err
		}
	}
	return nil
}

func (mf *MockFilesystem) Remove(name string) error {
	fi, err := mf.resolve(name)
	if err != nil {
		return err
	}
	if len(fi.Children()) > 0 {
		return GetPathError(name, "Directory contains children")
	}
	delete(fi.Parent().Children(), fi.name)
	fi.Parent().modified = time.Now()
	return nil
}

func (mf *MockFilesystem) RemoveAll(path string) error {
	fi, err := mf.resolve(path)
	if err != nil {
		return err
	}
	delete(fi.Parent().Children(), fi.name)
	fi.Parent().modified = time.Now()
	return nil
}

func (mf *MockFilesystem) Rename(oldname string, newname string) error {
	return errors.New("Not implemented")
}

func (mf *MockFilesystem) Create(name string) (file File, err error) {
	path := mf.getpath(name)
	dir, filename := filepath.Split(path)
	fi, err := mf.resolve(dir)
	if err != nil {
		return nil, err
	}
	fi.children[filename] = &MockFileInfo{
		name:       filename,
		filesystem: mf,
		mode:       0666,
		modified:   time.Now(),
		data:       nil,
		parent:     fi,
		children:   nil,
	}
	fi.modified = time.Now()
	f := &MockFile{
		filesystem: mf,
		closed: false,
		path: path,
	}
	return f, nil
}

func (mf *MockFilesystem) Open(name string) (file File, err error) {
	_, err = mf.resolve(name)
	if err != nil {
		return nil, err
	}
	f := &MockFile{
		filesystem: mf,
		closed: false,
		path: name,
	}
	return f, nil
}

func (mf *MockFilesystem) OpenFile(name string, flag int, perm os.FileMode) (file File, err error) {
	return nil, errors.New("Not implemented")
}

type MockFile struct {
	path       string
	closed     bool
	filesystem *MockFilesystem
}

func (mf *MockFile) stat() (mfi *MockFileInfo, err error) {
	return mf.filesystem.resolve(mf.path)
}

func (mf *MockFile) Chdir() error {
	return mf.filesystem.Chdir(filepath.Dir(mf.path))
}

func (mf *MockFile) Chmod(mode os.FileMode) error {
	var (
		mfi *MockFileInfo
		err error
	)
	if mfi, err = mf.stat(); err != nil {
		return err
	}
	mfi.mode = mode
	return nil
}

func (mf *MockFile) Close() error {
	mf.closed = true
	return nil
}

func (mf *MockFile) Name() string {
	_, filename := filepath.Split(mf.path)
	return filename
}

func (mf *MockFile) Read(b []byte) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) ReadAt(b []byte, off int64) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) Readdir(n int) (fi []os.FileInfo, err error) {
	// TODO: Enable returning additional elements in subsequent calls.
	var (
		mfi *MockFileInfo
	)
	fi = make([]os.FileInfo, 0)
	if mfi, err = mf.stat(); err != nil {
		return nil, err
	}
	limit := len(mfi.children)
	if n > 0 {
		limit = n
	}
	if len(mfi.children) < limit {
		err = io.EOF
	}
	if len(mfi.children) == 0 {
		return
	}
	i := 0
	for _, child := range mfi.children {
		if i == limit {
			break
		}
		fi = append(fi, child)
		i++
	}
	return
}

func (mf *MockFile) Readdirnames(n int) (names []string, err error) {
	fi, err := mf.Readdir(n)
	names = make([]string, len(fi))
	for i, f := range fi {
		names[i] = f.Name()
	}
	return names, err
}

func (mf *MockFile) Seek(offset int64, whence int) (ret int64, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) Stat() (fi os.FileInfo, err error) {
	return mf.stat()
}

func (mf *MockFile) Sync() (err error) {
	return errors.New("Not implemented")
}

func (mf *MockFile) Truncate(size int64) error {
	return errors.New("Not implemented")
}

func (mf *MockFile) Write(b []byte) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) WriteAt(b []byte, off int64) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) WriteString(s string) (ret int, err error) {
	return 0, errors.New("Not implemented")
}

type MockFileInfo struct {
	data       *[]byte
	name       string
	filesystem *MockFilesystem
	mode       os.FileMode
	modified   time.Time
	parent     *MockFileInfo
	children   map[string]*MockFileInfo
}

// Not needed for interface.
func (mfi *MockFileInfo) path() string {
	ptr := mfi
	filepath := mfi.name
	for ptr != mfi.filesystem.root {
		ptr = ptr.parent
		filepath = path.Join(ptr.name, filepath)
	}
	return filepath
}

func (mfi *MockFileInfo) Parent() *MockFileInfo {
	return mfi.parent
}

func (mfi *MockFileInfo) Children() map[string]*MockFileInfo {
	return mfi.children
}

func (mfi *MockFileInfo) Child(name string) *MockFileInfo {
	return mfi.children[name]
}

func (mfi *MockFileInfo) Name() string {
	return mfi.name
}

func (mfi *MockFileInfo) Size() int64 {
	return int64(len(*mfi.data))
}

func (mfi *MockFileInfo) Mode() os.FileMode {
	return mfi.mode
}

func (mfi *MockFileInfo) ModTime() time.Time {
	return mfi.modified
}

func (mfi *MockFileInfo) IsDir() bool {
	return mfi.mode.IsDir()
}

func (mfi *MockFileInfo) Sys() interface{} {
	return nil
}
<commit_msg>Implements sync<commit_after>// Copyright 2012 Arne Roomann-Kurrik
//
// 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 fauxfile

import (
	"errors"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

func GetPathError(path string, message string) error {
	return &os.PathError{
		Path: path,
		Err:  errors.New(message),
	}
}

type MockFilesystem struct {
	cwd  *MockFileInfo
	root *MockFileInfo
}

func NewMockFilesystem() *MockFilesystem {
	root := &MockFileInfo{
		name:       "/",
		filesystem: nil,
		mode:       os.ModeDir | 0755,
		modified:   time.Now(),
		data:       nil,
		parent:     nil,
		children:   map[string]*MockFileInfo{},
	}
	mf := &MockFilesystem{
		cwd:  root,
		root: root,
	}
	root.filesystem = mf
	return mf
}

func (mf *MockFilesystem) getpath(path string) string {
	if filepath.IsAbs(path) {
		return filepath.Clean(path)
	}
	return filepath.Join(mf.cwd.path(), path)
}

func (mf *MockFilesystem) resolve(path string) (*MockFileInfo, error) {
	path = mf.getpath(path)
	parts := strings.Split(path, string(filepath.Separator))
	ptr := mf.root
	for _, part := range parts {
		if part == "" {
			continue
		}
		if child := ptr.Child(part); child != nil {
			ptr = ptr.Child(part)
		} else {
			return nil, GetPathError(path, "Path does not exist")
		}
	}
	return ptr, nil
}

func (mf *MockFilesystem) exists(path string) bool {
	_, err := mf.resolve(path)
	return err == nil
}

func (mf *MockFilesystem) Chdir(dir string) error {
	fi, err := mf.resolve(dir)
	if err == nil {
		if fi.IsDir() {
			mf.cwd = fi
			return nil
		} else {
			return GetPathError(dir, "Path is not a directory")
		}
	}
	return err
}

func (mf *MockFilesystem) Mkdir(name string, perm os.FileMode) error {
	path := mf.getpath(name)
	parentpath, dirname := filepath.Split(path)
	fi, err := mf.resolve(parentpath)
	if err != nil {
		return err
	}
	if child := fi.Child(dirname); child != nil {
		return GetPathError(path, "Path already exists")
	}
	fi.children[dirname] = &MockFileInfo{
		name:       dirname,
		filesystem: mf,
		mode:       perm | os.ModeDir,
		modified:   time.Now(),
		data:       nil,
		parent:     fi,
		children:   map[string]*MockFileInfo{},
	}
	fi.modified = time.Now() // Update directory's timestamp.
	return nil
}

func (mf *MockFilesystem) MkdirAll(path string, perm os.FileMode) error {
	path = mf.getpath(path)
	parts := strings.Split(path, string(filepath.Separator))
	base := "/"
	for _, part := range parts {
		if part == "" {
			continue
		}
		base = filepath.Join(base, part)
		if err := mf.Mkdir(base, perm); err != nil {
			return err
		}
	}
	return nil
}

func (mf *MockFilesystem) Remove(name string) error {
	fi, err := mf.resolve(name)
	if err != nil {
		return err
	}
	if len(fi.Children()) > 0 {
		return GetPathError(name, "Directory contains children")
	}
	delete(fi.Parent().Children(), fi.name)
	fi.Parent().modified = time.Now()
	return nil
}

func (mf *MockFilesystem) RemoveAll(path string) error {
	fi, err := mf.resolve(path)
	if err != nil {
		return err
	}
	delete(fi.Parent().Children(), fi.name)
	fi.Parent().modified = time.Now()
	return nil
}

func (mf *MockFilesystem) Rename(oldname string, newname string) error {
	return errors.New("Not implemented")
}

func (mf *MockFilesystem) Create(name string) (file File, err error) {
	path := mf.getpath(name)
	dir, filename := filepath.Split(path)
	fi, err := mf.resolve(dir)
	if err != nil {
		return nil, err
	}
	fi.children[filename] = &MockFileInfo{
		name:       filename,
		filesystem: mf,
		mode:       0666,
		modified:   time.Now(),
		data:       nil,
		parent:     fi,
		children:   nil,
	}
	fi.modified = time.Now()
	f := &MockFile{
		filesystem: mf,
		closed: false,
		path: path,
	}
	return f, nil
}

func (mf *MockFilesystem) Open(name string) (file File, err error) {
	_, err = mf.resolve(name)
	if err != nil {
		return nil, err
	}
	f := &MockFile{
		filesystem: mf,
		closed: false,
		path: name,
	}
	return f, nil
}

func (mf *MockFilesystem) OpenFile(name string, flag int, perm os.FileMode) (file File, err error) {
	return nil, errors.New("Not implemented")
}

type MockFile struct {
	path       string
	closed     bool
	filesystem *MockFilesystem
}

func (mf *MockFile) stat() (mfi *MockFileInfo, err error) {
	return mf.filesystem.resolve(mf.path)
}

func (mf *MockFile) Chdir() error {
	return mf.filesystem.Chdir(filepath.Dir(mf.path))
}

func (mf *MockFile) Chmod(mode os.FileMode) error {
	var (
		mfi *MockFileInfo
		err error
	)
	if mfi, err = mf.stat(); err != nil {
		return err
	}
	mfi.mode = mode
	return nil
}

func (mf *MockFile) Close() error {
	mf.closed = true
	return nil
}

func (mf *MockFile) Name() string {
	_, filename := filepath.Split(mf.path)
	return filename
}

func (mf *MockFile) Read(b []byte) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) ReadAt(b []byte, off int64) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) Readdir(n int) (fi []os.FileInfo, err error) {
	// TODO: Enable returning additional elements in subsequent calls.
	var (
		mfi *MockFileInfo
	)
	fi = make([]os.FileInfo, 0)
	if mfi, err = mf.stat(); err != nil {
		return nil, err
	}
	limit := len(mfi.children)
	if n > 0 {
		limit = n
	}
	if len(mfi.children) < limit {
		err = io.EOF
	}
	if len(mfi.children) == 0 {
		return
	}
	i := 0
	for _, child := range mfi.children {
		if i == limit {
			break
		}
		fi = append(fi, child)
		i++
	}
	return
}

func (mf *MockFile) Readdirnames(n int) (names []string, err error) {
	fi, err := mf.Readdir(n)
	names = make([]string, len(fi))
	for i, f := range fi {
		names[i] = f.Name()
	}
	return names, err
}

func (mf *MockFile) Seek(offset int64, whence int) (ret int64, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) Stat() (fi os.FileInfo, err error) {
	return mf.stat()
}

func (mf *MockFile) Sync() (err error) {
	return nil
}

func (mf *MockFile) Truncate(size int64) error {
	return errors.New("Not implemented")
}

func (mf *MockFile) Write(b []byte) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) WriteAt(b []byte, off int64) (n int, err error) {
	return 0, errors.New("Not implemented")
}

func (mf *MockFile) WriteString(s string) (ret int, err error) {
	return 0, errors.New("Not implemented")
}

type MockFileInfo struct {
	data       *[]byte
	name       string
	filesystem *MockFilesystem
	mode       os.FileMode
	modified   time.Time
	parent     *MockFileInfo
	children   map[string]*MockFileInfo
}

// Not needed for interface.
func (mfi *MockFileInfo) path() string {
	ptr := mfi
	filepath := mfi.name
	for ptr != mfi.filesystem.root {
		ptr = ptr.parent
		filepath = path.Join(ptr.name, filepath)
	}
	return filepath
}

func (mfi *MockFileInfo) Parent() *MockFileInfo {
	return mfi.parent
}

func (mfi *MockFileInfo) Children() map[string]*MockFileInfo {
	return mfi.children
}

func (mfi *MockFileInfo) Child(name string) *MockFileInfo {
	return mfi.children[name]
}

func (mfi *MockFileInfo) Name() string {
	return mfi.name
}

func (mfi *MockFileInfo) Size() int64 {
	return int64(len(*mfi.data))
}

func (mfi *MockFileInfo) Mode() os.FileMode {
	return mfi.mode
}

func (mfi *MockFileInfo) ModTime() time.Time {
	return mfi.modified
}

func (mfi *MockFileInfo) IsDir() bool {
	return mfi.mode.IsDir()
}

func (mfi *MockFileInfo) Sys() interface{} {
	return nil
}
<|endoftext|>


<commit_before>package master

import (
	"fmt"
	"github.com/KIT-MAMID/mamid/model"
	"github.com/KIT-MAMID/mamid/msp"
	"github.com/Sirupsen/logrus"
	"github.com/jinzhu/gorm"
	"sync"
	"time"
)

var monitorLog = logrus.WithField("module", "monitor")

type Monitor struct {
	DB              *model.DB
	BusWriteChannel chan<- interface{}
	MSPClient       msp.MSPClient
	Interval        time.Duration
}

func (m *Monitor) Run() {
	ticker := time.NewTicker(m.Interval)
	quit := make(chan struct{})
	go func() {
		for {
			select {
			case <-ticker.C:
				monitorLog.Info("Monitor running")

				//Get all slaves from database
				tx := m.DB.Begin()
				var slaves []model.Slave
				err := tx.Find(&slaves).Error
				if err != nil {
					monitorLog.WithError(err).Error("Could not get slaves")
				}
				tx.Rollback()

				wg := sync.WaitGroup{}

				//Observe active slaves
				for _, slave := range slaves {
					if slave.ConfiguredState == model.SlaveStateActive {
						wg.Add(1)
						go func() {
							m.observeSlave(slave)
							wg.Done()
						}()
					}
				}

				//Wait for all slaves to be observed
				wg.Wait()

				//Check degradation of replica sets
				m.observeReplicaSets()

			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()
}

func mongodTuple(s model.Slave, m msp.Mongod) string {
	return fmt.Sprintf("(%s(id=%d),%d,%s)", s.Hostname, s.ID, m.Port, m.ReplicaSetName)
}

func (m *Monitor) observeSlave(slave model.Slave) {

	//Request mongod states from slave
	observedMongods, mspError := m.MSPClient.RequestStatus(msp.HostPort{slave.Hostname, msp.PortNumber(slave.Port)})

	// Notify about reachablility
	comErr := msp.Error{}
	if mspError != nil {
		comErr = *mspError
	}
	m.BusWriteChannel <- model.ConnectionStatus{
		Slave:              slave,
		Unreachable:        mspError != nil && mspError.Identifier == msp.CommunicationError,
		CommunicationError: comErr,
	}
	// TODO do we need to write this to the DB (currently there is no field for this in model.Slave)

	if mspError != nil {
		//TODO Handle other slave errors => check identifiers != CommunicationError
		//monitorLog.Error("monitor: error observing slave: %#v", mspError)
		return
	}

	tx := m.DB.Begin()

	if err := m.updateObservedStateInDB(tx, slave, observedMongods); err != nil {
		monitorLog.WithError(err).Error()
		tx.Rollback()
		return
	}

	if err := m.handleUnobservedMongodsOfSlave(tx, slave, observedMongods); err != nil {
		monitorLog.WithError(err).Error()
		tx.Rollback()
		return
	}

	tx.Commit()

	// Read-only transaction
	tx = m.DB.Begin()
	defer tx.Rollback()
	if err := m.sendMongodMismatchStatusToBus(tx, slave); err != nil {
		monitorLog.WithError(err).Error()
	}

}

// Update database Mongod.ObservedState with newly observedMongods
// Errors returned by this method should be handled by aborting the transaction tx
func (m *Monitor) updateObservedStateInDB(tx *gorm.DB, slave model.Slave, observedMongods []msp.Mongod) (criticalError error) {

	for _, observedMongod := range observedMongods {

		monitorLog.Debug("monitor: updating observed state for mongod `%s` in database`", mongodTuple(slave, observedMongod))

		var dbMongod model.Mongod

		dbMongodRes := tx.First(&dbMongod, &model.Mongod{
			ParentSlaveID: slave.ID,
			Port:          model.PortNumber(observedMongod.Port),
			ReplSetName:   observedMongod.ReplicaSetName,
		})

		if dbMongodRes.RecordNotFound() {
			return fmt.Errorf("monitor: internal inconsistency: did not find corresponding database Mongod to observed Mongod `%s`: %s",
				mongodTuple(slave, observedMongod), dbMongodRes.Error)
		} else if dbMongodRes.Error != nil {
			return fmt.Errorf("monitor: database error when querying for Mongod corresponding to observed Mongod `%s`: %s",
				mongodTuple(slave, observedMongod), dbMongodRes.Error)
		}

		//Get desired state if it exists
		relatedResult := tx.Model(&dbMongod).Related(&dbMongod.DesiredState, "DesiredState")
		if !relatedResult.RecordNotFound() && relatedResult.Error != nil {
			return fmt.Errorf("monitor: internal inconsistency: could not get desired state for Mongod `%s`: %s",
				mongodTuple(slave, observedMongod), relatedResult.Error.Error())
		}

		//Get observed state if it exists
		relatedResult = tx.Model(&dbMongod).Related(&dbMongod.ObservedState, "ObservedState")
		if !relatedResult.RecordNotFound() && relatedResult.Error != nil {
			return fmt.Errorf("monitor: database error when querying for observed state of Mongod `%s`: %s",
				mongodTuple(slave, observedMongod), relatedResult.Error)
		}

		// Update database representation of observation
		if observedMongod.StatusError == nil {
			//TODO Finish this
			//Put observations into model
			dbMongod.ObservedState.ExecutionState = mspMongodStateToModelExecutionState(observedMongod.State)
			dbMongod.ObservedState.IsShardingConfigServer = observedMongod.ShardingConfigServer
			dbMongod.ObservationError = model.MSPError{}
		} else {
			dbMongod.ObservationError = mspErrorToModelMSPError(observedMongod.StatusError)
		}

		// Persist updated database representation
		//TODO Only update observed state and errors to prevent collisions with cluster allocator
		saveErr := tx.Save(&dbMongod).Error
		if saveErr != nil {
			return fmt.Errorf("monitor: error persisting updated observed state for mongod `%s`: %s",
				mongodTuple(slave, observedMongod), saveErr.Error())
		}

		monitorLog.Debug("monitor: finished updating observed state for mongod `%s` in database`", mongodTuple(slave, observedMongod))

	}

	return nil

}

// Remove observed state of mongods the slave does not report
// Errors returned by this method should be handled by aborting the transaction tx
func (m *Monitor) handleUnobservedMongodsOfSlave(tx *gorm.DB, slave model.Slave, observedMongods []msp.Mongod) (err error) {

	var modelMongods []model.Mongod
	if err := tx.Model(&slave).Related(&modelMongods, "Mongods").Error; err != nil {
		return err
	}

outer:
	for _, modelMongod := range modelMongods {

		//Check if slave reported this mongod
		for _, observedMongod := range observedMongods {
			if modelMongod.Port == model.PortNumber(observedMongod.Port) &&
				modelMongod.ReplSetName == observedMongod.ReplicaSetName {
				continue outer
			}
		}

		//Else remove observed state
		deleteErr := tx.Delete(&model.MongodState{}, "id = ?", modelMongod.ObservedStateID).Error
		if deleteErr != nil {
			return deleteErr
		}
	}

	return nil

}

// Check every Mongod of the Slave for mismatches between DesiredState and ObservedState
// and send an appropriate MongodMismatchStatus to the Bus
func (m *Monitor) sendMongodMismatchStatusToBus(tx *gorm.DB, slave model.Slave) (err error) {

	var modelMongods []model.Mongod
	if err := tx.Model(&slave).Related(&modelMongods, "Mongods").Error; err != nil {
		return err
	}

	for _, modelMongod := range modelMongods {

		if err := tx.Model(modelMongod).Related(&modelMongod.DesiredState, "DesiredState").Error; err != nil {
			return fmt.Errorf("monitor: error fetching DesiredState for mongod `%v`: %s", modelMongod, err)
		}

		observedStateRes := tx.Model(modelMongod).Related(&modelMongod.ObservedState, "ObservedState")
		if !observedStateRes.RecordNotFound() && observedStateRes.Error != nil {
			return fmt.Errorf("monitor: error fetching ObservedState for mongod `%v`: %s", modelMongod, observedStateRes.Error)
		}

		if observedStateRes.RecordNotFound() {
			// This happens when
			// 	a new Mongod with no observations is found
			// 	a Mongod with desired state = destroyed is still in the database
			m.BusWriteChannel <- model.MongodMatchStatus{
				Mismatch: true,
				Mongod:   modelMongod,
			}
		} else {
			m.BusWriteChannel <- compareStates(modelMongod)
		}

	}

	return nil
}

func compareStates(mongod model.Mongod) (m model.MongodMatchStatus) {
	//TODO Finish this
	m.Mismatch =
		mongod.DesiredState.ExecutionState != mongod.ObservedState.ExecutionState ||
			mongod.DesiredState.IsShardingConfigServer != mongod.ObservedState.IsShardingConfigServer
	m.Mongod = mongod
	return
}

func mspMongodStateToModelExecutionState(e msp.MongodState) model.MongodExecutionState {
	switch e {
	case msp.MongodStateDestroyed:
		return model.MongodExecutionStateDestroyed
	case msp.MongodStateNotRunning:
		return model.MongodExecutionStateNotRunning
	case msp.MongodStateRecovering:
		return model.MongodExecutionStateRecovering
	case msp.MongodStateRunning:
		return model.MongodExecutionStateRunning
	default:
		return 0 // Invalid
	}
}

func mspErrorToModelMSPError(mspError *msp.Error) model.MSPError {
	return model.MSPError{
		Identifier:      mspError.Identifier,
		Description:     mspError.Description,
		LongDescription: mspError.LongDescription,
	}
}

func (m *Monitor) observeReplicaSets() {
	tx := m.DB.Begin()

	// Get replica sets and the count of their actually configured members from the database
	replicaSetsWithMemberCounts, err := tx.Raw(`SELECT
				r.*,
				(SELECT COUNT(*) FROM replica_set_configured_members WHERE replica_set_id = r.id AND persistent_storage = ?)
					AS configured_persistent_members,
				(SELECT COUNT(*) FROM replica_set_configured_members WHERE replica_set_id = r.id AND persistent_storage = ?)
					AS configured_volatile_members,
				(SELECT COUNT(*) FROM replica_set_effective_members WHERE replica_set_id = r.id AND persistent_storage = ?)
					AS actual_persistent_members,
				(SELECT COUNT(*) FROM replica_set_effective_members WHERE replica_set_id = r.id AND persistent_storage = ?)
					AS actual_volatile_members
				FROM replica_sets r
				`, true, false, true, false).Rows()
	tx.Rollback()
	if err != nil {
		monitorLog.WithError(err).Error("Error getting configured and actual member counts of replica sets")
		return
	}

	for replicaSetsWithMemberCounts.Next() {
		var replicaSet model.ReplicaSet
		tx.ScanRows(replicaSetsWithMemberCounts, &replicaSet)

		memberCounts := struct {
			ConfiguredPersistentMembers uint
			ConfiguredVolatileMembers   uint
			ActualPersistentMembers     uint
			ActualVolatileMembers       uint
		}{}
		tx.ScanRows(replicaSetsWithMemberCounts, &memberCounts)

		unsatisfied := memberCounts.ConfiguredVolatileMembers > memberCounts.ActualVolatileMembers ||
			memberCounts.ConfiguredPersistentMembers > memberCounts.ActualPersistentMembers

		m.BusWriteChannel <- model.ObservedReplicaSetConstraintStatus{
			Unsatisfied:               unsatisfied,
			ReplicaSet:                replicaSet,
			ConfiguredPersistentCount: memberCounts.ConfiguredPersistentMembers,
			ConfiguredVolatileCount:   memberCounts.ConfiguredVolatileMembers,
			ActualPersistentCount:     memberCounts.ActualPersistentMembers,
			ActualVolatileCount:       memberCounts.ActualVolatileMembers,
		}
	}
}
<commit_msg>FIX: monitor: call by value to observeSlave goroutine<commit_after>package master

import (
	"fmt"
	"github.com/KIT-MAMID/mamid/model"
	"github.com/KIT-MAMID/mamid/msp"
	"github.com/Sirupsen/logrus"
	"github.com/jinzhu/gorm"
	"sync"
	"time"
)

var monitorLog = logrus.WithField("module", "monitor")

type Monitor struct {
	DB              *model.DB
	BusWriteChannel chan<- interface{}
	MSPClient       msp.MSPClient
	Interval        time.Duration
}

func (m *Monitor) Run() {
	ticker := time.NewTicker(m.Interval)
	quit := make(chan struct{})
	go func() {
		for {
			select {
			case <-ticker.C:
				monitorLog.Info("Monitor running")

				//Get all slaves from database
				tx := m.DB.Begin()
				var slaves []model.Slave
				err := tx.Find(&slaves).Error
				if err != nil {
					monitorLog.WithError(err).Error("Could not get slaves")
				}
				tx.Rollback()

				wg := sync.WaitGroup{}

				//Observe active slaves
				for _, slave := range slaves {
					if slave.ConfiguredState == model.SlaveStateActive {
						wg.Add(1)
						go func(s model.Slave) {
							m.observeSlave(s)
							wg.Done()
						}(slave)
					}
				}

				//Wait for all slaves to be observed
				wg.Wait()

				//Check degradation of replica sets
				m.observeReplicaSets()

			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()
}

func mongodTuple(s model.Slave, m msp.Mongod) string {
	return fmt.Sprintf("(%s(id=%d),%d,%s)", s.Hostname, s.ID, m.Port, m.ReplicaSetName)
}

func (m *Monitor) observeSlave(slave model.Slave) {

	//Request mongod states from slave
	observedMongods, mspError := m.MSPClient.RequestStatus(msp.HostPort{slave.Hostname, msp.PortNumber(slave.Port)})

	// Notify about reachablility
	comErr := msp.Error{}
	if mspError != nil {
		comErr = *mspError
	}
	m.BusWriteChannel <- model.ConnectionStatus{
		Slave:              slave,
		Unreachable:        mspError != nil && mspError.Identifier == msp.CommunicationError,
		CommunicationError: comErr,
	}
	// TODO do we need to write this to the DB (currently there is no field for this in model.Slave)

	if mspError != nil {
		//TODO Handle other slave errors => check identifiers != CommunicationError
		//monitorLog.Error("monitor: error observing slave: %#v", mspError)
		return
	}

	tx := m.DB.Begin()

	if err := m.updateObservedStateInDB(tx, slave, observedMongods); err != nil {
		monitorLog.WithError(err).Error()
		tx.Rollback()
		return
	}

	if err := m.handleUnobservedMongodsOfSlave(tx, slave, observedMongods); err != nil {
		monitorLog.WithError(err).Error()
		tx.Rollback()
		return
	}

	tx.Commit()

	// Read-only transaction
	tx = m.DB.Begin()
	defer tx.Rollback()
	if err := m.sendMongodMismatchStatusToBus(tx, slave); err != nil {
		monitorLog.WithError(err).Error()
	}

}

// Update database Mongod.ObservedState with newly observedMongods
// Errors returned by this method should be handled by aborting the transaction tx
func (m *Monitor) updateObservedStateInDB(tx *gorm.DB, slave model.Slave, observedMongods []msp.Mongod) (criticalError error) {

	for _, observedMongod := range observedMongods {

		monitorLog.Debug("monitor: updating observed state for mongod `%s` in database`", mongodTuple(slave, observedMongod))

		var dbMongod model.Mongod

		dbMongodRes := tx.First(&dbMongod, &model.Mongod{
			ParentSlaveID: slave.ID,
			Port:          model.PortNumber(observedMongod.Port),
			ReplSetName:   observedMongod.ReplicaSetName,
		})

		if dbMongodRes.RecordNotFound() {
			return fmt.Errorf("monitor: internal inconsistency: did not find corresponding database Mongod to observed Mongod `%s`: %s",
				mongodTuple(slave, observedMongod), dbMongodRes.Error)
		} else if dbMongodRes.Error != nil {
			return fmt.Errorf("monitor: database error when querying for Mongod corresponding to observed Mongod `%s`: %s",
				mongodTuple(slave, observedMongod), dbMongodRes.Error)
		}

		//Get desired state if it exists
		relatedResult := tx.Model(&dbMongod).Related(&dbMongod.DesiredState, "DesiredState")
		if !relatedResult.RecordNotFound() && relatedResult.Error != nil {
			return fmt.Errorf("monitor: internal inconsistency: could not get desired state for Mongod `%s`: %s",
				mongodTuple(slave, observedMongod), relatedResult.Error.Error())
		}

		//Get observed state if it exists
		relatedResult = tx.Model(&dbMongod).Related(&dbMongod.ObservedState, "ObservedState")
		if !relatedResult.RecordNotFound() && relatedResult.Error != nil {
			return fmt.Errorf("monitor: database error when querying for observed state of Mongod `%s`: %s",
				mongodTuple(slave, observedMongod), relatedResult.Error)
		}

		// Update database representation of observation
		if observedMongod.StatusError == nil {
			//TODO Finish this
			//Put observations into model
			dbMongod.ObservedState.ExecutionState = mspMongodStateToModelExecutionState(observedMongod.State)
			dbMongod.ObservedState.IsShardingConfigServer = observedMongod.ShardingConfigServer
			dbMongod.ObservationError = model.MSPError{}
		} else {
			dbMongod.ObservationError = mspErrorToModelMSPError(observedMongod.StatusError)
		}

		// Persist updated database representation
		//TODO Only update observed state and errors to prevent collisions with cluster allocator
		saveErr := tx.Save(&dbMongod).Error
		if saveErr != nil {
			return fmt.Errorf("monitor: error persisting updated observed state for mongod `%s`: %s",
				mongodTuple(slave, observedMongod), saveErr.Error())
		}

		monitorLog.Debug("monitor: finished updating observed state for mongod `%s` in database`", mongodTuple(slave, observedMongod))

	}

	return nil

}

// Remove observed state of mongods the slave does not report
// Errors returned by this method should be handled by aborting the transaction tx
func (m *Monitor) handleUnobservedMongodsOfSlave(tx *gorm.DB, slave model.Slave, observedMongods []msp.Mongod) (err error) {

	var modelMongods []model.Mongod
	if err := tx.Model(&slave).Related(&modelMongods, "Mongods").Error; err != nil {
		return err
	}

outer:
	for _, modelMongod := range modelMongods {

		//Check if slave reported this mongod
		for _, observedMongod := range observedMongods {
			if modelMongod.Port == model.PortNumber(observedMongod.Port) &&
				modelMongod.ReplSetName == observedMongod.ReplicaSetName {
				continue outer
			}
		}

		//Else remove observed state
		deleteErr := tx.Delete(&model.MongodState{}, "id = ?", modelMongod.ObservedStateID).Error
		if deleteErr != nil {
			return deleteErr
		}
	}

	return nil

}

// Check every Mongod of the Slave for mismatches between DesiredState and ObservedState
// and send an appropriate MongodMismatchStatus to the Bus
func (m *Monitor) sendMongodMismatchStatusToBus(tx *gorm.DB, slave model.Slave) (err error) {

	var modelMongods []model.Mongod
	if err := tx.Model(&slave).Related(&modelMongods, "Mongods").Error; err != nil {
		return err
	}

	for _, modelMongod := range modelMongods {

		if err := tx.Model(modelMongod).Related(&modelMongod.DesiredState, "DesiredState").Error; err != nil {
			return fmt.Errorf("monitor: error fetching DesiredState for mongod `%v`: %s", modelMongod, err)
		}

		observedStateRes := tx.Model(modelMongod).Related(&modelMongod.ObservedState, "ObservedState")
		if !observedStateRes.RecordNotFound() && observedStateRes.Error != nil {
			return fmt.Errorf("monitor: error fetching ObservedState for mongod `%v`: %s", modelMongod, observedStateRes.Error)
		}

		if observedStateRes.RecordNotFound() {
			// This happens when
			// 	a new Mongod with no observations is found
			// 	a Mongod with desired state = destroyed is still in the database
			m.BusWriteChannel <- model.MongodMatchStatus{
				Mismatch: true,
				Mongod:   modelMongod,
			}
		} else {
			m.BusWriteChannel <- compareStates(modelMongod)
		}

	}

	return nil
}

func compareStates(mongod model.Mongod) (m model.MongodMatchStatus) {
	//TODO Finish this
	m.Mismatch =
		mongod.DesiredState.ExecutionState != mongod.ObservedState.ExecutionState ||
			mongod.DesiredState.IsShardingConfigServer != mongod.ObservedState.IsShardingConfigServer
	m.Mongod = mongod
	return
}

func mspMongodStateToModelExecutionState(e msp.MongodState) model.MongodExecutionState {
	switch e {
	case msp.MongodStateDestroyed:
		return model.MongodExecutionStateDestroyed
	case msp.MongodStateNotRunning:
		return model.MongodExecutionStateNotRunning
	case msp.MongodStateRecovering:
		return model.MongodExecutionStateRecovering
	case msp.MongodStateRunning:
		return model.MongodExecutionStateRunning
	default:
		return 0 // Invalid
	}
}

func mspErrorToModelMSPError(mspError *msp.Error) model.MSPError {
	return model.MSPError{
		Identifier:      mspError.Identifier,
		Description:     mspError.Description,
		LongDescription: mspError.LongDescription,
	}
}

func (m *Monitor) observeReplicaSets() {
	tx := m.DB.Begin()

	// Get replica sets and the count of their actually configured members from the database
	replicaSetsWithMemberCounts, err := tx.Raw(`SELECT
				r.*,
				(SELECT COUNT(*) FROM replica_set_configured_members WHERE replica_set_id = r.id AND persistent_storage = ?)
					AS configured_persistent_members,
				(SELECT COUNT(*) FROM replica_set_configured_members WHERE replica_set_id = r.id AND persistent_storage = ?)
					AS configured_volatile_members,
				(SELECT COUNT(*) FROM replica_set_effective_members WHERE replica_set_id = r.id AND persistent_storage = ?)
					AS actual_persistent_members,
				(SELECT COUNT(*) FROM replica_set_effective_members WHERE replica_set_id = r.id AND persistent_storage = ?)
					AS actual_volatile_members
				FROM replica_sets r
				`, true, false, true, false).Rows()
	tx.Rollback()
	if err != nil {
		monitorLog.WithError(err).Error("Error getting configured and actual member counts of replica sets")
		return
	}

	for replicaSetsWithMemberCounts.Next() {
		var replicaSet model.ReplicaSet
		tx.ScanRows(replicaSetsWithMemberCounts, &replicaSet)

		memberCounts := struct {
			ConfiguredPersistentMembers uint
			ConfiguredVolatileMembers   uint
			ActualPersistentMembers     uint
			ActualVolatileMembers       uint
		}{}
		tx.ScanRows(replicaSetsWithMemberCounts, &memberCounts)

		unsatisfied := memberCounts.ConfiguredVolatileMembers > memberCounts.ActualVolatileMembers ||
			memberCounts.ConfiguredPersistentMembers > memberCounts.ActualPersistentMembers

		m.BusWriteChannel <- model.ObservedReplicaSetConstraintStatus{
			Unsatisfied:               unsatisfied,
			ReplicaSet:                replicaSet,
			ConfiguredPersistentCount: memberCounts.ConfiguredPersistentMembers,
			ConfiguredVolatileCount:   memberCounts.ConfiguredVolatileMembers,
			ActualPersistentCount:     memberCounts.ActualPersistentMembers,
			ActualVolatileCount:       memberCounts.ActualVolatileMembers,
		}
	}
}
<|endoftext|>


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

package query

// show all metrics
// show tags WHERE predicate

import (
	"github.com/square/metrics/assert"
	"testing"
)

// these queries should successfully parse,
// with a corresponding command.
var inputs = []string{
	// describes
	"describe all",
	"describe x",
	"describe cpu_usage",
	"describe cpu_usage where key = 'value'",
	"describe cpu_usage where key = 'value\\''",
	"describe cpu_usage where key != 'value'",
	"describe cpu_usage where (key = 'value')",
	"describe cpu_usage where not (key = 'value')",
	"describe cpu_usage where not key = 'value'",
	"describe cpu_usage where (key = 'value' and key = 'value')",
	"describe cpu_usage where key = 'value' or key = 'value'",
	"describe cpu_usage where key in ('value', 'value')",
	"describe cpu_usage where key matches 'abc'",
	"describe nodes.cpu.usage where datacenter='sjc1b' and type='idle' and host matches 'fwd'",
	// selects - numbers
	"select 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 from 0 to 0",
	"select 10, 100, 1000 from 0 to 0",
	"select 10.1, 10.01, 10.001 from 0 to 0",
	"select -10.1, -10.01, -10.001 from 0 to 0",
	"select 1.0e1, 1.0e2, 1.0e10, 1.0e0 from 0 to 0",
	"select 1.0e-5, 1.0e+5 from 0 to 0",
	// selects - trying out arithmetic
	"select x from 0 to 0",
	"select x-y-z from 0 to 0",
	"select (x)-(y)-(z) from 0 to 0",
	"select 0 from 0 to 0",
	"select x, y from 0 to 0",
	"select 1 + 2 * 3 + 4 from 0 to 0",
	"select x * (y + 123), z from 0 to 0",
	// testing escaping
	"select `x` from 0 to 0",
	// selects - timestamps
	"select x * (y + 123), z from '2014-01-01' to '2014-01-02'",
	"select x * (y + 123), z from 0 to 10000",
	// selects - function calls
	"select foo(x) from 0 to 0",
	"select bar(x, y) from 0 to 0",
	"select baz(x, y, z+1+foo(1)) from 0 to 0",
}

var parseOnly = []string{
	// selects - aggregate functions
	"select scalar.max(x) from 0 to 0",
	"select aggregate.max(x, y) from 0 to 0",
	"select aggregate.max(x group by foo) + 3 from 0 to 0",
	// selects - where clause
	"select x where y = 'z' from 0 to 0",
	// selects - per-identifier where clause
	"select x + z[y = 'z'] from 0 to 0",
	"select x[y = 'z'] from 0 to 0",
	// selects - complicated queries
	"select aggregate.max(x[y = 'z'] group by foo) from 0 to 0",
	"select cpu.user + cpu.kernel where host = 'apa3.sjc2b' from 0 to 0",
}

// these queries should fail with a syntax error.
var syntaxErrorQuery = []string{
	"select (",
	"select )",
	"describe (",
	"describe invalid_regex where key matches 'ab['",
}

func TestParse_success(t *testing.T) {
	for _, row := range inputs {
		if err := checkSyntaxError(t, row); err != nil {
			t.Errorf("[%s] failed to parse: %s", row, err.Error())
		}
	}
	for _, row := range parseOnly {
		if err := checkSyntaxError(t, row); err != nil {
			t.Errorf("[%s] failed to parse: %s", row, err.Error())
		}
	}
}

func TestParse_syntaxError(t *testing.T) {
	for _, row := range syntaxErrorQuery {
		_, err := Parse(row)
		if err == nil {
			t.Errorf("[%s] should have failed to parse", row)
		} else if _, ok := err.(SyntaxErrors); !ok {
			t.Logf("[%s] Expected SyntaxErrors, got: %s", row, err.Error())
		}
	}
}

func TestCompile(t *testing.T) {
	for _, row := range inputs {
		a := assert.New(t).Contextf(row)
		p := Parser{Buffer: row}
		p.Init()
		a.CheckError(p.Parse())
		p.Execute()
		testParserResult(a, p)
	}
}

// Helper functions
// ================

func checkSyntaxError(t *testing.T, input string) error {
	p := Parser{Buffer: input}
	p.Init()
	return p.Parse()
}

func testParserResult(a assert.Assert, p Parser) {
	a.EqInt(len(p.assertions), 0)
	if len(p.assertions) != 0 {
		for _, err := range p.assertions {
			a.Errorf("assertion error: %s", err.Error())
		}
	}
	if len(p.nodeStack) != 0 {
		for _, node := range p.nodeStack {
			a.Errorf("node error:\n%s", PrintNode(node))
		}
	}
}
<commit_msg>adding tests<commit_after>// Copyright 2015 Square Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package query

// show all metrics
// show tags WHERE predicate

import (
	"github.com/square/metrics/assert"
	"testing"
)

// these queries should successfully parse,
// with a corresponding command.
var inputs = []string{
	// describes
	"describe all",
	"describe x",
	"describe cpu_usage",
	"describe cpu_usage where key = 'value'",
	"describe cpu_usage where key = 'value\\''",
	"describe cpu_usage where key != 'value'",
	"describe cpu_usage where (key = 'value')",
	"describe cpu_usage where not (key = 'value')",
	"describe cpu_usage where not key = 'value'",
	"describe cpu_usage where (key = 'value' and key = 'value')",
	"describe cpu_usage where key = 'value' or key = 'value'",
	"describe cpu_usage where key in ('value', 'value')",
	"describe cpu_usage where key matches 'abc'",
	"describe nodes.cpu.usage where datacenter='sjc1b' and type='idle' and host matches 'fwd'",
	// selects - numbers
	"select 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 from 0 to 0",
	"select 10, 100, 1000 from 0 to 0",
	"select 10.1, 10.01, 10.001 from 0 to 0",
	"select -10.1, -10.01, -10.001 from 0 to 0",
	"select 1.0e1, 1.0e2, 1.0e10, 1.0e0 from 0 to 0",
	"select 1.0e-5, 1.0e+5 from 0 to 0",
	// selects - trying out arithmetic
	"select x from 0 to 0",
	"select x-y-z from 0 to 0",
	"select (x)-(y)-(z) from 0 to 0",
	"select 0 from 0 to 0",
	"select x, y from 0 to 0",
	"select 1 + 2 * 3 + 4 from 0 to 0",
	"select x * (y + 123), z from 0 to 0",
	// testing escaping
	"select `x` from 0 to 0",
	// selects - timestamps
	"select x * (y + 123), z from '2014-01-01' to '2014-01-02'",
	"select x * (y + 123), z from 0 to 10000",
	// selects - function calls
	"select foo(x) from 0 to 0",
	"select bar(x, y) from 0 to 0",
	"select baz(x, y, z+1+foo(1)) from 0 to 0",
	// selects - testing out property values
	"select x",
	"select x from 0",
	"select x resolution 10",
	"select x sample 'max'",
}

var parseOnly = []string{
	// selects - aggregate functions
	"select scalar.max(x) from 0 to 0",
	"select aggregate.max(x, y) from 0 to 0",
	"select aggregate.max(x group by foo) + 3 from 0 to 0",
	// selects - where clause
	"select x where y = 'z' from 0 to 0",
	// selects - per-identifier where clause
	"select x + z[y = 'z'] from 0 to 0",
	"select x[y = 'z'] from 0 to 0",
	// selects - complicated queries
	"select aggregate.max(x[y = 'z'] group by foo) from 0 to 0",
	"select cpu.user + cpu.kernel where host = 'apa3.sjc2b' from 0 to 0",
}

// these queries should fail with a syntax error.
var syntaxErrorQuery = []string{
	"select (",
	"select )",
	"describe (",
	"describe invalid_regex where key matches 'ab['",
	"select x invalid_property 0",
}

func TestParse_success(t *testing.T) {
	for _, row := range inputs {
		if err := checkSyntaxError(t, row); err != nil {
			t.Errorf("[%s] failed to parse: %s", row, err.Error())
		}
	}
	for _, row := range parseOnly {
		if err := checkSyntaxError(t, row); err != nil {
			t.Errorf("[%s] failed to parse: %s", row, err.Error())
		}
	}
}

func TestParse_syntaxError(t *testing.T) {
	for _, row := range syntaxErrorQuery {
		_, err := Parse(row)
		if err == nil {
			t.Errorf("[%s] should have failed to parse", row)
		} else if _, ok := err.(SyntaxErrors); !ok {
			t.Logf("[%s] Expected SyntaxErrors, got: %s", row, err.Error())
		}
	}
}

func TestCompile(t *testing.T) {
	for _, row := range inputs {
		a := assert.New(t).Contextf(row)
		p := Parser{Buffer: row}
		p.Init()
		a.CheckError(p.Parse())
		p.Execute()
		testParserResult(a, p)
	}
}

// Helper functions
// ================

func checkSyntaxError(t *testing.T, input string) error {
	p := Parser{Buffer: input}
	p.Init()
	return p.Parse()
}

func testParserResult(a assert.Assert, p Parser) {
	a.EqInt(len(p.assertions), 0)
	if len(p.assertions) != 0 {
		for _, err := range p.assertions {
			a.Errorf("assertion error: %s", err.Error())
		}
	}
	if len(p.nodeStack) != 0 {
		for _, node := range p.nodeStack {
			a.Errorf("node error:\n%s", PrintNode(node))
		}
	}
}
<|endoftext|>


<commit_before>package queue

import (
	"os"
	"testing"

	"github.com/hashicorp/go-multierror"
	"github.com/stretchr/testify/suite"
	"github.com/stripe/aws-go/aws"
	"github.com/stripe/aws-go/gen/sqs"

	"github.com/nabeken/aws-go-sqs/queue/option"
)

func testSQSQueue(name string) (*Queue, error) {
	return New(
		sqs.New(aws.DetectCreds("", "", ""), "ap-northeast-1", nil),
		name,
	)
}

type SendMessageBatchSuite struct {
	suite.Suite

	queue *Queue
}

func (s *SendMessageBatchSuite) SetupSuite() {
	name := os.Getenv("TEST_SQS_QUEUE_NAME")
	if len(name) == 0 {
		s.T().Fatal("TEST_SQS_QUEUE_NAME must be set")
	}

	q, err := testSQSQueue(name)
	if err != nil {
		s.T().Fatal(err)
	}

	s.queue = q
}

func (s *SendMessageBatchSuite) SetupTest() {
}

func (s *SendMessageBatchSuite) TearDownTest() {
	// don't care of the result
	s.queue.PurgeQueue()
}

func (s *SendMessageBatchSuite) TearDownSuite() {
}

func (s *SendMessageBatchSuite) TestSendMessageBatch() {
	attrs := map[string]interface{}{
		"ATTR1": "STRING!!",
		"ATTR2": 12345,
	}

	batchMessages := []BatchMessage{
		BatchMessage{
			Body:    "body1",
			Options: []option.SendMessageRequest{option.MessageAttributes(attrs)},
		},
		BatchMessage{
			Body:    "body2",
			Options: []option.SendMessageRequest{option.MessageAttributes(attrs)},
		},
	}

	if err := s.queue.SendMessageBatch(batchMessages...); !s.NoError(err) {
		return
	}

	messages, err := s.queue.ReceiveMessage(
		option.MaxNumberOfMessages(5),
		option.UseAllAttribute(),
	)
	if !s.NoError(err) {
		return
	}

	s.Len(messages, 2)

	for i, m := range messages {
		s.Len(m.MessageAttributes, 2)
		for k, a := range m.MessageAttributes {
			mav := option.MessageAttributeValue(attrs[k])
			s.Equal(mav.StringValue, a.StringValue)
		}
		s.Equal(batchMessages[i].Body, *m.Body)
		s.queue.DeleteMessage(m.ReceiptHandle)
	}
}

func (s *SendMessageBatchSuite) TestSendMessageBatchError() {
	attrs := map[string]interface{}{
		"error": "",
	}

	batchMessages := []BatchMessage{
		BatchMessage{
			Body: "success",
		},
		BatchMessage{
			Body:    "failed",
			Options: []option.SendMessageRequest{option.MessageAttributes(attrs)},
		},
	}

	if err := s.queue.SendMessageBatch(batchMessages...); s.Error(err) {
		if merr, ok := err.(*multierror.Error); s.True(ok, "error must be *multierror.Error") {
			s.Len(merr.Errors, 1)
			if berr, ok := merr.Errors[0].(*BatchError); s.True(ok, "error must be *BatchError") {
				s.Equal(1, berr.Index, "batchMessages[1] must be error")
				s.Equal("InvalidParameterValue", berr.Code)
				s.Equal(true, berr.SenderFault)
			}
		}
	}

	messages, err := s.queue.ReceiveMessage(
		option.MaxNumberOfMessages(5),
		option.UseAllAttribute(),
	)
	if !s.NoError(err) {
		return
	}

	s.Len(messages, 1)
	for _, m := range messages {
		s.queue.DeleteMessage(m.ReceiptHandle)
	}
}

func TestSendMessageBatchSuite(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test")
	}

	suite.Run(t, new(SendMessageBatchSuite))
}
<commit_msg>queue_test: Perform PurgeQueue after teardown suite<commit_after>package queue

import (
	"os"
	"testing"

	"github.com/hashicorp/go-multierror"
	"github.com/stretchr/testify/suite"
	"github.com/stripe/aws-go/aws"
	"github.com/stripe/aws-go/gen/sqs"

	"github.com/nabeken/aws-go-sqs/queue/option"
)

func testSQSQueue(name string) (*Queue, error) {
	return New(
		sqs.New(aws.DetectCreds("", "", ""), "ap-northeast-1", nil),
		name,
	)
}

type SendMessageBatchSuite struct {
	suite.Suite

	queue *Queue
}

func (s *SendMessageBatchSuite) SetupSuite() {
	name := os.Getenv("TEST_SQS_QUEUE_NAME")
	if len(name) == 0 {
		s.T().Fatal("TEST_SQS_QUEUE_NAME must be set")
	}

	q, err := testSQSQueue(name)
	if err != nil {
		s.T().Fatal(err)
	}

	s.queue = q
}

func (s *SendMessageBatchSuite) SetupTest() {
}

func (s *SendMessageBatchSuite) TearDownTest() {
}

func (s *SendMessageBatchSuite) TearDownSuite() {
	// don't care of the result
	s.queue.PurgeQueue()
}

func (s *SendMessageBatchSuite) TestSendMessageBatch() {
	attrs := map[string]interface{}{
		"ATTR1": "STRING!!",
		"ATTR2": 12345,
	}

	batchMessages := []BatchMessage{
		BatchMessage{
			Body:    "body1",
			Options: []option.SendMessageRequest{option.MessageAttributes(attrs)},
		},
		BatchMessage{
			Body:    "body2",
			Options: []option.SendMessageRequest{option.MessageAttributes(attrs)},
		},
	}

	if err := s.queue.SendMessageBatch(batchMessages...); !s.NoError(err) {
		return
	}

	messages, err := s.queue.ReceiveMessage(
		option.MaxNumberOfMessages(5),
		option.UseAllAttribute(),
	)
	if !s.NoError(err) {
		return
	}

	s.Len(messages, 2)

	for i, m := range messages {
		s.Len(m.MessageAttributes, 2)
		for k, a := range m.MessageAttributes {
			mav := option.MessageAttributeValue(attrs[k])
			s.Equal(mav.StringValue, a.StringValue)
		}
		s.Equal(batchMessages[i].Body, *m.Body)
		s.queue.DeleteMessage(m.ReceiptHandle)
	}
}

func (s *SendMessageBatchSuite) TestSendMessageBatchError() {
	attrs := map[string]interface{}{
		"error": "",
	}

	batchMessages := []BatchMessage{
		BatchMessage{
			Body: "success",
		},
		BatchMessage{
			Body:    "failed",
			Options: []option.SendMessageRequest{option.MessageAttributes(attrs)},
		},
	}

	if err := s.queue.SendMessageBatch(batchMessages...); s.Error(err) {
		if merr, ok := err.(*multierror.Error); s.True(ok, "error must be *multierror.Error") {
			s.Len(merr.Errors, 1)
			if berr, ok := merr.Errors[0].(*BatchError); s.True(ok, "error must be *BatchError") {
				s.Equal(1, berr.Index, "batchMessages[1] must be error")
				s.Equal("InvalidParameterValue", berr.Code)
				s.Equal(true, berr.SenderFault)
			}
		}
	}

	messages, err := s.queue.ReceiveMessage(
		option.MaxNumberOfMessages(5),
		option.UseAllAttribute(),
	)
	if !s.NoError(err) {
		return
	}

	s.Len(messages, 1)
	for _, m := range messages {
		s.queue.DeleteMessage(m.ReceiptHandle)
	}
}

func TestSendMessageBatchSuite(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test")
	}

	suite.Run(t, new(SendMessageBatchSuite))
}
<|endoftext|>