<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package builder

import (
	"fmt"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"syscall"

	"mynewt.apache.org/newt/newt/parse"
	"mynewt.apache.org/newt/newt/pkg"
	"mynewt.apache.org/newt/newt/project"
	"mynewt.apache.org/newt/util"
)

func (t *TargetBuilder) Load(extraJtagCmd string) error {

	err := t.PrepBuild()

	if err != nil {
		return err
	}

	if t.LoaderBuilder != nil {
		err = t.AppBuilder.Load(1, extraJtagCmd)
		if err == nil {
			err = t.LoaderBuilder.Load(0, extraJtagCmd)
		}
	} else {
		err = t.AppBuilder.Load(0, extraJtagCmd)
	}

	return err
}

func RunOptionalCheck(checkScript string, env []string) error {
	if checkScript == "" {
		return nil
	}

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGTERM)

	cmd := []string{
		checkScript,
	}

	/* Handle Ctrl-C, terminate newt, as it is the
	   intended behavior */
	go func() {
		sig := <-sigs
		fmt.Println(sig)
		os.Exit(0)
	}()

	util.StatusMessage(util.VERBOSITY_SILENT,
		"Optional target check: %s\n", strings.Join(cmd, " "))
	util.ShellInteractiveCommand(cmd, env, true)

	/* Unregister SIGTERM handler */
	signal.Reset(syscall.SIGTERM)
	return nil
}

func Load(binBaseName string, bspPkg *pkg.BspPackage,
	extraEnvSettings map[string]string) error {

	if bspPkg.DownloadScript == "" {
		return nil
	}

	bspPath := bspPkg.BasePath()

	sortedKeys := make([]string, 0, len(extraEnvSettings))
	for k, _ := range extraEnvSettings {
		sortedKeys = append(sortedKeys, k)
	}
	sort.Strings(sortedKeys)

	env := []string{}
	for _, key := range sortedKeys {
		env = append(env, fmt.Sprintf("%s=%s", key, extraEnvSettings[key]))
	}

	coreRepo := project.GetProject().FindRepo("apache-mynewt-core")
	env = append(env, fmt.Sprintf("CORE_PATH=%s", coreRepo.Path()))
	env = append(env, fmt.Sprintf("BSP_PATH=%s", bspPath))
	env = append(env, fmt.Sprintf("BIN_BASENAME=%s", binBaseName))
	env = append(env, fmt.Sprintf("BIN_ROOT=%s", BinRoot()))
	env = append(env, fmt.Sprintf("MYNEWT_PROJECT_ROOT=%s", ProjectRoot()))

	RunOptionalCheck(bspPkg.OptChkScript, env)
	// bspPath, binBaseName are passed in command line for backwards
	// compatibility
	cmd := []string{
		bspPkg.DownloadScript,
		bspPath,
		binBaseName,
	}

	util.StatusMessage(util.VERBOSITY_VERBOSE, "Load command: %s\n",
		strings.Join(cmd, " "))
	util.StatusMessage(util.VERBOSITY_VERBOSE, "Environment:\n")
	for _, v := range env {
		util.StatusMessage(util.VERBOSITY_VERBOSE, "* %s\n", v)
	}
	if _, err := util.ShellCommand(cmd, env); err != nil {
		return err
	}
	util.StatusMessage(util.VERBOSITY_VERBOSE, "Successfully loaded image.\n")

	return nil
}

func (b *Builder) Load(imageSlot int, extraJtagCmd string) error {
	if b.appPkg == nil {
		return util.NewNewtError("app package not specified")
	}

	/* Populate the package list and feature sets. */
	err := b.targetBuilder.PrepBuild()
	if err != nil {
		return err
	}

	envSettings := map[string]string{
		"IMAGE_SLOT": strconv.Itoa(imageSlot),
		"FEATURES":   b.FeatureString(),
	}
	if extraJtagCmd != "" {
		envSettings["EXTRA_JTAG_CMD"] = extraJtagCmd
	}
	settings := b.cfg.SettingValues()

	var flashTargetArea string
	if parse.ValueIsTrue(settings["BOOT_LOADER"]) {
		envSettings["BOOT_LOADER"] = "1"

		flashTargetArea = "FLASH_AREA_BOOTLOADER"
		util.StatusMessage(util.VERBOSITY_DEFAULT,
			"Loading bootloader\n")
	} else {
		if imageSlot == 0 {
			flashTargetArea = "FLASH_AREA_IMAGE_0"
		} else if imageSlot == 1 {
			flashTargetArea = "FLASH_AREA_IMAGE_1"
		}
		util.StatusMessage(util.VERBOSITY_DEFAULT,
			"Loading %s image into slot %d\n", b.buildName, imageSlot+1)
	}

	bspPkg := b.targetBuilder.bspPkg
	tgtArea := bspPkg.FlashMap.Areas[flashTargetArea]
	if tgtArea.Name == "" {
		return util.NewNewtError(fmt.Sprintf("No flash target area %s\n",
			flashTargetArea))
	}
	envSettings["FLASH_OFFSET"] = "0x" + strconv.FormatInt(int64(tgtArea.Offset), 16)
	envSettings["FLASH_AREA_SIZE"] = "0x" + strconv.FormatInt(int64(tgtArea.Size), 16)

	// Add all syscfg settings to the environment with the MYNEWT_VAL_ prefix.
	for k, v := range settings {
		envSettings["MYNEWT_VAL_"+k] = v
	}

	// Convert the binary path from absolute to relative.  This is required for
	// compatibility with unix-in-windows environemnts (e.g., cygwin).
	binPath := util.TryRelPath(b.AppBinBasePath())
	if err := Load(binPath, b.targetBuilder.bspPkg, envSettings); err != nil {
		return err
	}

	return nil
}

func (t *TargetBuilder) Debug(extraJtagCmd string, reset bool, noGDB bool) error {
	if err := t.PrepBuild(); err != nil {
		return err
	}

	if t.LoaderBuilder == nil {
		return t.AppBuilder.Debug(extraJtagCmd, reset, noGDB)
	}
	return t.LoaderBuilder.Debug(extraJtagCmd, reset, noGDB)
}

func (b *Builder) debugBin(binPath string, extraJtagCmd string, reset bool,
	noGDB bool) error {
	/*
	 * Populate the package list and feature sets.
	 */
	err := b.targetBuilder.PrepBuild()
	if err != nil {
		return err
	}

	bspPath := b.bspPkg.rpkg.Lpkg.BasePath()
	binBaseName := binPath
	featureString := b.FeatureString()
	bspPkg := b.targetBuilder.bspPkg

	coreRepo := project.GetProject().FindRepo("apache-mynewt-core")
	envSettings := []string{
		fmt.Sprintf("CORE_PATH=%s", coreRepo.Path()),
		fmt.Sprintf("BSP_PATH=%s", bspPath),
		fmt.Sprintf("BIN_BASENAME=%s", binBaseName),
		fmt.Sprintf("FEATURES=%s", featureString),
		fmt.Sprintf("MYNEWT_PROJECT_ROOT=%s", ProjectRoot()),
	}
	if extraJtagCmd != "" {
		envSettings = append(envSettings,
			fmt.Sprintf("EXTRA_JTAG_CMD=%s", extraJtagCmd))
	}
	if reset == true {
		envSettings = append(envSettings, fmt.Sprintf("RESET=true"))
	}
	if noGDB == true {
		envSettings = append(envSettings, fmt.Sprintf("NO_GDB=1"))
	}

	os.Chdir(project.GetProject().Path())

	RunOptionalCheck(bspPkg.OptChkScript, envSettings)
	// bspPath, binBaseName are passed in command line for backwards
	// compatibility
	cmdLine := []string{
		b.targetBuilder.bspPkg.DebugScript, bspPath, binBaseName,
	}

	fmt.Printf("%s\n", cmdLine)
	return util.ShellInteractiveCommand(cmdLine, envSettings, false)
}

func (b *Builder) Debug(extraJtagCmd string, reset bool, noGDB bool) error {
	if b.appPkg == nil {
		return util.NewNewtError("app package not specified")
	}

	// Convert the binary path from absolute to relative.  This is required for
	// Windows compatibility.
	binPath := util.TryRelPath(b.AppBinBasePath())
	return b.debugBin(binPath, extraJtagCmd, reset, noGDB)
}
<commit_msg>Allow BSP to override scripts with syscfg<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package builder

import (
	"fmt"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"syscall"

	"mynewt.apache.org/newt/newt/parse"
	"mynewt.apache.org/newt/newt/pkg"
	"mynewt.apache.org/newt/newt/project"
	"mynewt.apache.org/newt/util"
)

func (t *TargetBuilder) loadLoader(slot int, extraJtagCmd string) error {
	if err := t.bspPkg.Reload(t.LoaderBuilder.cfg.SettingValues()); err != nil {
		return err
	}

	return t.LoaderBuilder.Load(slot, extraJtagCmd)
}

func (t *TargetBuilder) loadApp(slot int, extraJtagCmd string) error {
	if err := t.bspPkg.Reload(t.AppBuilder.cfg.SettingValues()); err != nil {
		return err
	}

	return t.AppBuilder.Load(slot, extraJtagCmd)
}

func (t *TargetBuilder) debugLoader(extraJtagCmd string, reset bool,
	noGDB bool) error {

	if err := t.bspPkg.Reload(t.LoaderBuilder.cfg.SettingValues()); err != nil {
		return err
	}

	return t.LoaderBuilder.Debug(extraJtagCmd, reset, noGDB)
}

func (t *TargetBuilder) debugApp(extraJtagCmd string, reset bool,
	noGDB bool) error {

	if err := t.bspPkg.Reload(t.AppBuilder.cfg.SettingValues()); err != nil {
		return err
	}

	return t.AppBuilder.Debug(extraJtagCmd, reset, noGDB)
}

func (t *TargetBuilder) Load(extraJtagCmd string) error {
	err := t.PrepBuild()
	if err != nil {
		return err
	}

	if t.LoaderBuilder != nil {
		err = t.loadApp(1, extraJtagCmd)
		if err == nil {
			err = t.loadLoader(0, extraJtagCmd)
		}
	} else {
		err = t.loadApp(0, extraJtagCmd)
	}

	return err
}

func RunOptionalCheck(checkScript string, env []string) error {
	if checkScript == "" {
		return nil
	}

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGTERM)

	cmd := []string{
		checkScript,
	}

	/* Handle Ctrl-C, terminate newt, as it is the
	   intended behavior */
	go func() {
		sig := <-sigs
		fmt.Println(sig)
		os.Exit(0)
	}()

	util.StatusMessage(util.VERBOSITY_SILENT,
		"Optional target check: %s\n", strings.Join(cmd, " "))
	util.ShellInteractiveCommand(cmd, env, true)

	/* Unregister SIGTERM handler */
	signal.Reset(syscall.SIGTERM)
	return nil
}

func Load(binBaseName string, bspPkg *pkg.BspPackage,
	extraEnvSettings map[string]string) error {

	if bspPkg.DownloadScript == "" {
		return nil
	}

	bspPath := bspPkg.BasePath()

	sortedKeys := make([]string, 0, len(extraEnvSettings))
	for k, _ := range extraEnvSettings {
		sortedKeys = append(sortedKeys, k)
	}
	sort.Strings(sortedKeys)

	env := []string{}
	for _, key := range sortedKeys {
		env = append(env, fmt.Sprintf("%s=%s", key, extraEnvSettings[key]))
	}

	coreRepo := project.GetProject().FindRepo("apache-mynewt-core")
	env = append(env, fmt.Sprintf("CORE_PATH=%s", coreRepo.Path()))
	env = append(env, fmt.Sprintf("BSP_PATH=%s", bspPath))
	env = append(env, fmt.Sprintf("BIN_BASENAME=%s", binBaseName))
	env = append(env, fmt.Sprintf("BIN_ROOT=%s", BinRoot()))
	env = append(env, fmt.Sprintf("MYNEWT_PROJECT_ROOT=%s", ProjectRoot()))

	RunOptionalCheck(bspPkg.OptChkScript, env)
	// bspPath, binBaseName are passed in command line for backwards
	// compatibility
	cmd := []string{
		bspPkg.DownloadScript,
		bspPath,
		binBaseName,
	}

	util.StatusMessage(util.VERBOSITY_VERBOSE, "Load command: %s\n",
		strings.Join(cmd, " "))
	util.StatusMessage(util.VERBOSITY_VERBOSE, "Environment:\n")
	for _, v := range env {
		util.StatusMessage(util.VERBOSITY_VERBOSE, "* %s\n", v)
	}
	if _, err := util.ShellCommand(cmd, env); err != nil {
		return err
	}
	util.StatusMessage(util.VERBOSITY_VERBOSE, "Successfully loaded image.\n")

	return nil
}

func (b *Builder) Load(imageSlot int, extraJtagCmd string) error {
	if b.appPkg == nil {
		return util.NewNewtError("app package not specified")
	}

	/* Populate the package list and feature sets. */
	err := b.targetBuilder.PrepBuild()
	if err != nil {
		return err
	}

	envSettings := map[string]string{
		"IMAGE_SLOT": strconv.Itoa(imageSlot),
		"FEATURES":   b.FeatureString(),
	}
	if extraJtagCmd != "" {
		envSettings["EXTRA_JTAG_CMD"] = extraJtagCmd
	}
	settings := b.cfg.SettingValues()

	var flashTargetArea string
	if parse.ValueIsTrue(settings["BOOT_LOADER"]) {
		envSettings["BOOT_LOADER"] = "1"

		flashTargetArea = "FLASH_AREA_BOOTLOADER"
		util.StatusMessage(util.VERBOSITY_DEFAULT,
			"Loading bootloader\n")
	} else {
		if imageSlot == 0 {
			flashTargetArea = "FLASH_AREA_IMAGE_0"
		} else if imageSlot == 1 {
			flashTargetArea = "FLASH_AREA_IMAGE_1"
		}
		util.StatusMessage(util.VERBOSITY_DEFAULT,
			"Loading %s image into slot %d\n", b.buildName, imageSlot+1)
	}

	bspPkg := b.targetBuilder.bspPkg
	tgtArea := bspPkg.FlashMap.Areas[flashTargetArea]
	if tgtArea.Name == "" {
		return util.NewNewtError(fmt.Sprintf("No flash target area %s\n",
			flashTargetArea))
	}
	envSettings["FLASH_OFFSET"] = "0x" + strconv.FormatInt(int64(tgtArea.Offset), 16)
	envSettings["FLASH_AREA_SIZE"] = "0x" + strconv.FormatInt(int64(tgtArea.Size), 16)

	// Add all syscfg settings to the environment with the MYNEWT_VAL_ prefix.
	for k, v := range settings {
		envSettings["MYNEWT_VAL_"+k] = v
	}

	// Convert the binary path from absolute to relative.  This is required for
	// compatibility with unix-in-windows environemnts (e.g., cygwin).
	binPath := util.TryRelPath(b.AppBinBasePath())
	if err := Load(binPath, b.targetBuilder.bspPkg, envSettings); err != nil {
		return err
	}

	return nil
}

func (t *TargetBuilder) Debug(extraJtagCmd string, reset bool, noGDB bool) error {
	if err := t.PrepBuild(); err != nil {
		return err
	}

	if t.LoaderBuilder == nil {
		return t.debugApp(extraJtagCmd, reset, noGDB)
	}
	return t.debugLoader(extraJtagCmd, reset, noGDB)
}

func (b *Builder) debugBin(binPath string, extraJtagCmd string, reset bool,
	noGDB bool) error {
	/*
	 * Populate the package list and feature sets.
	 */
	err := b.targetBuilder.PrepBuild()
	if err != nil {
		return err
	}

	bspPath := b.bspPkg.rpkg.Lpkg.BasePath()
	binBaseName := binPath
	featureString := b.FeatureString()
	bspPkg := b.targetBuilder.bspPkg

	coreRepo := project.GetProject().FindRepo("apache-mynewt-core")
	envSettings := []string{
		fmt.Sprintf("CORE_PATH=%s", coreRepo.Path()),
		fmt.Sprintf("BSP_PATH=%s", bspPath),
		fmt.Sprintf("BIN_BASENAME=%s", binBaseName),
		fmt.Sprintf("FEATURES=%s", featureString),
		fmt.Sprintf("MYNEWT_PROJECT_ROOT=%s", ProjectRoot()),
	}
	if extraJtagCmd != "" {
		envSettings = append(envSettings,
			fmt.Sprintf("EXTRA_JTAG_CMD=%s", extraJtagCmd))
	}
	if reset == true {
		envSettings = append(envSettings, fmt.Sprintf("RESET=true"))
	}
	if noGDB == true {
		envSettings = append(envSettings, fmt.Sprintf("NO_GDB=1"))
	}

	os.Chdir(project.GetProject().Path())

	RunOptionalCheck(bspPkg.OptChkScript, envSettings)
	// bspPath, binBaseName are passed in command line for backwards
	// compatibility
	cmdLine := []string{
		b.targetBuilder.bspPkg.DebugScript, bspPath, binBaseName,
	}

	fmt.Printf("%s\n", cmdLine)
	return util.ShellInteractiveCommand(cmdLine, envSettings, false)
}

func (b *Builder) Debug(extraJtagCmd string, reset bool, noGDB bool) error {
	if b.appPkg == nil {
		return util.NewNewtError("app package not specified")
	}

	// Convert the binary path from absolute to relative.  This is required for
	// Windows compatibility.
	binPath := util.TryRelPath(b.AppBinBasePath())
	return b.debugBin(binPath, extraJtagCmd, reset, noGDB)
}
<|endoftext|>


<commit_before>package test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"testing"
	"time"

	"github.com/gruntwork-io/terratest/modules/aws"
	"github.com/gruntwork-io/terratest/modules/logger"
	"github.com/gruntwork-io/terratest/modules/random"
	"github.com/gruntwork-io/terratest/modules/retry"
	"github.com/gruntwork-io/terratest/modules/terraform"
	"github.com/gruntwork-io/terratest/modules/test-structure"
)

const REPO_ROOT = "../"

const ENV_VAR_AWS_REGION = "AWS_DEFAULT_REGION"

const VAR_AMI_ID = "ami_id"

const CLUSTER_COLOCATED_EXAMPLE_PATH = "/"
const CLUSTER_COLOCATED_EXAMPLE_VAR_CLUSTER_NAME = "cluster_name"
const CLUSTER_COLOCATED_EXAMPLE_VAR_NUM_SERVERS = "num_servers"
const CLUSTER_COLOCATED_EXAMPLE_VAR_NUM_CLIENTS = "num_clients"
const CLUSTER_COLOCATED_EXAMPLE_OUTPUT_SERVER_ASG_NAME = "asg_name_servers"

const CLUSTER_SEPARATE_EXAMPLE_PATH = "examples/nomad-consul-separate-cluster"
const CLUSTER_SEPARATE_EXAMPLE_VAR_NOMAD_CLUSTER_NAME = "nomad_cluster_name"
const CLUSTER_SEPARATE_EXAMPLE_VAR_CONSUL_CLUSTER_NAME = "consul_cluster_name"
const CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_NOMAD_SERVERS = "num_nomad_servers"
const CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_CONSUL_SERVERS = "num_consul_servers"
const CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_NOMAD_CLIENTS = "num_nomad_clients"
const CLUSTER_SEPARATE_EXAMPLE_OUTPUT_NOMAD_SERVER_ASG_NAME = "asg_name_nomad_servers"

const DEFAULT_NUM_SERVERS = 3
const DEFAULT_NUM_CLIENTS = 6

const SAVED_AWS_REGION = "AwsRegion"
const SAVED_UNIQUE_ID = "UniqueId"

// Test the Nomad/Consul colocated cluster example by:
//
// 1. Copying the code in this repo to a temp folder so tests on the Terraform code can run in parallel without the
//    state files overwriting each other.
// 2. Building the AMI in the nomad-consul-ami example with the given build name
// 3. Deploying that AMI using the example Terraform code
// 4. Checking that the Nomad cluster comes up within a reasonable time period and can respond to requests
func runNomadClusterColocatedTest(t *testing.T, packerBuildName string) {
	examplesDir := test_structure.CopyTerraformFolderToTemp(t, REPO_ROOT, CLUSTER_COLOCATED_EXAMPLE_PATH)

	defer test_structure.RunTestStage(t, "teardown", func() {
		terraformOptions := test_structure.LoadTerraformOptions(t, examplesDir)
		terraform.Destroy(t, terraformOptions)

		amiId := test_structure.LoadAmiId(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)
		aws.DeleteAmi(t, awsRegion, amiId)
	})

	test_structure.RunTestStage(t, "setup_ami", func() {
		awsRegion := getRandomRegion(t)
		test_structure.SaveString(t, examplesDir, SAVED_AWS_REGION, awsRegion)

		uniqueId := random.UniqueId()
		test_structure.SaveString(t, examplesDir, SAVED_UNIQUE_ID, uniqueId)

		amiId := buildAmi(t, filepath.Join(examplesDir, "examples", "nomad-consul-ami", "nomad-consul.json"), packerBuildName, awsRegion, uniqueId)
		test_structure.SaveAmiId(t, examplesDir, amiId)
	})

	test_structure.RunTestStage(t, "deploy", func() {
		amiId := test_structure.LoadAmiId(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)
		uniqueId := test_structure.LoadString(t, examplesDir, SAVED_UNIQUE_ID)

		terraformOptions := &terraform.Options{
			TerraformDir: examplesDir,
			Vars: map[string]interface{}{
				CLUSTER_COLOCATED_EXAMPLE_VAR_CLUSTER_NAME: fmt.Sprintf("test-%s", uniqueId),
				CLUSTER_COLOCATED_EXAMPLE_VAR_NUM_SERVERS:  DEFAULT_NUM_SERVERS,
				CLUSTER_COLOCATED_EXAMPLE_VAR_NUM_CLIENTS:  DEFAULT_NUM_CLIENTS,
				VAR_AMI_ID: amiId,
			},
			EnvVars: map[string]string{
				ENV_VAR_AWS_REGION: awsRegion,
			},
		}
		test_structure.SaveTerraformOptions(t, examplesDir, terraformOptions)

		terraform.InitAndApply(t, terraformOptions)
	})

	test_structure.RunTestStage(t, "validate", func() {
		terraformOptions := test_structure.LoadTerraformOptions(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)

		checkNomadClusterIsWorking(t, CLUSTER_COLOCATED_EXAMPLE_OUTPUT_SERVER_ASG_NAME, terraformOptions, awsRegion)
	})
}

// Test the Nomad/Consul separate clusters example by:
//
// 1. Copying the code in this repo to a temp folder so tests on the Terraform code can run in parallel without the
//    state files overwriting each other.
// 2. Building the AMI in the nomad-consul-ami example with the given build name
// 3. Deploying that AMI using the example Terraform code
// 4. Checking that the Nomad cluster comes up within a reasonable time period and can respond to requests
func runNomadClusterSeparateTest(t *testing.T, packerBuildName string) {
	examplesDir := test_structure.CopyTerraformFolderToTemp(t, REPO_ROOT, "/")

	defer test_structure.RunTestStage(t, "teardown", func() {
		terraformOptions := test_structure.LoadTerraformOptions(t, examplesDir)
		terraform.Destroy(t, terraformOptions)

		amiId := test_structure.LoadAmiId(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)
		aws.DeleteAmi(t, awsRegion, amiId)
	})

	test_structure.RunTestStage(t, "setup_ami", func() {
		awsRegion := getRandomRegion(t)
		test_structure.SaveString(t, examplesDir, SAVED_AWS_REGION, awsRegion)

		uniqueId := random.UniqueId()
		test_structure.SaveString(t, examplesDir, SAVED_UNIQUE_ID, uniqueId)

		amiId := buildAmi(t, filepath.Join(examplesDir, "examples", "nomad-consul-ami", "nomad-consul.json"), packerBuildName, awsRegion, uniqueId)
		test_structure.SaveAmiId(t, examplesDir, amiId)
	})

	test_structure.RunTestStage(t, "deploy", func() {
		amiId := test_structure.LoadAmiId(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)
		uniqueId := test_structure.LoadString(t, examplesDir, SAVED_UNIQUE_ID)

		terraformOptions := &terraform.Options{
			TerraformDir: filepath.Join(examplesDir, "examples", "nomad-consul-separate-cluster"),
			Vars: map[string]interface{}{
				CLUSTER_SEPARATE_EXAMPLE_VAR_NOMAD_CLUSTER_NAME:  fmt.Sprintf("test-%s", uniqueId),
				CLUSTER_SEPARATE_EXAMPLE_VAR_CONSUL_CLUSTER_NAME: fmt.Sprintf("test-%s", uniqueId),
				CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_NOMAD_SERVERS:   DEFAULT_NUM_SERVERS,
				CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_CONSUL_SERVERS:  DEFAULT_NUM_SERVERS,
				CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_NOMAD_CLIENTS:   DEFAULT_NUM_CLIENTS,
				VAR_AMI_ID: amiId,
			},
			EnvVars: map[string]string{
				ENV_VAR_AWS_REGION: awsRegion,
			},
		}
		test_structure.SaveTerraformOptions(t, examplesDir, terraformOptions)

		terraform.InitAndApply(t, terraformOptions)
	})

	test_structure.RunTestStage(t, "validate", func() {
		terraformOptions := test_structure.LoadTerraformOptions(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)

		checkNomadClusterIsWorking(t, CLUSTER_SEPARATE_EXAMPLE_OUTPUT_NOMAD_SERVER_ASG_NAME, terraformOptions, awsRegion)
	})
}

// Check that the Nomad cluster comes up within a reasonable time period and can respond to requests
func checkNomadClusterIsWorking(t *testing.T, asgNameOutputVar string, terraformOptions *terraform.Options, awsRegion string) {
	asgName := terraform.Output(t, terraformOptions, asgNameOutputVar)
	nodeIpAddress := getIpAddressOfAsgInstance(t, asgName, awsRegion)
	testNomadCluster(t, nodeIpAddress)
}

// Use a Nomad client to connect to the given node and use it to verify that:
//
// 1. The Nomad cluster has deployed
// 2. The cluster has the expected number of server nodes
// 2. The cluster has the expected number of client nodes
func testNomadCluster(t *testing.T, nodeIpAddress string) {
	maxRetries := 90
	sleepBetweenRetries := 10 * time.Second

	response := retry.DoWithRetry(t, "Check Nomad cluster has expected number of servers and clients", maxRetries, sleepBetweenRetries, func() (string, error) {
		clients, err := callNomadApi(t, nodeIpAddress, "v1/nodes")
		if err != nil {
			return "", err
		}

		if len(clients) != DEFAULT_NUM_CLIENTS {
			return "", fmt.Errorf("Expected the cluster to have %d clients, but found %d", DEFAULT_NUM_CLIENTS, len(clients))
		}

		servers, err := callNomadApi(t, nodeIpAddress, "v1/status/peers")
		if err != nil {
			return "", err
		}

		if len(servers) != DEFAULT_NUM_SERVERS {
			return "", fmt.Errorf("Expected the cluster to have %d servers, but found %d", DEFAULT_NUM_SERVERS, len(servers))
		}

		return fmt.Sprintf("Got back expected number of clients (%d) and servers (%d)", len(clients), len(servers)), nil
	})

	logger.Logf(t, "Nomad cluster is properly deployed: %s", response)
}

// A quick, hacky way to call the Nomad HTTP API: https://www.nomadproject.io/docs/http/index.html
func callNomadApi(t *testing.T, nodeIpAddress string, path string) ([]interface{}, error) {
	url := fmt.Sprintf("http://%s:4646/%s", nodeIpAddress, path)
	logger.Logf(t, "Making an HTTP GET to URL %s", url)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	logger.Logf(t, "Response from Nomad for URL %s: %s", url, string(body))

	result := []interface{}{}
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, err
	}

	return result, nil
}
<commit_msg>Use different tags for different clusters<commit_after>package test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"testing"
	"time"

	"github.com/gruntwork-io/terratest/modules/aws"
	"github.com/gruntwork-io/terratest/modules/logger"
	"github.com/gruntwork-io/terratest/modules/random"
	"github.com/gruntwork-io/terratest/modules/retry"
	"github.com/gruntwork-io/terratest/modules/terraform"
	"github.com/gruntwork-io/terratest/modules/test-structure"
)

const REPO_ROOT = "../"

const ENV_VAR_AWS_REGION = "AWS_DEFAULT_REGION"

const VAR_AMI_ID = "ami_id"

const CLUSTER_COLOCATED_EXAMPLE_PATH = "/"
const CLUSTER_COLOCATED_EXAMPLE_VAR_CLUSTER_NAME = "cluster_name"
const CLUSTER_COLOCATED_EXAMPLE_VAR_CLUSTER_TAG_VALUE = "cluster_tag_value"
const CLUSTER_COLOCATED_EXAMPLE_VAR_NUM_SERVERS = "num_servers"
const CLUSTER_COLOCATED_EXAMPLE_VAR_NUM_CLIENTS = "num_clients"
const CLUSTER_COLOCATED_EXAMPLE_OUTPUT_SERVER_ASG_NAME = "asg_name_servers"

const CLUSTER_SEPARATE_EXAMPLE_PATH = "examples/nomad-consul-separate-cluster"
const CLUSTER_SEPARATE_EXAMPLE_VAR_NOMAD_CLUSTER_NAME = "nomad_cluster_name"
const CLUSTER_SEPARATE_EXAMPLE_VAR_CONSUL_CLUSTER_NAME = "consul_cluster_name"
const CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_NOMAD_SERVERS = "num_nomad_servers"
const CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_CONSUL_SERVERS = "num_consul_servers"
const CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_NOMAD_CLIENTS = "num_nomad_clients"
const CLUSTER_SEPARATE_EXAMPLE_OUTPUT_NOMAD_SERVER_ASG_NAME = "asg_name_nomad_servers"

const DEFAULT_NUM_SERVERS = 3
const DEFAULT_NUM_CLIENTS = 6

const SAVED_AWS_REGION = "AwsRegion"
const SAVED_UNIQUE_ID = "UniqueId"

// Test the Nomad/Consul colocated cluster example by:
//
// 1. Copying the code in this repo to a temp folder so tests on the Terraform code can run in parallel without the
//    state files overwriting each other.
// 2. Building the AMI in the nomad-consul-ami example with the given build name
// 3. Deploying that AMI using the example Terraform code
// 4. Checking that the Nomad cluster comes up within a reasonable time period and can respond to requests
func runNomadClusterColocatedTest(t *testing.T, packerBuildName string) {
	examplesDir := test_structure.CopyTerraformFolderToTemp(t, REPO_ROOT, CLUSTER_COLOCATED_EXAMPLE_PATH)

	defer test_structure.RunTestStage(t, "teardown", func() {
		terraformOptions := test_structure.LoadTerraformOptions(t, examplesDir)
		terraform.Destroy(t, terraformOptions)

		amiId := test_structure.LoadAmiId(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)
		aws.DeleteAmi(t, awsRegion, amiId)
	})

	test_structure.RunTestStage(t, "setup_ami", func() {
		awsRegion := getRandomRegion(t)
		test_structure.SaveString(t, examplesDir, SAVED_AWS_REGION, awsRegion)

		uniqueId := random.UniqueId()
		test_structure.SaveString(t, examplesDir, SAVED_UNIQUE_ID, uniqueId)

		amiId := buildAmi(t, filepath.Join(examplesDir, "examples", "nomad-consul-ami", "nomad-consul.json"), packerBuildName, awsRegion, uniqueId)
		test_structure.SaveAmiId(t, examplesDir, amiId)
	})

	test_structure.RunTestStage(t, "deploy", func() {
		amiId := test_structure.LoadAmiId(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)
		uniqueId := test_structure.LoadString(t, examplesDir, SAVED_UNIQUE_ID)

		terraformOptions := &terraform.Options{
			TerraformDir: examplesDir,
			Vars: map[string]interface{}{
				CLUSTER_COLOCATED_EXAMPLE_VAR_CLUSTER_NAME:      fmt.Sprintf("test-%s", uniqueId),
				CLUSTER_COLOCATED_EXAMPLE_VAR_CLUSTER_TAG_VALUE: fmt.Sprintf("auto-join-%s", uniqueId),
				CLUSTER_COLOCATED_EXAMPLE_VAR_NUM_SERVERS:       DEFAULT_NUM_SERVERS,
				CLUSTER_COLOCATED_EXAMPLE_VAR_NUM_CLIENTS:       DEFAULT_NUM_CLIENTS,
				VAR_AMI_ID: amiId,
			},
			EnvVars: map[string]string{
				ENV_VAR_AWS_REGION: awsRegion,
			},
		}
		test_structure.SaveTerraformOptions(t, examplesDir, terraformOptions)

		terraform.InitAndApply(t, terraformOptions)
	})

	test_structure.RunTestStage(t, "validate", func() {
		terraformOptions := test_structure.LoadTerraformOptions(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)

		checkNomadClusterIsWorking(t, CLUSTER_COLOCATED_EXAMPLE_OUTPUT_SERVER_ASG_NAME, terraformOptions, awsRegion)
	})
}

// Test the Nomad/Consul separate clusters example by:
//
// 1. Copying the code in this repo to a temp folder so tests on the Terraform code can run in parallel without the
//    state files overwriting each other.
// 2. Building the AMI in the nomad-consul-ami example with the given build name
// 3. Deploying that AMI using the example Terraform code
// 4. Checking that the Nomad cluster comes up within a reasonable time period and can respond to requests
func runNomadClusterSeparateTest(t *testing.T, packerBuildName string) {
	examplesDir := test_structure.CopyTerraformFolderToTemp(t, REPO_ROOT, "/")

	defer test_structure.RunTestStage(t, "teardown", func() {
		terraformOptions := test_structure.LoadTerraformOptions(t, examplesDir)
		terraform.Destroy(t, terraformOptions)

		amiId := test_structure.LoadAmiId(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)
		aws.DeleteAmi(t, awsRegion, amiId)
	})

	test_structure.RunTestStage(t, "setup_ami", func() {
		awsRegion := getRandomRegion(t)
		test_structure.SaveString(t, examplesDir, SAVED_AWS_REGION, awsRegion)

		uniqueId := random.UniqueId()
		test_structure.SaveString(t, examplesDir, SAVED_UNIQUE_ID, uniqueId)

		amiId := buildAmi(t, filepath.Join(examplesDir, "examples", "nomad-consul-ami", "nomad-consul.json"), packerBuildName, awsRegion, uniqueId)
		test_structure.SaveAmiId(t, examplesDir, amiId)
	})

	test_structure.RunTestStage(t, "deploy", func() {
		amiId := test_structure.LoadAmiId(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)
		uniqueId := test_structure.LoadString(t, examplesDir, SAVED_UNIQUE_ID)

		terraformOptions := &terraform.Options{
			TerraformDir: filepath.Join(examplesDir, "examples", "nomad-consul-separate-cluster"),
			Vars: map[string]interface{}{
				CLUSTER_SEPARATE_EXAMPLE_VAR_NOMAD_CLUSTER_NAME:  fmt.Sprintf("test-%s", uniqueId),
				CLUSTER_SEPARATE_EXAMPLE_VAR_CONSUL_CLUSTER_NAME: fmt.Sprintf("test-%s", uniqueId),
				CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_NOMAD_SERVERS:   DEFAULT_NUM_SERVERS,
				CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_CONSUL_SERVERS:  DEFAULT_NUM_SERVERS,
				CLUSTER_SEPARATE_EXAMPLE_VAR_NUM_NOMAD_CLIENTS:   DEFAULT_NUM_CLIENTS,
				VAR_AMI_ID: amiId,
			},
			EnvVars: map[string]string{
				ENV_VAR_AWS_REGION: awsRegion,
			},
		}
		test_structure.SaveTerraformOptions(t, examplesDir, terraformOptions)

		terraform.InitAndApply(t, terraformOptions)
	})

	test_structure.RunTestStage(t, "validate", func() {
		terraformOptions := test_structure.LoadTerraformOptions(t, examplesDir)
		awsRegion := test_structure.LoadString(t, examplesDir, SAVED_AWS_REGION)

		checkNomadClusterIsWorking(t, CLUSTER_SEPARATE_EXAMPLE_OUTPUT_NOMAD_SERVER_ASG_NAME, terraformOptions, awsRegion)
	})
}

// Check that the Nomad cluster comes up within a reasonable time period and can respond to requests
func checkNomadClusterIsWorking(t *testing.T, asgNameOutputVar string, terraformOptions *terraform.Options, awsRegion string) {
	asgName := terraform.Output(t, terraformOptions, asgNameOutputVar)
	nodeIpAddress := getIpAddressOfAsgInstance(t, asgName, awsRegion)
	testNomadCluster(t, nodeIpAddress)
}

// Use a Nomad client to connect to the given node and use it to verify that:
//
// 1. The Nomad cluster has deployed
// 2. The cluster has the expected number of server nodes
// 2. The cluster has the expected number of client nodes
func testNomadCluster(t *testing.T, nodeIpAddress string) {
	maxRetries := 90
	sleepBetweenRetries := 10 * time.Second

	response := retry.DoWithRetry(t, "Check Nomad cluster has expected number of servers and clients", maxRetries, sleepBetweenRetries, func() (string, error) {
		clients, err := callNomadApi(t, nodeIpAddress, "v1/nodes")
		if err != nil {
			return "", err
		}

		if len(clients) != DEFAULT_NUM_CLIENTS {
			return "", fmt.Errorf("Expected the cluster to have %d clients, but found %d", DEFAULT_NUM_CLIENTS, len(clients))
		}

		servers, err := callNomadApi(t, nodeIpAddress, "v1/status/peers")
		if err != nil {
			return "", err
		}

		if len(servers) != DEFAULT_NUM_SERVERS {
			return "", fmt.Errorf("Expected the cluster to have %d servers, but found %d", DEFAULT_NUM_SERVERS, len(servers))
		}

		return fmt.Sprintf("Got back expected number of clients (%d) and servers (%d)", len(clients), len(servers)), nil
	})

	logger.Logf(t, "Nomad cluster is properly deployed: %s", response)
}

// A quick, hacky way to call the Nomad HTTP API: https://www.nomadproject.io/docs/http/index.html
func callNomadApi(t *testing.T, nodeIpAddress string, path string) ([]interface{}, error) {
	url := fmt.Sprintf("http://%s:4646/%s", nodeIpAddress, path)
	logger.Logf(t, "Making an HTTP GET to URL %s", url)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	logger.Logf(t, "Response from Nomad for URL %s: %s", url, string(body))

	result := []interface{}{}
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, err
	}

	return result, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"io"
	"time"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
)

// BlockVersion is the current latest supported block version.
const BlockVersion = 4

// MaxBlockHeaderPayload is the maximum number of bytes a block header can be.
// Version 4 bytes + Timestamp 4 bytes + Bits 4 bytes + Nonce 4 bytes +
// PrevBlock and MerkleRoot hashes.
const MaxBlockHeaderPayload = 16 + (chainhash.HashSize * 2)

// BlockHeader defines information about a block and is used in the bitcoin
// block (MsgBlock) and headers (MsgHeaders) messages.
type BlockHeader struct {
	// Version of the block.  This is not the same as the protocol version.
	Version int32

	// Hash of the previous block in the block chain.
	PrevBlock chainhash.Hash

	// Merkle tree reference to hash of all transactions for the block.
	MerkleRoot chainhash.Hash

	// Time the block was created.  This is, unfortunately, encoded as a
	// uint32 on the wire and therefore is limited to 2106.
	Timestamp time.Time

	// Difficulty target for the block.
	Bits uint32

	// Nonce used to generate the block.
	Nonce uint32
}

// blockHeaderLen is a constant that represents the number of bytes for a block
// header.
const blockHeaderLen = 80

// BlockHash computes the block identifier hash for the given block header.
func (h *BlockHeader) BlockHash() chainhash.Hash {
	// Encode the header and double sha256 everything prior to the number of
	// transactions.  Ignore the error returns since there is no way the
	// encode could fail except being out of memory which would cause a
	// run-time panic.
	var buf bytes.Buffer
	_ = writeBlockHeader(&buf, 0, h)

	return chainhash.DoubleHashH(buf.Bytes())
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
// See Deserialize for decoding block headers stored to disk, such as in a
// database, as opposed to decoding block headers from the wire.
func (h *BlockHeader) BtcDecode(r io.Reader, pver uint32) error {
	return readBlockHeader(r, pver, h)
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
// See Serialize for encoding block headers to be stored to disk, such as in a
// database, as opposed to encoding block headers for the wire.
func (h *BlockHeader) BtcEncode(w io.Writer, pver uint32) error {
	return writeBlockHeader(w, pver, h)
}

// Deserialize decodes a block header from r into the receiver using a format
// that is suitable for long-term storage such as a database while respecting
// the Version field.
func (h *BlockHeader) Deserialize(r io.Reader) error {
	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of readBlockHeader.
	return readBlockHeader(r, 0, h)
}

// Serialize encodes a block header from r into the receiver using a format
// that is suitable for long-term storage such as a database while respecting
// the Version field.
func (h *BlockHeader) Serialize(w io.Writer) error {
	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of writeBlockHeader.
	return writeBlockHeader(w, 0, h)
}

// NewBlockHeader returns a new BlockHeader using the provided previous block
// hash, merkle root hash, difficulty bits, and nonce used to generate the
// block with defaults for the remaining fields.
func NewBlockHeader(prevHash *chainhash.Hash, merkleRootHash *chainhash.Hash,
	bits uint32, nonce uint32) *BlockHeader {

	// Limit the timestamp to one second precision since the protocol
	// doesn't support better.
	return &BlockHeader{
		Version:    BlockVersion,
		PrevBlock:  *prevHash,
		MerkleRoot: *merkleRootHash,
		Timestamp:  time.Unix(time.Now().Unix(), 0),
		Bits:       bits,
		Nonce:      nonce,
	}
}

// readBlockHeader reads a bitcoin block header from r.  See Deserialize for
// decoding block headers stored to disk, such as in a database, as opposed to
// decoding from the wire.
func readBlockHeader(r io.Reader, pver uint32, bh *BlockHeader) error {
	return readElements(r, &bh.Version, &bh.PrevBlock, &bh.MerkleRoot,
		(*uint32Time)(&bh.Timestamp), &bh.Bits, &bh.Nonce)
}

// writeBlockHeader writes a bitcoin block header to w.  See Serialize for
// encoding block headers to be stored to disk, such as in a database, as
// opposed to encoding for the wire.
func writeBlockHeader(w io.Writer, pver uint32, bh *BlockHeader) error {
	sec := uint32(bh.Timestamp.Unix())
	return writeElements(w, bh.Version, &bh.PrevBlock, &bh.MerkleRoot,
		sec, bh.Bits, bh.Nonce)
}
<commit_msg>wire: reduce allocs in BlockHash()<commit_after>// Copyright (c) 2013-2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"io"
	"time"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
)

// BlockVersion is the current latest supported block version.
const BlockVersion = 4

// MaxBlockHeaderPayload is the maximum number of bytes a block header can be.
// Version 4 bytes + Timestamp 4 bytes + Bits 4 bytes + Nonce 4 bytes +
// PrevBlock and MerkleRoot hashes.
const MaxBlockHeaderPayload = 16 + (chainhash.HashSize * 2)

// BlockHeader defines information about a block and is used in the bitcoin
// block (MsgBlock) and headers (MsgHeaders) messages.
type BlockHeader struct {
	// Version of the block.  This is not the same as the protocol version.
	Version int32

	// Hash of the previous block in the block chain.
	PrevBlock chainhash.Hash

	// Merkle tree reference to hash of all transactions for the block.
	MerkleRoot chainhash.Hash

	// Time the block was created.  This is, unfortunately, encoded as a
	// uint32 on the wire and therefore is limited to 2106.
	Timestamp time.Time

	// Difficulty target for the block.
	Bits uint32

	// Nonce used to generate the block.
	Nonce uint32
}

// blockHeaderLen is a constant that represents the number of bytes for a block
// header.
const blockHeaderLen = 80

// BlockHash computes the block identifier hash for the given block header.
func (h *BlockHeader) BlockHash() chainhash.Hash {
	// Encode the header and double sha256 everything prior to the number of
	// transactions.  Ignore the error returns since there is no way the
	// encode could fail except being out of memory which would cause a
	// run-time panic.
	buf := bytes.NewBuffer(make([]byte, 0, MaxBlockHeaderPayload))
	_ = writeBlockHeader(buf, 0, h)

	return chainhash.DoubleHashH(buf.Bytes())
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
// See Deserialize for decoding block headers stored to disk, such as in a
// database, as opposed to decoding block headers from the wire.
func (h *BlockHeader) BtcDecode(r io.Reader, pver uint32) error {
	return readBlockHeader(r, pver, h)
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
// See Serialize for encoding block headers to be stored to disk, such as in a
// database, as opposed to encoding block headers for the wire.
func (h *BlockHeader) BtcEncode(w io.Writer, pver uint32) error {
	return writeBlockHeader(w, pver, h)
}

// Deserialize decodes a block header from r into the receiver using a format
// that is suitable for long-term storage such as a database while respecting
// the Version field.
func (h *BlockHeader) Deserialize(r io.Reader) error {
	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of readBlockHeader.
	return readBlockHeader(r, 0, h)
}

// Serialize encodes a block header from r into the receiver using a format
// that is suitable for long-term storage such as a database while respecting
// the Version field.
func (h *BlockHeader) Serialize(w io.Writer) error {
	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of writeBlockHeader.
	return writeBlockHeader(w, 0, h)
}

// NewBlockHeader returns a new BlockHeader using the provided previous block
// hash, merkle root hash, difficulty bits, and nonce used to generate the
// block with defaults for the remaining fields.
func NewBlockHeader(prevHash *chainhash.Hash, merkleRootHash *chainhash.Hash,
	bits uint32, nonce uint32) *BlockHeader {

	// Limit the timestamp to one second precision since the protocol
	// doesn't support better.
	return &BlockHeader{
		Version:    BlockVersion,
		PrevBlock:  *prevHash,
		MerkleRoot: *merkleRootHash,
		Timestamp:  time.Unix(time.Now().Unix(), 0),
		Bits:       bits,
		Nonce:      nonce,
	}
}

// readBlockHeader reads a bitcoin block header from r.  See Deserialize for
// decoding block headers stored to disk, such as in a database, as opposed to
// decoding from the wire.
func readBlockHeader(r io.Reader, pver uint32, bh *BlockHeader) error {
	return readElements(r, &bh.Version, &bh.PrevBlock, &bh.MerkleRoot,
		(*uint32Time)(&bh.Timestamp), &bh.Bits, &bh.Nonce)
}

// writeBlockHeader writes a bitcoin block header to w.  See Serialize for
// encoding block headers to be stored to disk, such as in a database, as
// opposed to encoding for the wire.
func writeBlockHeader(w io.Writer, pver uint32, bh *BlockHeader) error {
	sec := uint32(bh.Timestamp.Unix())
	return writeElements(w, bh.Version, &bh.PrevBlock, &bh.MerkleRoot,
		sec, bh.Bits, bh.Nonce)
}
<|endoftext|>


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

package logs

import (
	"time"

	"github.com/golang/protobuf/proto"
	ds "github.com/luci/gae/service/datastore"
	"github.com/luci/luci-go/appengine/logdog/coordinator"
	"github.com/luci/luci-go/common/api/logdog_coordinator/logs/v1"
	"github.com/luci/luci-go/common/config"
	"github.com/luci/luci-go/common/grpcutil"
	"github.com/luci/luci-go/common/logdog/types"
	log "github.com/luci/luci-go/common/logging"
	"github.com/luci/luci-go/common/proto/logdog/logpb"
	"github.com/luci/luci-go/common/retry"
	"github.com/luci/luci-go/server/logdog/storage"
	"github.com/luci/luci-go/server/logdog/storage/archive"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
)

const (
	// getInitialArraySize is the initial amount of log slots to allocate for a
	// Get request.
	getInitialArraySize = 256

	// getBytesLimit is the maximum amount of data that we are willing to query.
	//
	// We will limit byte responses to 16MB, based on the following constraints:
	//	- AppEngine cannot respond with more than 32MB of data. This includes JSON
	//	  overhead, including notation and base64 data expansion.
	//	- `urlfetch`, which is used for Google Cloud Storage (archival) responses,
	//	  cannot handle responses larger than 32MB.
	getBytesLimit = 16 * 1024 * 1024
)

// Get returns state and log data for a single log stream.
func (s *server) Get(c context.Context, req *logdog.GetRequest) (*logdog.GetResponse, error) {
	return s.getImpl(c, req, false)
}

// Tail returns the last log entry for a given log stream.
func (s *server) Tail(c context.Context, req *logdog.TailRequest) (*logdog.GetResponse, error) {
	r := logdog.GetRequest{
		Project: req.Project,
		Path:    req.Path,
		State:   req.State,
	}
	return s.getImpl(c, &r, true)
}

// getImpl is common code shared between Get and Tail endpoints.
func (s *server) getImpl(c context.Context, req *logdog.GetRequest, tail bool) (*logdog.GetResponse, error) {
	log.Fields{
		"project": req.Project,
		"path":    req.Path,
		"index":   req.Index,
		"tail":    tail,
	}.Debugf(c, "Received get request.")

	path := types.StreamPath(req.Path)
	if err := path.Validate(); err != nil {
		log.WithError(err).Errorf(c, "Invalid path supplied.")
		return nil, grpcutil.Errf(codes.InvalidArgument, "invalid path value")
	}

	di := ds.Get(c)
	ls := &coordinator.LogStream{ID: coordinator.LogStreamID(path)}
	lst := ls.State(di)
	log.Fields{
		"id": ls.ID,
	}.Debugf(c, "Loading stream.")

	if err := di.GetMulti([]interface{}{ls, lst}); err != nil {
		if isNoSuchEntity(err) {
			log.Errorf(c, "Log stream does not exist.")
			return nil, grpcutil.Errf(codes.NotFound, "path not found")
		}

		log.WithError(err).Errorf(c, "Failed to look up log stream.")
		return nil, grpcutil.Internal
	}

	// If this log entry is Purged and we're not admin, pretend it doesn't exist.
	if ls.Purged {
		if authErr := coordinator.IsAdminUser(c); authErr != nil {
			log.Fields{
				log.ErrorKey: authErr,
			}.Warningf(c, "Non-superuser requested purged log.")
			return nil, grpcutil.Errf(codes.NotFound, "path not found")
		}
	}

	// If nothing was requested, return nothing.
	resp := logdog.GetResponse{}
	if !(req.State || tail) && req.LogCount < 0 {
		return &resp, nil
	}

	if req.State {
		resp.State = buildLogStreamState(ls, lst)

		var err error
		resp.Desc, err = ls.DescriptorValue()
		if err != nil {
			log.WithError(err).Errorf(c, "Failed to deserialize descriptor protobuf.")
			return nil, grpcutil.Internal
		}
	}

	// Retrieve requested logs from storage, if requested.
	if tail || req.LogCount >= 0 {
		var err error
		resp.Logs, err = s.getLogs(c, req, tail, ls, lst)
		if err != nil {
			log.WithError(err).Errorf(c, "Failed to get logs.")
			return nil, grpcutil.Internal
		}
	}

	log.Fields{
		"logCount": len(resp.Logs),
	}.Debugf(c, "Get request completed successfully.")
	return &resp, nil
}

func (s *server) getLogs(c context.Context, req *logdog.GetRequest, tail bool, ls *coordinator.LogStream,
	lst *coordinator.LogStreamState) ([]*logpb.LogEntry, error) {
	byteLimit := int(req.ByteCount)
	if byteLimit <= 0 || byteLimit > getBytesLimit {
		byteLimit = getBytesLimit
	}

	svc := coordinator.GetServices(c)
	var st storage.Storage
	if lst.ArchivalState() == coordinator.NotArchived {
		log.Debugf(c, "Log is not archived. Fetching from intermediate storage.")

		// Logs are not archived. Fetch from intermediate storage.
		var err error
		st, err = svc.IntermediateStorage(c)
		if err != nil {
			return nil, err
		}
	} else {
		log.Fields{
			"indexURL":    lst.ArchiveIndexURL,
			"streamURL":   lst.ArchiveStreamURL,
			"archiveTime": lst.ArchivedTime,
		}.Debugf(c, "Log is archived. Fetching from archive storage.")

		var err error
		gs, err := svc.GSClient(c)
		if err != nil {
			log.WithError(err).Errorf(c, "Failed to create Google Storage client.")
			return nil, err
		}
		defer func() {
			if err := gs.Close(); err != nil {
				log.WithError(err).Warningf(c, "Failed to close Google Storage client.")
			}
		}()

		st, err = archive.New(c, archive.Options{
			IndexURL:  lst.ArchiveIndexURL,
			StreamURL: lst.ArchiveStreamURL,
			Client:    gs,
			MaxBytes:  byteLimit,
		})
		if err != nil {
			log.WithError(err).Errorf(c, "Failed to create Google Storage storage instance.")
			return nil, err
		}
	}
	defer st.Close()

	project, path := coordinator.Project(c), ls.Path()

	var fetchedLogs [][]byte
	var err error
	if tail {
		fetchedLogs, err = getTail(c, st, project, path)
	} else {
		fetchedLogs, err = getHead(c, req, st, project, path, byteLimit)
	}
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to fetch log records.")
		return nil, err
	}

	logEntries := make([]*logpb.LogEntry, len(fetchedLogs))
	for idx, ld := range fetchedLogs {
		// Deserialize the log entry, then convert it to output value.
		le := logpb.LogEntry{}
		if err := proto.Unmarshal(ld, &le); err != nil {
			log.Fields{
				log.ErrorKey: err,
				"index":      idx,
			}.Errorf(c, "Failed to generate response log entry.")
			return nil, err
		}
		logEntries[idx] = &le
	}
	return logEntries, nil
}

func getHead(c context.Context, req *logdog.GetRequest, st storage.Storage, project config.ProjectName,
	path types.StreamPath, byteLimit int) ([][]byte, error) {
	log.Fields{
		"project":       project,
		"path":          path,
		"index":         req.Index,
		"count":         req.LogCount,
		"bytes":         req.ByteCount,
		"noncontiguous": req.NonContiguous,
	}.Debugf(c, "Issuing Get request.")

	// Allocate result logs array.
	logCount := int(req.LogCount)
	asz := getInitialArraySize
	if logCount > 0 && logCount < asz {
		asz = logCount
	}
	logs := make([][]byte, 0, asz)

	sreq := storage.GetRequest{
		Project: project,
		Path:    path,
		Index:   types.MessageIndex(req.Index),
		Limit:   logCount,
	}

	count := 0
	err := retry.Retry(c, retry.TransientOnly(retry.Default), func() error {
		// Issue the Get request. This may return a transient error, in which case
		// we will retry.
		return st.Get(sreq, func(idx types.MessageIndex, ld []byte) bool {
			if count > 0 && byteLimit-len(ld) < 0 {
				// Not the first log, and we've exceeded our byte limit.
				return false
			}
			byteLimit -= len(ld)

			if !(req.NonContiguous || idx == sreq.Index) {
				return false
			}
			logs = append(logs, ld)
			sreq.Index = idx + 1
			count++
			return !(logCount > 0 && count >= logCount)
		})
	}, func(err error, delay time.Duration) {
		log.Fields{
			log.ErrorKey:   err,
			"delay":        delay,
			"initialIndex": req.Index,
			"nextIndex":    sreq.Index,
			"count":        len(logs),
		}.Warningf(c, "Transient error while loading logs; retrying.")
	})
	switch err {
	case nil:
		return logs, nil

	case storage.ErrDoesNotExist:
		return nil, nil

	default:
		log.Fields{
			log.ErrorKey:   err,
			"initialIndex": req.Index,
			"nextIndex":    sreq.Index,
			"count":        len(logs),
		}.Errorf(c, "Failed to execute range request.")
		return nil, err
	}
}

func getTail(c context.Context, st storage.Storage, project config.ProjectName, path types.StreamPath) ([][]byte, error) {
	log.Fields{
		"project": project,
		"path":    path,
	}.Debugf(c, "Issuing Tail request.")

	var data []byte
	err := retry.Retry(c, retry.TransientOnly(retry.Default), func() (err error) {
		data, _, err = st.Tail(project, path)
		return
	}, func(err error, delay time.Duration) {
		log.Fields{
			log.ErrorKey: err,
			"delay":      delay,
		}.Warningf(c, "Transient error while fetching tail log; retrying.")
	})
	switch err {
	case nil:
		return [][]byte{data}, nil

	case storage.ErrDoesNotExist:
		return nil, nil

	default:
		log.WithError(err).Errorf(c, "Failed to fetch tail log.")
		return nil, err
	}
}
<commit_msg>LogDog: Use proper test for archival state.<commit_after>// Copyright 2015 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package logs

import (
	"time"

	"github.com/golang/protobuf/proto"
	ds "github.com/luci/gae/service/datastore"
	"github.com/luci/luci-go/appengine/logdog/coordinator"
	"github.com/luci/luci-go/common/api/logdog_coordinator/logs/v1"
	"github.com/luci/luci-go/common/config"
	"github.com/luci/luci-go/common/grpcutil"
	"github.com/luci/luci-go/common/logdog/types"
	log "github.com/luci/luci-go/common/logging"
	"github.com/luci/luci-go/common/proto/logdog/logpb"
	"github.com/luci/luci-go/common/retry"
	"github.com/luci/luci-go/server/logdog/storage"
	"github.com/luci/luci-go/server/logdog/storage/archive"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
)

const (
	// getInitialArraySize is the initial amount of log slots to allocate for a
	// Get request.
	getInitialArraySize = 256

	// getBytesLimit is the maximum amount of data that we are willing to query.
	//
	// We will limit byte responses to 16MB, based on the following constraints:
	//	- AppEngine cannot respond with more than 32MB of data. This includes JSON
	//	  overhead, including notation and base64 data expansion.
	//	- `urlfetch`, which is used for Google Cloud Storage (archival) responses,
	//	  cannot handle responses larger than 32MB.
	getBytesLimit = 16 * 1024 * 1024
)

// Get returns state and log data for a single log stream.
func (s *server) Get(c context.Context, req *logdog.GetRequest) (*logdog.GetResponse, error) {
	return s.getImpl(c, req, false)
}

// Tail returns the last log entry for a given log stream.
func (s *server) Tail(c context.Context, req *logdog.TailRequest) (*logdog.GetResponse, error) {
	r := logdog.GetRequest{
		Project: req.Project,
		Path:    req.Path,
		State:   req.State,
	}
	return s.getImpl(c, &r, true)
}

// getImpl is common code shared between Get and Tail endpoints.
func (s *server) getImpl(c context.Context, req *logdog.GetRequest, tail bool) (*logdog.GetResponse, error) {
	log.Fields{
		"project": req.Project,
		"path":    req.Path,
		"index":   req.Index,
		"tail":    tail,
	}.Debugf(c, "Received get request.")

	path := types.StreamPath(req.Path)
	if err := path.Validate(); err != nil {
		log.WithError(err).Errorf(c, "Invalid path supplied.")
		return nil, grpcutil.Errf(codes.InvalidArgument, "invalid path value")
	}

	di := ds.Get(c)
	ls := &coordinator.LogStream{ID: coordinator.LogStreamID(path)}
	lst := ls.State(di)
	log.Fields{
		"id": ls.ID,
	}.Debugf(c, "Loading stream.")

	if err := di.GetMulti([]interface{}{ls, lst}); err != nil {
		if isNoSuchEntity(err) {
			log.Errorf(c, "Log stream does not exist.")
			return nil, grpcutil.Errf(codes.NotFound, "path not found")
		}

		log.WithError(err).Errorf(c, "Failed to look up log stream.")
		return nil, grpcutil.Internal
	}

	// If this log entry is Purged and we're not admin, pretend it doesn't exist.
	if ls.Purged {
		if authErr := coordinator.IsAdminUser(c); authErr != nil {
			log.Fields{
				log.ErrorKey: authErr,
			}.Warningf(c, "Non-superuser requested purged log.")
			return nil, grpcutil.Errf(codes.NotFound, "path not found")
		}
	}

	// If nothing was requested, return nothing.
	resp := logdog.GetResponse{}
	if !(req.State || tail) && req.LogCount < 0 {
		return &resp, nil
	}

	if req.State {
		resp.State = buildLogStreamState(ls, lst)

		var err error
		resp.Desc, err = ls.DescriptorValue()
		if err != nil {
			log.WithError(err).Errorf(c, "Failed to deserialize descriptor protobuf.")
			return nil, grpcutil.Internal
		}
	}

	// Retrieve requested logs from storage, if requested.
	if tail || req.LogCount >= 0 {
		var err error
		resp.Logs, err = s.getLogs(c, req, tail, ls, lst)
		if err != nil {
			log.WithError(err).Errorf(c, "Failed to get logs.")
			return nil, grpcutil.Internal
		}
	}

	log.Fields{
		"logCount": len(resp.Logs),
	}.Debugf(c, "Get request completed successfully.")
	return &resp, nil
}

func (s *server) getLogs(c context.Context, req *logdog.GetRequest, tail bool, ls *coordinator.LogStream,
	lst *coordinator.LogStreamState) ([]*logpb.LogEntry, error) {
	byteLimit := int(req.ByteCount)
	if byteLimit <= 0 || byteLimit > getBytesLimit {
		byteLimit = getBytesLimit
	}

	svc := coordinator.GetServices(c)
	var st storage.Storage
	if !lst.ArchivalState().Archived() {
		log.Debugf(c, "Log is not archived. Fetching from intermediate storage.")

		// Logs are not archived. Fetch from intermediate storage.
		var err error
		st, err = svc.IntermediateStorage(c)
		if err != nil {
			return nil, err
		}
	} else {
		log.Fields{
			"indexURL":    lst.ArchiveIndexURL,
			"streamURL":   lst.ArchiveStreamURL,
			"archiveTime": lst.ArchivedTime,
		}.Debugf(c, "Log is archived. Fetching from archive storage.")

		var err error
		gs, err := svc.GSClient(c)
		if err != nil {
			log.WithError(err).Errorf(c, "Failed to create Google Storage client.")
			return nil, err
		}
		defer func() {
			if err := gs.Close(); err != nil {
				log.WithError(err).Warningf(c, "Failed to close Google Storage client.")
			}
		}()

		st, err = archive.New(c, archive.Options{
			IndexURL:  lst.ArchiveIndexURL,
			StreamURL: lst.ArchiveStreamURL,
			Client:    gs,
			MaxBytes:  byteLimit,
		})
		if err != nil {
			log.WithError(err).Errorf(c, "Failed to create Google Storage storage instance.")
			return nil, err
		}
	}
	defer st.Close()

	project, path := coordinator.Project(c), ls.Path()

	var fetchedLogs [][]byte
	var err error
	if tail {
		fetchedLogs, err = getTail(c, st, project, path)
	} else {
		fetchedLogs, err = getHead(c, req, st, project, path, byteLimit)
	}
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to fetch log records.")
		return nil, err
	}

	logEntries := make([]*logpb.LogEntry, len(fetchedLogs))
	for idx, ld := range fetchedLogs {
		// Deserialize the log entry, then convert it to output value.
		le := logpb.LogEntry{}
		if err := proto.Unmarshal(ld, &le); err != nil {
			log.Fields{
				log.ErrorKey: err,
				"index":      idx,
			}.Errorf(c, "Failed to generate response log entry.")
			return nil, err
		}
		logEntries[idx] = &le
	}
	return logEntries, nil
}

func getHead(c context.Context, req *logdog.GetRequest, st storage.Storage, project config.ProjectName,
	path types.StreamPath, byteLimit int) ([][]byte, error) {
	log.Fields{
		"project":       project,
		"path":          path,
		"index":         req.Index,
		"count":         req.LogCount,
		"bytes":         req.ByteCount,
		"noncontiguous": req.NonContiguous,
	}.Debugf(c, "Issuing Get request.")

	// Allocate result logs array.
	logCount := int(req.LogCount)
	asz := getInitialArraySize
	if logCount > 0 && logCount < asz {
		asz = logCount
	}
	logs := make([][]byte, 0, asz)

	sreq := storage.GetRequest{
		Project: project,
		Path:    path,
		Index:   types.MessageIndex(req.Index),
		Limit:   logCount,
	}

	count := 0
	err := retry.Retry(c, retry.TransientOnly(retry.Default), func() error {
		// Issue the Get request. This may return a transient error, in which case
		// we will retry.
		return st.Get(sreq, func(idx types.MessageIndex, ld []byte) bool {
			if count > 0 && byteLimit-len(ld) < 0 {
				// Not the first log, and we've exceeded our byte limit.
				return false
			}
			byteLimit -= len(ld)

			if !(req.NonContiguous || idx == sreq.Index) {
				return false
			}
			logs = append(logs, ld)
			sreq.Index = idx + 1
			count++
			return !(logCount > 0 && count >= logCount)
		})
	}, func(err error, delay time.Duration) {
		log.Fields{
			log.ErrorKey:   err,
			"delay":        delay,
			"initialIndex": req.Index,
			"nextIndex":    sreq.Index,
			"count":        len(logs),
		}.Warningf(c, "Transient error while loading logs; retrying.")
	})
	switch err {
	case nil:
		return logs, nil

	case storage.ErrDoesNotExist:
		return nil, nil

	default:
		log.Fields{
			log.ErrorKey:   err,
			"initialIndex": req.Index,
			"nextIndex":    sreq.Index,
			"count":        len(logs),
		}.Errorf(c, "Failed to execute range request.")
		return nil, err
	}
}

func getTail(c context.Context, st storage.Storage, project config.ProjectName, path types.StreamPath) ([][]byte, error) {
	log.Fields{
		"project": project,
		"path":    path,
	}.Debugf(c, "Issuing Tail request.")

	var data []byte
	err := retry.Retry(c, retry.TransientOnly(retry.Default), func() (err error) {
		data, _, err = st.Tail(project, path)
		return
	}, func(err error, delay time.Duration) {
		log.Fields{
			log.ErrorKey: err,
			"delay":      delay,
		}.Warningf(c, "Transient error while fetching tail log; retrying.")
	})
	switch err {
	case nil:
		return [][]byte{data}, nil

	case storage.ErrDoesNotExist:
		return nil, nil

	default:
		log.WithError(err).Errorf(c, "Failed to fetch tail log.")
		return nil, err
	}
}
<|endoftext|>


<commit_before>package test

import (
	"math/rand"
	"strconv"
	"testing"
	"time"

	"fmt"

	"github.com/heqzha/dcache/core"
)

func TestSrvTable(t *testing.T) {
	tb1 := core.SrvTable{
		"127.0.0.1:1001": 0,
		"127.0.0.1:1002": 0,
		"127.0.0.1:1003": 0,
	}
	tb2 := core.SrvTable{
		"127.0.0.1:1004": 0,
		"127.0.0.1:1005": 0,
	}
	tb1.Add(tb2)
	fmt.Println(tb1.String())

	tb2 = tb1.Clone()
	fmt.Println(tb2.String())
}

func TestSrvGroup(t *testing.T) {
	group := core.SrvGroup{}
	group.NewGroup("test1")
	tb1 := core.SrvTable{
		"127.0.0.1:1001": 0,
		"127.0.0.1:1002": 0,
		"127.0.0.1:1003": 0,
	}
	tb2 := core.SrvTable{
		"127.0.0.1:1004": 0,
		"127.0.0.1:1005": 0,
	}
	group.SetTable("test1", &tb1)

	group.SetTable("test2", &tb2)

	fmt.Println(group.GetTable("test1"))

	fmt.Println(group.GetTable("test2"))
}

func TestSGM(t *testing.T) {
	var (
		sgm  core.SGM
		sgm2 core.SGM
	)
	sgm.Init()
	sgm.RegisterLocalAddr("group1", "127.0.0.1:1001")
	sgm2.Init()
	sgm2.RegisterLocalAddr("group2", "127.0.0.2:1001")

	for index := 1001; index < 1004; index++ {
		sgm.Register("group1", "127.0.0.1:"+strconv.Itoa(index))
	}
	for index := 1002; index < 1006; index++ {
		sgm.Register("group2", "127.0.0.1:"+strconv.Itoa(index))
	}
	fmt.Println(sgm.GetGroupNames())
	tb1g1, _ := sgm.GetTable("group1")
	tb1g2, _ := sgm.GetTable("group2")
	clk1 := sgm.GetClock()
	fmt.Println(tb1g1.String())
	fmt.Println(tb1g2.String())
	fmt.Println(clk1.ReturnVCString())

	time.Sleep(2 * time.Second)

	sgm.UpdateTimestamp("group1", "127.0.0.1:1001")

	for index := 1001; index < 1003; index++ {
		sgm2.Register("group1", "127.0.0.2:"+strconv.Itoa(index))
	}
	tb2gp1, _ := sgm2.GetTable("group1")
	clk2 := sgm2.GetClock()
	fmt.Println(tb2gp1.String())
	fmt.Println(clk2.ReturnVCString())

	sgm2.Merge(sgm)
	fmt.Println("Condition:", sgm2.CompareReadable(sgm))
	tb2mgp1, _ := sgm2.GetTable("group1")
	tb2mgp2, _ := sgm2.GetTable("group2")
	clk2 = sgm2.GetClock()
	fmt.Println(tb2mgp1.String())
	fmt.Println(tb2mgp2.String())
	fmt.Println(clk2.ReturnVCString())

	sgm.Merge(sgm2)
	fmt.Println("Condition:", sgm.CompareReadable(sgm2))
	tb1mgp1, _ := sgm.GetTable("group1")
	tb1mgp2, _ := sgm.GetTable("group2")
	clk1 = sgm.GetClock()
	fmt.Println(tb1mgp1.String())
	fmt.Println(tb1mgp2.String())
	fmt.Println(clk1.ReturnVCString())
}

func TestSTMDumpLoad(t *testing.T) {
	var (
		sgm  core.SGM
		sgm2 core.SGM
	)
	sgm.Init()
	sgm.RegisterLocalAddr("group1", "127.0.0.1:1001")
	sgm2.Init()
	sgm2.RegisterLocalAddr("group2", "127.0.0.2:1001")

	for index := 1001; index < 1004; index++ {
		sgm.Register("group1", "127.0.0.1:"+strconv.Itoa(index))
	}

	for index := 1003; index < 1007; index++ {
		sgm.Register("group2", "127.0.0.1:"+strconv.Itoa(index))
	}

	d, err := sgm.Dump()
	if err != nil {
		t.Error(err)
		return
	}

	if err := sgm2.Load(d); err != nil {
		t.Error(err)
		return
	}
	tableg1, _ := sgm2.GetTable("group1")
	tableg2, _ := sgm2.GetTable("group2")
	fmt.Println(tableg1.String())
	fmt.Println(tableg2.String())
	fmt.Println(sgm.CompareReadable(sgm2))
}

const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

func RandStringBytes(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = letterBytes[rand.Intn(len(letterBytes))]
	}
	return string(b)
}

func TestSTHash(t *testing.T) {
	var (
		sth core.SGHash
		stm core.SGM
	)
	stm.Init()
	stm.RegisterLocalAddr("group1", "127.0.0.1:1001")

	for index := 1001; index < 1010; index++ {
		stm.Register("group1", "127.0.0.1:"+strconv.Itoa(index))
	}

	sth.Init(3, nil)
	sth.Load(stm.GetGroup())

	for index := 0; index < 20; index++ {
		key := RandStringBytes(10)
		srv := sth.Pick("group1", key)
		fmt.Println(key, srv)
	}
}

func TestUnregister(t *testing.T) {
	var (
		sgm core.SGM
	)
	sgm.Init()
	sgm.RegisterLocalAddr("group1", "127.0.0.1:1000")

	sgm.Register("group1", "127.0.0.1:1001")
	tb, _ := sgm.GetTable("group1")
	t.Log(tb.String())
	sgm.Unregister("group1", "127.0.0.1:1001")
	tb, _ = sgm.GetTable("group1")
	t.Log(tb.String())
}
<commit_msg>Add sgh load pick test codes<commit_after>package test

import (
	"math/rand"
	"strconv"
	"testing"
	"time"

	"fmt"

	"github.com/heqzha/dcache/core"
)

func TestSrvTable(t *testing.T) {
	tb1 := core.SrvTable{
		"127.0.0.1:1001": 0,
		"127.0.0.1:1002": 0,
		"127.0.0.1:1003": 0,
	}
	tb2 := core.SrvTable{
		"127.0.0.1:1004": 0,
		"127.0.0.1:1005": 0,
	}
	tb1.Add(tb2)
	fmt.Println(tb1.String())

	tb2 = tb1.Clone()
	fmt.Println(tb2.String())
}

func TestSrvGroup(t *testing.T) {
	group := core.SrvGroup{}
	group.NewGroup("test1")
	tb1 := core.SrvTable{
		"127.0.0.1:1001": 0,
		"127.0.0.1:1002": 0,
		"127.0.0.1:1003": 0,
	}
	tb2 := core.SrvTable{
		"127.0.0.1:1004": 0,
		"127.0.0.1:1005": 0,
	}
	group.SetTable("test1", &tb1)

	group.SetTable("test2", &tb2)

	fmt.Println(group.GetTable("test1"))

	fmt.Println(group.GetTable("test2"))
}

func TestSGM(t *testing.T) {
	var (
		sgm  core.SGM
		sgm2 core.SGM
	)
	sgm.Init()
	sgm.RegisterLocalAddr("group1", "127.0.0.1:1001")
	sgm2.Init()
	sgm2.RegisterLocalAddr("group2", "127.0.0.2:1001")

	for index := 1001; index < 1004; index++ {
		sgm.Register("group1", "127.0.0.1:"+strconv.Itoa(index))
	}
	for index := 1002; index < 1006; index++ {
		sgm.Register("group2", "127.0.0.1:"+strconv.Itoa(index))
	}
	fmt.Println(sgm.GetGroupNames())
	tb1g1, _ := sgm.GetTable("group1")
	tb1g2, _ := sgm.GetTable("group2")
	clk1 := sgm.GetClock()
	fmt.Println(tb1g1.String())
	fmt.Println(tb1g2.String())
	fmt.Println(clk1.ReturnVCString())

	time.Sleep(2 * time.Second)

	sgm.UpdateTimestamp("group1", "127.0.0.1:1001")

	for index := 1001; index < 1003; index++ {
		sgm2.Register("group1", "127.0.0.2:"+strconv.Itoa(index))
	}
	tb2gp1, _ := sgm2.GetTable("group1")
	clk2 := sgm2.GetClock()
	fmt.Println(tb2gp1.String())
	fmt.Println(clk2.ReturnVCString())

	sgm2.Merge(sgm)
	fmt.Println("Condition:", sgm2.CompareReadable(sgm))
	tb2mgp1, _ := sgm2.GetTable("group1")
	tb2mgp2, _ := sgm2.GetTable("group2")
	clk2 = sgm2.GetClock()
	fmt.Println(tb2mgp1.String())
	fmt.Println(tb2mgp2.String())
	fmt.Println(clk2.ReturnVCString())

	sgm.Merge(sgm2)
	fmt.Println("Condition:", sgm.CompareReadable(sgm2))
	tb1mgp1, _ := sgm.GetTable("group1")
	tb1mgp2, _ := sgm.GetTable("group2")
	clk1 = sgm.GetClock()
	fmt.Println(tb1mgp1.String())
	fmt.Println(tb1mgp2.String())
	fmt.Println(clk1.ReturnVCString())
}

func TestSTMDumpLoad(t *testing.T) {
	var (
		sgm  core.SGM
		sgm2 core.SGM
	)
	sgm.Init()
	sgm.RegisterLocalAddr("group1", "127.0.0.1:1001")
	sgm2.Init()
	sgm2.RegisterLocalAddr("group2", "127.0.0.2:1001")

	for index := 1001; index < 1004; index++ {
		sgm.Register("group1", "127.0.0.1:"+strconv.Itoa(index))
	}

	for index := 1003; index < 1007; index++ {
		sgm.Register("group2", "127.0.0.1:"+strconv.Itoa(index))
	}

	d, err := sgm.Dump()
	if err != nil {
		t.Error(err)
		return
	}

	if err := sgm2.Load(d); err != nil {
		t.Error(err)
		return
	}
	tableg1, _ := sgm2.GetTable("group1")
	tableg2, _ := sgm2.GetTable("group2")
	fmt.Println(tableg1.String())
	fmt.Println(tableg2.String())
	fmt.Println(sgm.CompareReadable(sgm2))
}

const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

func RandStringBytes(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = letterBytes[rand.Intn(len(letterBytes))]
	}
	return string(b)
}

func TestSTHash(t *testing.T) {
	var (
		sth core.SGHash
		stm core.SGM
	)
	stm.Init()
	stm.RegisterLocalAddr("group1", "127.0.0.1:1001")

	for index := 1001; index < 1010; index++ {
		stm.Register("group1", "127.0.0.1:"+strconv.Itoa(index))
	}

	sth.Init(3, nil)
	sth.Load(stm.GetGroup())

	for index := 0; index < 20; index++ {
		key := RandStringBytes(10)
		srv := sth.Pick("group1", key)
		fmt.Println(key, srv)
	}
}

func TestSTHashLoadPick(t *testing.T) {
	var (
		stm1 core.SGM
		stm2 core.SGM
		stm3 core.SGM

		sth1 core.SGHash
		sth2 core.SGHash
		sth3 core.SGHash
	)
	stm1.Init()
	stm1.RegisterLocalAddr("group1", "127.0.0.1:1001")
	stm2.Init()
	stm2.RegisterLocalAddr("group1", "127.0.0.1:1002")
	stm1.Register(stm2.GetLocalGroup(), stm2.GetLocalAddr())
	stm2.Merge(stm1)
	stm1.Merge(stm2)

	stm3.Init()
	stm3.RegisterLocalAddr("group1", "127.0.0.1:1003")
	stm2.Register(stm3.GetLocalGroup(), stm3.GetLocalAddr())
	stm3.Merge(stm2)
	stm2.Merge(stm3)
	stm1.Merge(stm2)
	stm2.Merge(stm1)

	tb1, _ := stm1.GetTable("group1")
	tb2, _ := stm2.GetTable("group1")
	tb3, _ := stm3.GetTable("group1")
	fmt.Println(tb1.String())
	fmt.Println(tb2.String())
	fmt.Println(tb3.String())

	sth1.Init(3, nil)
	sth1.Load(stm1.GetGroup())

	sth2.Init(3, nil)
	sth2.Load(stm2.GetGroup())

	sth3.Init(3, nil)
	sth3.Load(stm3.GetGroup())

	keys := []string{}
	for index := 0; index < 10; index++ {
		keys = append(keys, RandStringBytes(2))
	}

	for _, k := range keys {
		addr1 := sth1.Pick("group1", k)
		addr2 := sth2.Pick("group1", k)
		addr3 := sth3.Pick("group1", k)
		fmt.Println(k, addr1, addr2, addr3)
	}
}

func TestUnregister(t *testing.T) {
	var (
		sgm core.SGM
	)
	sgm.Init()
	sgm.RegisterLocalAddr("group1", "127.0.0.1:1000")

	sgm.Register("group1", "127.0.0.1:1001")
	tb, _ := sgm.GetTable("group1")
	t.Log(tb.String())
	sgm.Unregister("group1", "127.0.0.1:1001")
	tb, _ = sgm.GetTable("group1")
	t.Log(tb.String())
}
<|endoftext|>


<commit_before>package data

import (
	"fmt"
	"math"

	bits "github.com/willf/bitset"

	cmn "github.com/RxnWeaver/rxnweaver/common"
)

// Atom represents a chemical atom.
//
// It has various physical and chemical properties, apart from the
// transient information attached to it during its participation in
// reactions.
type Atom struct {
	mol    *Molecule // Containing molecule of this atom.
	atNum  uint8     // Atomic number of this atom's element.
	symbol string    // Symbol, in case of a different isotope.
	iId    uint16    // Serial input ID of this atom.
	nId    uint16    // Normalised ID of this atom.

	X float32 // X-coordinate of this atom.
	Y float32 // Y-coordinate of this atom.
	Z float32 // Z-coordinate of this atom.

	numH    uint8 // Number of implicit + explicit H atoms attached to this atom.
	charge  int8  // Residual net charge of this atom.
	valence int8  // Current valence configuration of this atom.

	pHash uint64 // A pseudo-hash of this atom, using some attributes.
	sHash uint64 // A pseudo-hash of this atom, using some attributes.

	bonds          *bits.BitSet // Bitmap of bonds of this atom.
	nbrs           []uint16     // Expanded list of neighbours of this atom.
	numSingleBonds uint8        // Number of single bonds this atom has.
	numDoubleBonds uint8        // Number of double bonds this atom has.
	numTripleBonds uint8        // Number of triple bonds this atom has.

	rings *bits.BitSet // Bitmap of IDs of rings this atom participates in.
	// Does this atom participate in at least one aromatic ring?
	isInAroRing bool
	// Is this atom a bridgehead of a bicyclic system of rings?
	isBridgeHead bool
	// Is this atom the sole common atom of all of its rings?
	isSpiro bool

	// The functional groups substituted on this atom.  They are listed in
	// descending order of importance.  The first is the primary feature.
	features [cmn.MaxFeatures]uint8
}

// newAtom constructs and initialises a new atom of the given element
// type, and belonging to the given molecule.
func newAtom(mol *Molecule, atNum uint8) *Atom {
	atom := new(Atom)
	atom.mol = mol
	atom.atNum = atNum
	atom.valence = cmn.PeriodicTable[cmn.ElemSyms[atNum]].Valence

	atom.bonds = bits.New(cmn.MaxBonds)
	atom.nbrs = make([]uint16, 0, cmn.MaxBonds)
	atom.rings = bits.New(cmn.MaxRings)

	return atom
}

// AtomicNumber answers the atomic number of this atom.
func (a *Atom) AtomicNumber() uint8 {
	return a.atNum
}

// Parent answers the parent molecule of this atom.
func (a *Atom) Parent() *Molecule {
	return a.mol
}

// numPiElectrons answers the number of delocalised pi electrons
// contributed by this atom.  This number is important for calculating
// the aromaticity of the rings this atom participates in.
func (a *Atom) numPiElectrons() int {
	mol := a.mol
	wtSum := 100*int16(a.numDoubleBonds) + 10*int16(a.numSingleBonds) + int16(a.charge)

	switch a.atNum {
	case 6:
		switch wtSum {
		case 19:
			return 2
		case 110:
			return 1
		case 120:
			var b *Bond
			for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
				b = mol.bondWithId(uint16(bid - 1))
				if b.bType == cmn.BondTypeDouble {
					break
				}
			}
			if b.isCyclic() {
				return 1
			}
			return 0
		default:
			return 0
		}

	case 7:
		switch wtSum {
		case 20:
			return 2
		case 30:
			return 2
		case 110:
			return 1
		case 121:
			return 1
		default:
			return 0
		}

	case 8:
		switch wtSum {
		case 20:
			return 2
		case 111:
			return 1
		default:
			return 0
		}

	case 16:
		switch wtSum {
		case 20:
			return 2
		case 111:
			return 1
		case 120:
			var b *Bond
			for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
				b = mol.bondWithId(uint16(bid - 1))
				if b.bType == cmn.BondTypeDouble {
					break
				}
			}
			oaid := b.otherAtom(a.iId)
			oa := mol.atomWithIid(oaid)
			if oa.atNum == 8 && !oa.isCyclic() {
				return 2
			}
			return 0
		case 220:
			c := 0
			for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
				b := mol.bondWithId(uint16(bid - 1))
				if b.bType == cmn.BondTypeDouble {
					oaid := b.otherAtom(a.iId)
					if !mol.atomWithIid(oaid).isCyclic() {
						c++
					}
				}
			}
			if c > 1 {
				return -1
			}
			return 0
		default:
			return 0
		}
	}

	return 0
}

// isCyclic answers if this atom participates in at least one ring.
func (a *Atom) isCyclic() bool {
	return a.rings.Count() > 0
}

// isJunction answers if this atom has more than 2 distinct
// neighbours.
func (a *Atom) isJunction() bool {
	return a.bonds.Count() > 2
}

// addBond adds the given bond to this atom, if it is not already
// present.  It also adjusts the list of its neighbours appropriately.
//
// Note that it does NOT check to see if the addition conforms to this
// atom's current valence configuration.
func (a *Atom) addBond(b *Bond) {
	for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
		if uint16(bid) == b.id {
			return
		}
	}

	a.bonds.Set(uint(b.id))
	nbrId := b.otherAtom(a.iId)
	n := int(b.bType)
	for i := 0; i < n; i++ {
		a.nbrs = append(a.nbrs, nbrId)
	}
}

// removeBond removes the given bond from this atom, and adjusts the
// list of neighbours appropriately.
//
// Note that it does NOT check to see if the removal conforms to this
// atom's current valence configuration.
func (a *Atom) removeBond(b *Bond) {
	nbrId := b.otherAtom(a.iId)

	wid := 0
	for _, nid := range a.nbrs {
		if nid == nbrId {
			continue
		}
		a.nbrs[wid] = nid
		wid++
	}
	a.nbrs = a.nbrs[:wid]

	a.bonds.Clear(uint(b.id))
}

// bondTo answers the bond that binds this atom to the given atom, if
// one such bond exists.  Answers `nil` otherwise.
func (a *Atom) bondTo(other uint16) *Bond {
	mol := a.mol
	for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
		b := mol.bondWithId(uint16(bid - 1))
		if b.otherAtom(a.iId) == other {
			return b
		}
	}

	return nil
}

// firstDoublyBondedNbr answers this atom's doubly-bonded neighbour
// having the highest priority.
//
// This method assumes that the molecule is already normalised!
// Calling it on a molecule that has not be normalised yet, leads to
// incorrect results.
func (a *Atom) firstDoublyBondedNbr() uint16 {
	if a.numDoubleBonds == 0 {
		return 0
	}

	mol := a.mol
	for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
		b := mol.bondWithId(uint16(bid - 1))
		if b.bType == cmn.BondTypeDouble {
			return b.otherAtom(a.iId)
		}
	}

	panic("Should never be here!")
}

// firstMultiplyBondedNbr answers this atom's doubly-bonded neighbour
// having the highest priority.
//
// This method assumes that the molecule is already normalised!
// Calling it on a molecule that has not be normalised yet, leads to
// incorrect results.
func (a *Atom) firstMultiplyBondedNbr() uint16 {
	if a.numDoubleBonds == 0 && a.numTripleBonds == 0 {
		return 0
	}

	mol := a.mol
	for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
		b := mol.bondWithId(uint16(bid - 1))
		if b.bType >= cmn.BondTypeDouble {
			return b.otherAtom(a.iId)
		}
	}

	panic("Should never be here!")
}

// inRingOfSize answers if this atom participates in at least one ring
// of the given size.
func (a *Atom) inRingOfSize(n int) bool {
	mol := a.mol
	for rid, ok := a.rings.NextSet(0); ok; rid, ok = a.rings.NextSet(rid + 1) {
		r := mol.ringWithId(uint8(rid))
		if r.size() == n {
			return true
		}
	}

	return false
}

// inRingLargerThan answers if this atom participates in at least one
// ring that is larger than the given number.
func (a *Atom) inRingLargerThan(n int) bool {
	mol := a.mol
	for rid, ok := a.rings.NextSet(0); ok; rid, ok = a.rings.NextSet(rid + 1) {
		r := mol.ringWithId(uint8(rid))
		if r.size() > n {
			return true
		}
	}

	return false
}

// smallestRing answers the smallest unique ring in which this atom
// participates.  Should no such unique ring exist, an error is
// answered.
func (a *Atom) smallestRing() (uint8, error) {
	if !a.isCyclic() {
		return 0, fmt.Errorf("Atom not cyclic.")
	}

	min := int(math.MaxUint8)
	c := 0
	var ret uint8

	mol := a.mol
	for rid, ok := a.rings.NextSet(0); ok; rid, ok = a.rings.NextSet(rid + 1) {
		r := mol.ringWithId(uint8(rid))
		size := r.size()
		if size == min {
			c++
		} else if size < min {
			ret = uint8(rid)
			c = 1
		}
	}

	if c > 1 {
		return 0, fmt.Errorf("Smallest ring size: %d, number of smallest rings: %d.", min, c)
	}

	return ret, nil
}

// isAromatic answers if this atom is part of an aromatic ring.
//
// Note that the actual aromaticity determination is handled by
// `Ring`.  This method merely answers the set flag.
func (a *Atom) isAromatic() bool {
	return a.isInAroRing
}

// inHetAromaticRing answers if this atom is part of an aromatic ring
// with at least one hetero atom.
func (a *Atom) inHetAromaticRing() bool {
	if a.isInAroRing && a.atNum != 6 {
		return true // Simplest case!
	}

	mol := a.mol
	for rid, ok := a.rings.NextSet(0); ok; rid, ok = a.rings.NextSet(rid + 1) {
		r := mol.ringWithId(uint8(rid))
		if r.isHetAromatic() {
			return true
		}
	}

	return false
}

// haveCommonRings answers if this atom has at least one ring in
// common with the given atom.
func (a *Atom) haveCommonRings(aiid uint16) bool {
	other := a.mol.atomWithIid(aiid)

	return a.rings.IntersectionCardinality(other.rings) > 0
}

// inSameRingsAs answers if this atom participates in exactly the same
// rings as the given atom.
func (a *Atom) inSameRingsAs(aiid uint16) bool {
	other := a.mol.atomWithIid(aiid)

	return a.rings.Equal(other.rings)
}

// inAllRingsOf answers if this atom participates in every ring in
// which the given atom does.
//
// Note that this atom may participate in more rings, as well.
func (a *Atom) inAllRingsOf(aiid uint16) bool {
	other := a.mol.atomWithIid(aiid)

	return other.rings.DifferenceCardinality(a.rings) == 0
}

// addRing adds the given ring to the list of this atom's rings.
func (a *Atom) addRing(r *Ring) {
	a.rings.Set(uint(r.id))
}

// removeRing removes the given ring from the list of this atom's
// rings.
//
// This method should, usually, not be called.  The composition of a
// ring never changes once it is completed.  However, during a
// reaction, a ring may get broken, leading to its death.  The
// constituent atoms are then notified of that death by calling this
// method.
func (a *Atom) removeRing(r *Ring) {
	a.rings.Clear(r.id)
}
<commit_msg>Fix an issue with bitmap input type in `Atom.removeRing`<commit_after>package data

import (
	"fmt"
	"math"

	bits "github.com/willf/bitset"

	cmn "github.com/RxnWeaver/rxnweaver/common"
)

// Atom represents a chemical atom.
//
// It has various physical and chemical properties, apart from the
// transient information attached to it during its participation in
// reactions.
type Atom struct {
	mol    *Molecule // Containing molecule of this atom.
	atNum  uint8     // Atomic number of this atom's element.
	symbol string    // Symbol, in case of a different isotope.
	iId    uint16    // Serial input ID of this atom.
	nId    uint16    // Normalised ID of this atom.

	X float32 // X-coordinate of this atom.
	Y float32 // Y-coordinate of this atom.
	Z float32 // Z-coordinate of this atom.

	numH    uint8 // Number of implicit + explicit H atoms attached to this atom.
	charge  int8  // Residual net charge of this atom.
	valence int8  // Current valence configuration of this atom.

	pHash uint64 // A pseudo-hash of this atom, using some attributes.
	sHash uint64 // A pseudo-hash of this atom, using some attributes.

	bonds          *bits.BitSet // Bitmap of bonds of this atom.
	nbrs           []uint16     // Expanded list of neighbours of this atom.
	numSingleBonds uint8        // Number of single bonds this atom has.
	numDoubleBonds uint8        // Number of double bonds this atom has.
	numTripleBonds uint8        // Number of triple bonds this atom has.

	rings *bits.BitSet // Bitmap of IDs of rings this atom participates in.
	// Does this atom participate in at least one aromatic ring?
	isInAroRing bool
	// Is this atom a bridgehead of a bicyclic system of rings?
	isBridgeHead bool
	// Is this atom the sole common atom of all of its rings?
	isSpiro bool

	// The functional groups substituted on this atom.  They are listed in
	// descending order of importance.  The first is the primary feature.
	features [cmn.MaxFeatures]uint8
}

// newAtom constructs and initialises a new atom of the given element
// type, and belonging to the given molecule.
func newAtom(mol *Molecule, atNum uint8) *Atom {
	atom := new(Atom)
	atom.mol = mol
	atom.atNum = atNum
	atom.valence = cmn.PeriodicTable[cmn.ElemSyms[atNum]].Valence

	atom.bonds = bits.New(cmn.MaxBonds)
	atom.nbrs = make([]uint16, 0, cmn.MaxBonds)
	atom.rings = bits.New(cmn.MaxRings)

	return atom
}

// AtomicNumber answers the atomic number of this atom.
func (a *Atom) AtomicNumber() uint8 {
	return a.atNum
}

// Parent answers the parent molecule of this atom.
func (a *Atom) Parent() *Molecule {
	return a.mol
}

// numPiElectrons answers the number of delocalised pi electrons
// contributed by this atom.  This number is important for calculating
// the aromaticity of the rings this atom participates in.
func (a *Atom) numPiElectrons() int {
	mol := a.mol
	wtSum := 100*int16(a.numDoubleBonds) + 10*int16(a.numSingleBonds) + int16(a.charge)

	switch a.atNum {
	case 6:
		switch wtSum {
		case 19:
			return 2
		case 110:
			return 1
		case 120:
			var b *Bond
			for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
				b = mol.bondWithId(uint16(bid - 1))
				if b.bType == cmn.BondTypeDouble {
					break
				}
			}
			if b.isCyclic() {
				return 1
			}
			return 0
		default:
			return 0
		}

	case 7:
		switch wtSum {
		case 20:
			return 2
		case 30:
			return 2
		case 110:
			return 1
		case 121:
			return 1
		default:
			return 0
		}

	case 8:
		switch wtSum {
		case 20:
			return 2
		case 111:
			return 1
		default:
			return 0
		}

	case 16:
		switch wtSum {
		case 20:
			return 2
		case 111:
			return 1
		case 120:
			var b *Bond
			for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
				b = mol.bondWithId(uint16(bid - 1))
				if b.bType == cmn.BondTypeDouble {
					break
				}
			}
			oaid := b.otherAtom(a.iId)
			oa := mol.atomWithIid(oaid)
			if oa.atNum == 8 && !oa.isCyclic() {
				return 2
			}
			return 0
		case 220:
			c := 0
			for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
				b := mol.bondWithId(uint16(bid - 1))
				if b.bType == cmn.BondTypeDouble {
					oaid := b.otherAtom(a.iId)
					if !mol.atomWithIid(oaid).isCyclic() {
						c++
					}
				}
			}
			if c > 1 {
				return -1
			}
			return 0
		default:
			return 0
		}
	}

	return 0
}

// isCyclic answers if this atom participates in at least one ring.
func (a *Atom) isCyclic() bool {
	return a.rings.Count() > 0
}

// isJunction answers if this atom has more than 2 distinct
// neighbours.
func (a *Atom) isJunction() bool {
	return a.bonds.Count() > 2
}

// addBond adds the given bond to this atom, if it is not already
// present.  It also adjusts the list of its neighbours appropriately.
//
// Note that it does NOT check to see if the addition conforms to this
// atom's current valence configuration.
func (a *Atom) addBond(b *Bond) {
	for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
		if uint16(bid) == b.id {
			return
		}
	}

	a.bonds.Set(uint(b.id))
	nbrId := b.otherAtom(a.iId)
	n := int(b.bType)
	for i := 0; i < n; i++ {
		a.nbrs = append(a.nbrs, nbrId)
	}
}

// removeBond removes the given bond from this atom, and adjusts the
// list of neighbours appropriately.
//
// Note that it does NOT check to see if the removal conforms to this
// atom's current valence configuration.
func (a *Atom) removeBond(b *Bond) {
	nbrId := b.otherAtom(a.iId)

	wid := 0
	for _, nid := range a.nbrs {
		if nid == nbrId {
			continue
		}
		a.nbrs[wid] = nid
		wid++
	}
	a.nbrs = a.nbrs[:wid]

	a.bonds.Clear(uint(b.id))
}

// bondTo answers the bond that binds this atom to the given atom, if
// one such bond exists.  Answers `nil` otherwise.
func (a *Atom) bondTo(other uint16) *Bond {
	mol := a.mol
	for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
		b := mol.bondWithId(uint16(bid - 1))
		if b.otherAtom(a.iId) == other {
			return b
		}
	}

	return nil
}

// firstDoublyBondedNbr answers this atom's doubly-bonded neighbour
// having the highest priority.
//
// This method assumes that the molecule is already normalised!
// Calling it on a molecule that has not be normalised yet, leads to
// incorrect results.
func (a *Atom) firstDoublyBondedNbr() uint16 {
	if a.numDoubleBonds == 0 {
		return 0
	}

	mol := a.mol
	for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
		b := mol.bondWithId(uint16(bid - 1))
		if b.bType == cmn.BondTypeDouble {
			return b.otherAtom(a.iId)
		}
	}

	panic("Should never be here!")
}

// firstMultiplyBondedNbr answers this atom's doubly-bonded neighbour
// having the highest priority.
//
// This method assumes that the molecule is already normalised!
// Calling it on a molecule that has not be normalised yet, leads to
// incorrect results.
func (a *Atom) firstMultiplyBondedNbr() uint16 {
	if a.numDoubleBonds == 0 && a.numTripleBonds == 0 {
		return 0
	}

	mol := a.mol
	for bid, ok := a.bonds.NextSet(0); ok; bid, ok = a.bonds.NextSet(bid + 1) {
		b := mol.bondWithId(uint16(bid - 1))
		if b.bType >= cmn.BondTypeDouble {
			return b.otherAtom(a.iId)
		}
	}

	panic("Should never be here!")
}

// inRingOfSize answers if this atom participates in at least one ring
// of the given size.
func (a *Atom) inRingOfSize(n int) bool {
	mol := a.mol
	for rid, ok := a.rings.NextSet(0); ok; rid, ok = a.rings.NextSet(rid + 1) {
		r := mol.ringWithId(uint8(rid))
		if r.size() == n {
			return true
		}
	}

	return false
}

// inRingLargerThan answers if this atom participates in at least one
// ring that is larger than the given number.
func (a *Atom) inRingLargerThan(n int) bool {
	mol := a.mol
	for rid, ok := a.rings.NextSet(0); ok; rid, ok = a.rings.NextSet(rid + 1) {
		r := mol.ringWithId(uint8(rid))
		if r.size() > n {
			return true
		}
	}

	return false
}

// smallestRing answers the smallest unique ring in which this atom
// participates.  Should no such unique ring exist, an error is
// answered.
func (a *Atom) smallestRing() (uint8, error) {
	if !a.isCyclic() {
		return 0, fmt.Errorf("Atom not cyclic.")
	}

	min := int(math.MaxUint8)
	c := 0
	var ret uint8

	mol := a.mol
	for rid, ok := a.rings.NextSet(0); ok; rid, ok = a.rings.NextSet(rid + 1) {
		r := mol.ringWithId(uint8(rid))
		size := r.size()
		if size == min {
			c++
		} else if size < min {
			ret = uint8(rid)
			c = 1
		}
	}

	if c > 1 {
		return 0, fmt.Errorf("Smallest ring size: %d, number of smallest rings: %d.", min, c)
	}

	return ret, nil
}

// isAromatic answers if this atom is part of an aromatic ring.
//
// Note that the actual aromaticity determination is handled by
// `Ring`.  This method merely answers the set flag.
func (a *Atom) isAromatic() bool {
	return a.isInAroRing
}

// inHetAromaticRing answers if this atom is part of an aromatic ring
// with at least one hetero atom.
func (a *Atom) inHetAromaticRing() bool {
	if a.isInAroRing && a.atNum != 6 {
		return true // Simplest case!
	}

	mol := a.mol
	for rid, ok := a.rings.NextSet(0); ok; rid, ok = a.rings.NextSet(rid + 1) {
		r := mol.ringWithId(uint8(rid))
		if r.isHetAromatic() {
			return true
		}
	}

	return false
}

// haveCommonRings answers if this atom has at least one ring in
// common with the given atom.
func (a *Atom) haveCommonRings(aiid uint16) bool {
	other := a.mol.atomWithIid(aiid)

	return a.rings.IntersectionCardinality(other.rings) > 0
}

// inSameRingsAs answers if this atom participates in exactly the same
// rings as the given atom.
func (a *Atom) inSameRingsAs(aiid uint16) bool {
	other := a.mol.atomWithIid(aiid)

	return a.rings.Equal(other.rings)
}

// inAllRingsOf answers if this atom participates in every ring in
// which the given atom does.
//
// Note that this atom may participate in more rings, as well.
func (a *Atom) inAllRingsOf(aiid uint16) bool {
	other := a.mol.atomWithIid(aiid)

	return other.rings.DifferenceCardinality(a.rings) == 0
}

// addRing adds the given ring to the list of this atom's rings.
func (a *Atom) addRing(r *Ring) {
	a.rings.Set(uint(r.id))
}

// removeRing removes the given ring from the list of this atom's
// rings.
//
// This method should, usually, not be called.  The composition of a
// ring never changes once it is completed.  However, during a
// reaction, a ring may get broken, leading to its death.  The
// constituent atoms are then notified of that death by calling this
// method.
func (a *Atom) removeRing(r *Ring) {
	a.rings.Clear(uint(r.id))
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"strings"
	"testing"

	mkr "github.com/mackerelio/mackerel-client-go"
)

func TestIsSameMonitor(t *testing.T) {
	a := &mkr.Monitor{ID: "12345", Name: "foo", Type: "connectivity"}
	b := &mkr.Monitor{Name: "foo", Type: "connectivity"}

	_, ret := isSameMonitor(a, b, true)
	if ret != true {
		t.Error("should recognize same monitors")
	}

	_, ret = isSameMonitor(a, b, false)
	if ret == true {
		t.Error("should not recognize same monitors")
	}
}

func TestValidateRoles(t *testing.T) {
	a := &mkr.Monitor{ID: "12345", Name: "foo", Type: "connectivity"}

	ret, err := validateRules([](*mkr.Monitor){a}, "test monitor")
	if ret != true {
		t.Errorf("should validate the rule: %s", err.Error())
	}

}

func TestDiffMonitors(t *testing.T) {
	a := &mkr.Monitor{ID: "12345", Name: "foo", Type: "external", URL: "http://example.com", Service: "bar", ResponseTimeCritical: 1000}
	b := &mkr.Monitor{ID: "12345", Name: "foo", Type: "external", URL: "http://example.com", Service: "bar"}

	ret := diffMonitor(a, b)

	correct := strings.Join([]string{
		"  {",
		"    \"name\": \"foo\",",
		"    \"type\": \"external\",",
		"    \"url\": \"http://example.com\",",
		"    \"service\": \"bar\",",
		"-   \"responseTimeCritical\": 1000.000000,",
		"+   \"responseTimeCritical\": 0.000000,",
		"  },",
	}, "\n")

	if ret != correct {
		t.Errorf("should validate the rule: %s\nbut result: %s", correct, ret)
	}
}

func TestMonitorSaveRules(t *testing.T) {
	tmpFile, err := ioutil.TempFile("", "")
	if err != nil {
		t.Errorf("should not raise error: %v", err)
	}
	a := &mkr.Monitor{
		ID:                   "12345",
		Name:                 "foo",
		Type:                 "external",
		URL:                  "http://example.com",
		Service:              "bar",
		ResponseTimeCritical: 1000,
	}
	monitorSaveRules([]*mkr.Monitor{a}, tmpFile.Name())

	byt, _ := ioutil.ReadFile(tmpFile.Name())
	content := string(byt)
	expected := `{
    "monitors": [
        {
            "id": "12345",
            "name": "foo",
            "type": "external",
            "url": "http://example.com",
            "service": "bar",
            "responseTimeCritical": 1000
        }
    ]
}
`
	if content != expected {
		t.Errorf("content should be:\n %s, but:\n %s", expected, content)
	}
}

func TestStringifyMonitor(t *testing.T) {
	a := &mkr.Monitor{ID: "12345", Name: "foo", Type: "connectivity"}
	expected := `+ {
+   "id": "12345",
+   "name": "foo",
+   "type": "connectivity"
+ },`

	r := stringifyMonitor(a, "+")
	if r != expected {
		t.Errorf("stringifyMonitor should be:\n%s\nbut:\n%s", expected, r)
	}
}
<commit_msg>add failing test<commit_after>package main

import (
	"io/ioutil"
	"strings"
	"testing"

	mkr "github.com/mackerelio/mackerel-client-go"
)

func TestIsSameMonitor(t *testing.T) {
	a := &mkr.Monitor{ID: "12345", Name: "foo", Type: "connectivity"}
	b := &mkr.Monitor{Name: "foo", Type: "connectivity"}

	_, ret := isSameMonitor(a, b, true)
	if ret != true {
		t.Error("should recognize same monitors")
	}

	_, ret = isSameMonitor(a, b, false)
	if ret == true {
		t.Error("should not recognize same monitors")
	}
}

func TestValidateRoles(t *testing.T) {
	a := &mkr.Monitor{ID: "12345", Name: "foo", Type: "connectivity"}

	ret, err := validateRules([](*mkr.Monitor){a}, "test monitor")
	if ret != true {
		t.Errorf("should validate the rule: %s", err.Error())
	}

}

func TestDiffMonitors(t *testing.T) {
	a := &mkr.Monitor{ID: "12345", Name: "foo", Type: "external", URL: "http://example.com", Service: "bar", ResponseTimeCritical: 1000}
	b := &mkr.Monitor{ID: "12345", Name: "foo", Type: "external", URL: "http://example.com", Service: "bar"}

	ret := diffMonitor(a, b)

	correct := strings.Join([]string{
		"  {",
		"    \"name\": \"foo\",",
		"    \"type\": \"external\",",
		"    \"url\": \"http://example.com\",",
		"    \"service\": \"bar\",",
		"-   \"responseTimeCritical\": 1000.000000,",
		"+   \"responseTimeCritical\": 0.000000,",
		"  },",
	}, "\n")

	if ret != correct {
		t.Errorf("should validate the rule: %s\nbut result: %s", correct, ret)
	}
}

func TestMonitorSaveRules(t *testing.T) {
	tmpFile, err := ioutil.TempFile("", "")
	if err != nil {
		t.Errorf("should not raise error: %v", err)
	}
	a := &mkr.Monitor{
		ID:                   "12345",
		Name:                 "foo",
		Type:                 "external",
		URL:                  "http://example.com",
		Service:              "bar",
		ResponseTimeCritical: 1000,
	}
	monitorSaveRules([]*mkr.Monitor{a}, tmpFile.Name())

	byt, _ := ioutil.ReadFile(tmpFile.Name())
	content := string(byt)
	expected := `{
    "monitors": [
        {
            "id": "12345",
            "name": "foo",
            "type": "external",
            "url": "http://example.com",
            "service": "bar",
            "responseTimeCritical": 1000
        }
    ]
}
`
	if content != expected {
		t.Errorf("content should be:\n %s, but:\n %s", expected, content)
	}
}

func TestStringifyMonitor(t *testing.T) {
	a := &mkr.Monitor{ID: "12345", Name: "foo", Type: "connectivity"}
	expected := `+ {
+   "id": "12345",
+   "name": "foo",
+   "type": "connectivity"
+ },`

	r := stringifyMonitor(a, "+")
	if r != expected {
		t.Errorf("stringifyMonitor should be:\n%s\nbut:\n%s", expected, r)
	}
}

func TestDiffMonitorsWithScopes(t *testing.T) {
	a := &mkr.Monitor{
		ID:   "12345",
		Name: "foo",
		Type: "connectivity",
	}
	b := &mkr.Monitor{
		ID:     "12345",
		Name:   "foo",
		Type:   "connectivity",
		Scopes: []string{"sss: notebook"},
	}
	diff := diffMonitor(a, b)
	expected := `  {
    "name": "foo",
    "type": "connectivity",
    "scopes": [
+     "sss: notebook",
    ],
  },`
	if diff != expected {
		t.Errorf("expected:\n%s\n, output:\n%s\n", expected, diff)
	}

}
<|endoftext|>


<commit_before>package data

import (
	_ "log"
	"math"
)

type Data struct {
	Attributes []float64
	Distances []float64
	Classification int
}

func New(attr []float64, numClass int) *Data {
	return &Data{
		Attributes: attr,
		Distances: make([]float64, numClass),
		Classification: 0,
	}
}

func (data *Data) updateDistances(centroids []*Data){
	for i, centroid := range(centroids){
		sumOfSquares := 0.0;
		for j := range(centroid.Attributes){
			sumOfSquares += math.Pow(data.Attributes[j] - centroid.Attributes[j], 2)
		}
		distance := math.Sqrt(sumOfSquares)
		data.Distances[i] = distance
	}
}

func (data *Data) UpdateClassification(centroids []*Data){
	data.updateDistances(centroids)
	minDistance := math.Inf(1)
	for i, distance := range(data.Distances) {
		if distance < minDistance {
			minDistance = distance
			data.Classification = i
		}
	}
}
<commit_msg>Added a weigthed averaging function for calculating the centroids<commit_after>package data

import (
	_ "log"
	"math"
)

type Data struct {
	Attributes []float64
	Distances []float64
	Classification int
}

func New(attr []float64, numClass int) *Data {
	return &Data{
		Attributes: attr,
		Distances: make([]float64, numClass),
		Classification: 0,
	}
}

func (data *Data) updateDistances(centroids []*Data){
	for i, centroid := range(centroids){
		sumOfSquares := 0.0;
		for j := range(centroid.Attributes){
			sumOfSquares += math.Pow(data.Attributes[j] - centroid.Attributes[j], 2)
		}
		distance := math.Sqrt(sumOfSquares)
		data.Distances[i] = distance
	}
}

func (data *Data) UpdateClassification(centroids []*Data){
	data.updateDistances(centroids)
	minDistance := math.Inf(1)
	for i, distance := range(data.Distances) {
		if distance < minDistance {
			minDistance = distance
			data.Classification = i
		}
	}
}

func (data* Data) Waverage(item *Data, weigth float64) {
	for i := range(data.Attributes){
		data.Attributes[i] = data.Attributes[i]*(1.0-weigth) + item.Attributes[i]*weigth
	}
}
<|endoftext|>