<commit_before>package testing

import (
	"bytes"
	"fmt"
	"testing"
	"time"

	"go.pedge.io/proto/stream"
	"go.pedge.io/protolog"

	"golang.org/x/net/context"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pkg/require"
	"github.com/pachyderm/pachyderm/src/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/pps"
)

const (
	defaultJobTimeoutSec = 30
)

var (
	baseSimpleTest = &simpleTest{
		inputFilePathToContent: map[string][]byte{
			"1": []byte("test1"),
			"2": []byte("test2"),
			"3": []byte("test3"),
		},
		outputFilePathToContent: map[string][]byte{
			"1.output": []byte("test1"),
			"2.output": []byte("test2"),
			"3.output": []byte("test3"),
		},
		transform: &pps.Transform{
			Image: "ubuntu:14.04",
			Cmd: []string{
				fmt.Sprintf("for file in /pfs/*; do cp \"${file}\" /pfs/$(basename \"${file}\").output; done"),
			},
		},
	}
)

func TestSimpleJob(t *testing.T) {
	t.Skip()
	runSimpleJobTest(t, baseSimpleTest.copy())
}

func TestSimplePipeline(t *testing.T) {
	t.Skip()
	runSimplePipelineTest(t, baseSimpleTest.copy())
}

func runSimpleJobTest(t *testing.T, simpleTest *simpleTest) {
	RunTest(t, simpleTest.runJob)
}

func runSimplePipelineTest(t *testing.T, simpleTest *simpleTest) {
	RunTest(t, simpleTest.runPipeline)
}

type simpleTest struct {
	inputRepoName           string
	outputRepoName          string
	inputFilePathToContent  map[string][]byte
	outputFilePathToContent map[string][]byte
	transform               *pps.Transform
	pipelineName            string
	jobTimeoutSec           int
	expectError             bool
}

func (s *simpleTest) copy() *simpleTest {
	c := &simpleTest{
		inputRepoName:           s.inputRepoName,
		outputRepoName:          s.outputRepoName,
		inputFilePathToContent:  make(map[string][]byte),
		outputFilePathToContent: make(map[string][]byte),
		pipelineName:            s.pipelineName,
		jobTimeoutSec:           s.jobTimeoutSec,
		expectError:             s.expectError,
	}
	for key, value := range s.inputFilePathToContent {
		buf := make([]byte, len(value))
		copy(buf, value)
		c.inputFilePathToContent[key] = buf
	}
	for key, value := range s.outputFilePathToContent {
		buf := make([]byte, len(value))
		copy(buf, value)
		c.outputFilePathToContent[key] = buf
	}
	if s.transform != nil {
		c.transform = &pps.Transform{
			Image: s.transform.Image,
			Cmd:   make([]string, len(s.transform.Cmd)),
		}
		for i, value := range s.transform.Cmd {
			c.transform.Cmd[i] = value
		}
	}
	return c
}

func (s *simpleTest) setup() {
	if s.inputRepoName == "" {
		s.inputRepoName = uuid.NewWithoutDashes()
	}
	if s.outputRepoName == "" {
		s.outputRepoName = uuid.NewWithoutDashes()
	}
	if s.pipelineName == "" {
		s.pipelineName = uuid.NewWithoutDashes()
	}
	if s.jobTimeoutSec == 0 {
		s.jobTimeoutSec = defaultJobTimeoutSec
	}
}

func (s *simpleTest) runJob(
	t *testing.T,
	pfsAPIClient pfs.APIClient,
	jobAPIClient pps.JobAPIClient,
	pipelineAPIClient pps.PipelineAPIClient,
) {
	s.setup()
	inputRepo := setupPFSRepo(t, pfsAPIClient, s.inputRepoName)
	outputRepo := setupPFSRepo(t, pfsAPIClient, s.outputRepoName)
	inputCommit := setupPFSInputCommit(t, pfsAPIClient, inputRepo, s.inputFilePathToContent)
	outputParentCommit := setupPFSOutputParentCommit(t, pfsAPIClient, outputRepo)
	job := createJob(t, jobAPIClient, s.transform, inputCommit, outputParentCommit)
	jobInfo := waitForJob(t, jobAPIClient, job, s.jobTimeoutSec, s.expectError)
	checkPFSOutput(t, pfsAPIClient, jobInfo.Output, s.outputFilePathToContent)
}

func (s *simpleTest) runPipeline(
	t *testing.T,
	pfsAPIClient pfs.APIClient,
	jobAPIClient pps.JobAPIClient,
	pipelineAPIClient pps.PipelineAPIClient,
) {
	s.setup()
	inputRepo := setupPFSRepo(t, pfsAPIClient, s.inputRepoName)
	outputRepo := setupPFSRepo(t, pfsAPIClient, s.outputRepoName)
	pipeline := createPipeline(t, pipelineAPIClient, s.pipelineName, s.transform, inputRepo, outputRepo)
	setupPFSInputCommit(t, pfsAPIClient, inputRepo, s.inputFilePathToContent)
	// TODO: all the waits heh
	time.Sleep(5 * time.Second)
	job := getJobForPipeline(t, jobAPIClient, pipeline)
	jobInfo := waitForJob(t, jobAPIClient, job, s.jobTimeoutSec, s.expectError)
	checkPFSOutput(t, pfsAPIClient, jobInfo.Output, s.outputFilePathToContent)
}

func setupPFSRepo(t *testing.T, pfsAPIClient pfs.APIClient, repoName string) *pfs.Repo {
	repo := &pfs.Repo{
		Name: repoName,
	}
	_, err := pfsAPIClient.CreateRepo(
		context.Background(),
		&pfs.CreateRepoRequest{
			Repo: repo,
		},
	)
	require.NoError(t, err)
	return repo
}

// TODO: handle directories in filePathToContent
func setupPFSInputCommit(t *testing.T, pfsAPIClient pfs.APIClient, repo *pfs.Repo, filePathToContent map[string][]byte) *pfs.Commit {
	commit, err := pfsAPIClient.StartCommit(
		context.Background(),
		&pfs.StartCommitRequest{
			Parent: &pfs.Commit{
				Repo: repo,
				Id:   pfs.InitialCommitID,
			},
		},
	)
	require.NoError(t, err)
	for filePath, content := range filePathToContent {
		apiPutFileClient, err := pfsAPIClient.PutFile(
			context.Background(),
		)
		require.NoError(t, err)
		err = apiPutFileClient.Send(
			&pfs.PutFileRequest{
				File: &pfs.File{
					Commit: commit,
					Path:   filePath,
				},
				FileType: pfs.FileType_FILE_TYPE_REGULAR,
				Value:    content,
			},
		)
		require.NoError(t, err)
		_, _ = apiPutFileClient.CloseAndRecv()
	}
	_, err = pfsAPIClient.FinishCommit(
		context.Background(),
		&pfs.FinishCommitRequest{
			Commit: commit,
		},
	)
	require.NoError(t, err)
	return commit
}

func setupPFSOutputParentCommit(t *testing.T, pfsAPIClient pfs.APIClient, repo *pfs.Repo) *pfs.Commit {
	return &pfs.Commit{Repo: repo, Id: pfs.InitialCommitID}
}

func createJob(t *testing.T, jobAPIClient pps.JobAPIClient, transform *pps.Transform, inputCommit *pfs.Commit, outputParentCommit *pfs.Commit) *pps.Job {
	job, err := jobAPIClient.CreateJob(
		context.Background(),
		&pps.CreateJobRequest{
			Spec: &pps.CreateJobRequest_Transform{
				Transform: transform,
			},
			Input:        inputCommit,
			OutputParent: outputParentCommit,
		},
	)
	require.NoError(t, err)
	return job
}

func waitForJob(t *testing.T, jobAPIClient pps.JobAPIClient, job *pps.Job, timeoutSec int, expectError bool) *pps.JobInfo {
	for i := 0; i < timeoutSec; i++ {
		time.Sleep(1 * time.Second)
		jobInfo, err := jobAPIClient.InspectJob(
			context.Background(),
			&pps.InspectJobRequest{
				Job: job,
			},
		)
		require.NoError(t, err)
		if len(jobInfo.JobStatus) == 0 {
			continue
		}
		jobStatus := jobInfo.JobStatus[0]
		protolog.Infof("status of job %s at %d seconds: %v", job.Id, i+1, jobInfo.JobStatus)
		switch jobStatus.Type {
		case pps.JobStatusType_JOB_STATUS_TYPE_ERROR:
			handleEndJobStatus(t, jobAPIClient, jobInfo)
			if !expectError {
				t.Fatalf("job %s had error", job.Id)
			}
			return jobInfo
		case pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS:
			handleEndJobStatus(t, jobAPIClient, jobInfo)
			if expectError {
				t.Fatalf("job %s did not have error", job.Id)
			}
			return jobInfo
		}
	}
	t.Fatalf("job %s did not finish in %d seconds", job.Id, timeoutSec)
	return nil
}

func handleEndJobStatus(t *testing.T, jobAPIClient pps.JobAPIClient, jobInfo *pps.JobInfo) {
	handleEndJobStatusOutputStream(t, jobAPIClient, jobInfo, pps.OutputStream_OUTPUT_STREAM_STDOUT)
	handleEndJobStatusOutputStream(t, jobAPIClient, jobInfo, pps.OutputStream_OUTPUT_STREAM_STDERR)
}

func handleEndJobStatusOutputStream(t *testing.T, jobAPIClient pps.JobAPIClient, jobInfo *pps.JobInfo, outputStream pps.OutputStream) {
	jobAPIGetJobLogsClient, err := jobAPIClient.GetJobLogs(
		context.Background(),
		&pps.GetJobLogsRequest{
			Job:          jobInfo.Job,
			OutputStream: outputStream,
		},
	)
	require.NoError(t, err)
	require.NoError(t, protostream.WriteFromStreamingBytesClient(jobAPIGetJobLogsClient, protolog.Writer()))
}

func createPipeline(t *testing.T, pipelineAPIClient pps.PipelineAPIClient, pipelineName string, transform *pps.Transform, inputRepo *pfs.Repo, outputRepo *pfs.Repo) *pps.Pipeline {
	pipeline := &pps.Pipeline{
		Name: pipelineName,
	}
	_, err := pipelineAPIClient.CreatePipeline(
		context.Background(),
		&pps.CreatePipelineRequest{
			Pipeline:  pipeline,
			Transform: transform,
			Input:     inputRepo,
			Output:    outputRepo,
		},
	)
	require.NoError(t, err)
	return pipeline
}

func getJobForPipeline(t *testing.T, jobAPIClient pps.JobAPIClient, pipeline *pps.Pipeline) *pps.Job {
	jobInfos, err := jobAPIClient.ListJob(
		context.Background(),
		&pps.ListJobRequest{
			Pipeline: pipeline,
		},
	)
	require.NoError(t, err)
	require.NotNil(t, jobInfos)
	require.Equal(t, 1, len(jobInfos.JobInfo))
	return jobInfos.JobInfo[0].Job
}

func checkPFSOutput(t *testing.T, pfsAPIClient pfs.APIClient, outputCommit *pfs.Commit, filePathToContent map[string][]byte) {
	for filePath, content := range filePathToContent {
		getContent, err := getPFSContent(pfsAPIClient, outputCommit, filePath)
		require.NoError(t, err)
		require.Equal(t, content, getContent)
	}
}

func getPFSContent(pfsAPIClient pfs.APIClient, commit *pfs.Commit, filePath string) ([]byte, error) {
	apiGetFileClient, err := pfsAPIClient.GetFile(
		context.Background(),
		&pfs.GetFileRequest{
			File: &pfs.File{
				Commit: commit,
				Path:   filePath,
			},
		},
	)
	if err != nil {
		return nil, err
	}
	buffer := bytes.NewBuffer(nil)
	if err := protostream.WriteFromStreamingBytesClient(apiGetFileClient, buffer); err != nil {
		return nil, err
	}
	return buffer.Bytes(), nil
}
<commit_msg>Make pps tests compile.<commit_after>package testing

import (
	"bytes"
	"fmt"
	"testing"
	"time"

	"go.pedge.io/proto/stream"
	"go.pedge.io/protolog"

	"golang.org/x/net/context"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pkg/require"
	"github.com/pachyderm/pachyderm/src/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/pps"
)

const (
	defaultJobTimeoutSec = 30
)

var (
	baseSimpleTest = &simpleTest{
		inputFilePathToContent: map[string][]byte{
			"1": []byte("test1"),
			"2": []byte("test2"),
			"3": []byte("test3"),
		},
		outputFilePathToContent: map[string][]byte{
			"1.output": []byte("test1"),
			"2.output": []byte("test2"),
			"3.output": []byte("test3"),
		},
		transform: &pps.Transform{
			Image: "ubuntu:14.04",
			Cmd: []string{
				fmt.Sprintf("for file in /pfs/*; do cp \"${file}\" /pfs/$(basename \"${file}\").output; done"),
			},
		},
	}
)

func TestSimpleJob(t *testing.T) {
	t.Skip()
	runSimpleJobTest(t, baseSimpleTest.copy())
}

func TestSimplePipeline(t *testing.T) {
	t.Skip()
	runSimplePipelineTest(t, baseSimpleTest.copy())
}

func runSimpleJobTest(t *testing.T, simpleTest *simpleTest) {
	RunTest(t, simpleTest.runJob)
}

func runSimplePipelineTest(t *testing.T, simpleTest *simpleTest) {
	RunTest(t, simpleTest.runPipeline)
}

type simpleTest struct {
	inputRepoName           string
	outputRepoName          string
	inputFilePathToContent  map[string][]byte
	outputFilePathToContent map[string][]byte
	transform               *pps.Transform
	pipelineName            string
	jobTimeoutSec           int
	expectError             bool
}

func (s *simpleTest) copy() *simpleTest {
	c := &simpleTest{
		inputRepoName:           s.inputRepoName,
		outputRepoName:          s.outputRepoName,
		inputFilePathToContent:  make(map[string][]byte),
		outputFilePathToContent: make(map[string][]byte),
		pipelineName:            s.pipelineName,
		jobTimeoutSec:           s.jobTimeoutSec,
		expectError:             s.expectError,
	}
	for key, value := range s.inputFilePathToContent {
		buf := make([]byte, len(value))
		copy(buf, value)
		c.inputFilePathToContent[key] = buf
	}
	for key, value := range s.outputFilePathToContent {
		buf := make([]byte, len(value))
		copy(buf, value)
		c.outputFilePathToContent[key] = buf
	}
	if s.transform != nil {
		c.transform = &pps.Transform{
			Image: s.transform.Image,
			Cmd:   make([]string, len(s.transform.Cmd)),
		}
		for i, value := range s.transform.Cmd {
			c.transform.Cmd[i] = value
		}
	}
	return c
}

func (s *simpleTest) setup() {
	if s.inputRepoName == "" {
		s.inputRepoName = uuid.NewWithoutDashes()
	}
	if s.outputRepoName == "" {
		s.outputRepoName = uuid.NewWithoutDashes()
	}
	if s.pipelineName == "" {
		s.pipelineName = uuid.NewWithoutDashes()
	}
	if s.jobTimeoutSec == 0 {
		s.jobTimeoutSec = defaultJobTimeoutSec
	}
}

func (s *simpleTest) runJob(
	t *testing.T,
	pfsAPIClient pfs.APIClient,
	jobAPIClient pps.JobAPIClient,
	pipelineAPIClient pps.PipelineAPIClient,
) {
	s.setup()
	inputRepo := setupPFSRepo(t, pfsAPIClient, s.inputRepoName)
	outputRepo := setupPFSRepo(t, pfsAPIClient, s.outputRepoName)
	inputCommit := setupPFSInputCommit(t, pfsAPIClient, inputRepo, s.inputFilePathToContent)
	outputParentCommit := setupPFSOutputParentCommit(t, pfsAPIClient, outputRepo)
	job := createJob(t, jobAPIClient, s.transform, inputCommit, outputParentCommit)
	jobInfo := waitForJob(t, jobAPIClient, job, s.jobTimeoutSec, s.expectError)
	checkPFSOutput(t, pfsAPIClient, jobInfo.OutputCommit, s.outputFilePathToContent)
}

func (s *simpleTest) runPipeline(
	t *testing.T,
	pfsAPIClient pfs.APIClient,
	jobAPIClient pps.JobAPIClient,
	pipelineAPIClient pps.PipelineAPIClient,
) {
	s.setup()
	inputRepo := setupPFSRepo(t, pfsAPIClient, s.inputRepoName)
	outputRepo := setupPFSRepo(t, pfsAPIClient, s.outputRepoName)
	pipeline := createPipeline(t, pipelineAPIClient, s.pipelineName, s.transform, inputRepo, outputRepo)
	setupPFSInputCommit(t, pfsAPIClient, inputRepo, s.inputFilePathToContent)
	// TODO: all the waits heh
	time.Sleep(5 * time.Second)
	job := getJobForPipeline(t, jobAPIClient, pipeline)
	jobInfo := waitForJob(t, jobAPIClient, job, s.jobTimeoutSec, s.expectError)
	checkPFSOutput(t, pfsAPIClient, jobInfo.OutputCommit, s.outputFilePathToContent)
}

func setupPFSRepo(t *testing.T, pfsAPIClient pfs.APIClient, repoName string) *pfs.Repo {
	repo := &pfs.Repo{
		Name: repoName,
	}
	_, err := pfsAPIClient.CreateRepo(
		context.Background(),
		&pfs.CreateRepoRequest{
			Repo: repo,
		},
	)
	require.NoError(t, err)
	return repo
}

// TODO: handle directories in filePathToContent
func setupPFSInputCommit(t *testing.T, pfsAPIClient pfs.APIClient, repo *pfs.Repo, filePathToContent map[string][]byte) *pfs.Commit {
	commit, err := pfsAPIClient.StartCommit(
		context.Background(),
		&pfs.StartCommitRequest{
			Parent: &pfs.Commit{
				Repo: repo,
				Id:   pfs.InitialCommitID,
			},
		},
	)
	require.NoError(t, err)
	for filePath, content := range filePathToContent {
		apiPutFileClient, err := pfsAPIClient.PutFile(
			context.Background(),
		)
		require.NoError(t, err)
		err = apiPutFileClient.Send(
			&pfs.PutFileRequest{
				File: &pfs.File{
					Commit: commit,
					Path:   filePath,
				},
				FileType: pfs.FileType_FILE_TYPE_REGULAR,
				Value:    content,
			},
		)
		require.NoError(t, err)
		_, _ = apiPutFileClient.CloseAndRecv()
	}
	_, err = pfsAPIClient.FinishCommit(
		context.Background(),
		&pfs.FinishCommitRequest{
			Commit: commit,
		},
	)
	require.NoError(t, err)
	return commit
}

func setupPFSOutputParentCommit(t *testing.T, pfsAPIClient pfs.APIClient, repo *pfs.Repo) *pfs.Commit {
	return &pfs.Commit{Repo: repo, Id: pfs.InitialCommitID}
}

func createJob(t *testing.T, jobAPIClient pps.JobAPIClient, transform *pps.Transform, inputCommit *pfs.Commit, outputParentCommit *pfs.Commit) *pps.Job {
	job, err := jobAPIClient.CreateJob(
		context.Background(),
		&pps.CreateJobRequest{
			Spec: &pps.CreateJobRequest_Transform{
				Transform: transform,
			},
			Input:        inputCommit,
			OutputParent: outputParentCommit,
		},
	)
	require.NoError(t, err)
	return job
}

func waitForJob(t *testing.T, jobAPIClient pps.JobAPIClient, job *pps.Job, timeoutSec int, expectError bool) *pps.JobInfo {
	for i := 0; i < timeoutSec; i++ {
		time.Sleep(1 * time.Second)
		jobInfo, err := jobAPIClient.InspectJob(
			context.Background(),
			&pps.InspectJobRequest{
				Job: job,
			},
		)
		require.NoError(t, err)
		if len(jobInfo.JobStatus) == 0 {
			continue
		}
		jobStatus := jobInfo.JobStatus[0]
		protolog.Infof("status of job %s at %d seconds: %v", job.Id, i+1, jobInfo.JobStatus)
		switch jobStatus.Type {
		case pps.JobStatusType_JOB_STATUS_TYPE_ERROR:
			handleEndJobStatus(t, jobAPIClient, jobInfo)
			if !expectError {
				t.Fatalf("job %s had error", job.Id)
			}
			return jobInfo
		case pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS:
			handleEndJobStatus(t, jobAPIClient, jobInfo)
			if expectError {
				t.Fatalf("job %s did not have error", job.Id)
			}
			return jobInfo
		}
	}
	t.Fatalf("job %s did not finish in %d seconds", job.Id, timeoutSec)
	return nil
}

func handleEndJobStatus(t *testing.T, jobAPIClient pps.JobAPIClient, jobInfo *pps.JobInfo) {
	handleEndJobStatusOutputStream(t, jobAPIClient, jobInfo, pps.OutputStream_OUTPUT_STREAM_STDOUT)
	handleEndJobStatusOutputStream(t, jobAPIClient, jobInfo, pps.OutputStream_OUTPUT_STREAM_STDERR)
}

func handleEndJobStatusOutputStream(t *testing.T, jobAPIClient pps.JobAPIClient, jobInfo *pps.JobInfo, outputStream pps.OutputStream) {
	jobAPIGetJobLogsClient, err := jobAPIClient.GetJobLogs(
		context.Background(),
		&pps.GetJobLogsRequest{
			Job:          jobInfo.Job,
			OutputStream: outputStream,
		},
	)
	require.NoError(t, err)
	require.NoError(t, protostream.WriteFromStreamingBytesClient(jobAPIGetJobLogsClient, protolog.Writer()))
}

func createPipeline(t *testing.T, pipelineAPIClient pps.PipelineAPIClient, pipelineName string, transform *pps.Transform, inputRepo *pfs.Repo, outputRepo *pfs.Repo) *pps.Pipeline {
	pipeline := &pps.Pipeline{
		Name: pipelineName,
	}
	_, err := pipelineAPIClient.CreatePipeline(
		context.Background(),
		&pps.CreatePipelineRequest{
			Pipeline:  pipeline,
			Transform: transform,
			Input:     inputRepo,
			Output:    outputRepo,
		},
	)
	require.NoError(t, err)
	return pipeline
}

func getJobForPipeline(t *testing.T, jobAPIClient pps.JobAPIClient, pipeline *pps.Pipeline) *pps.Job {
	jobInfos, err := jobAPIClient.ListJob(
		context.Background(),
		&pps.ListJobRequest{
			Pipeline: pipeline,
		},
	)
	require.NoError(t, err)
	require.NotNil(t, jobInfos)
	require.Equal(t, 1, len(jobInfos.JobInfo))
	return jobInfos.JobInfo[0].Job
}

func checkPFSOutput(t *testing.T, pfsAPIClient pfs.APIClient, outputCommit *pfs.Commit, filePathToContent map[string][]byte) {
	for filePath, content := range filePathToContent {
		getContent, err := getPFSContent(pfsAPIClient, outputCommit, filePath)
		require.NoError(t, err)
		require.Equal(t, content, getContent)
	}
}

func getPFSContent(pfsAPIClient pfs.APIClient, commit *pfs.Commit, filePath string) ([]byte, error) {
	apiGetFileClient, err := pfsAPIClient.GetFile(
		context.Background(),
		&pfs.GetFileRequest{
			File: &pfs.File{
				Commit: commit,
				Path:   filePath,
			},
		},
	)
	if err != nil {
		return nil, err
	}
	buffer := bytes.NewBuffer(nil)
	if err := protostream.WriteFromStreamingBytesClient(apiGetFileClient, buffer); err != nil {
		return nil, err
	}
	return buffer.Bytes(), nil
}
<|endoftext|>


<commit_before>package testing

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"go.pedge.io/proto/test"
	"go.pedge.io/protolog"

	"github.com/fsouza/go-dockerclient"
	"github.com/satori/go.uuid"
	"go.pachyderm.com/pachyderm/src/pfs"
	pfstesting "go.pachyderm.com/pachyderm/src/pfs/testing"
	"go.pachyderm.com/pachyderm/src/pkg/container"
	"go.pachyderm.com/pachyderm/src/pkg/require"
	"go.pachyderm.com/pachyderm/src/pps"
	"go.pachyderm.com/pachyderm/src/pps/jobserver"
	"go.pachyderm.com/pachyderm/src/pps/persist"
	persistserver "go.pachyderm.com/pachyderm/src/pps/persist/server"
	"go.pachyderm.com/pachyderm/src/pps/pipelineserver"
	"google.golang.org/grpc"
)

const (
	testNumServers = 1
)

func TestCreateAndGetPipeline(t *testing.T) {
	t.Skip()
	runTest(t, testCreateAndGetPipeline)
}

func testCreateAndGetPipeline(t *testing.T, jobAPIClient pps.JobAPIClient, pipelineAPIClient pps.PipelineAPIClient) {
	expectedPipeline := &pps.Pipeline{
		Name: "foo",
		Transform: &pps.Transform{
			Image: "ubuntu:14.04",
			Cmd: []string{
				"which",
				"bash",
			},
		},
		PipelineInput: []*pps.PipelineInput{
			&pps.PipelineInput{
				Input: &pps.PipelineInput_HostDir{
					HostDir: "/path/to/dir",
				},
			},
		},
	}
	pipeline, err := pipelineAPIClient.CreatePipeline(
		context.Background(),
		&pps.CreatePipelineRequest{
			Pipeline: expectedPipeline,
		},
	)
	require.NoError(t, err)
	require.Equal(t, expectedPipeline, pipeline)
	getPipeline, err := pipelineAPIClient.GetPipeline(
		context.Background(),
		&pps.GetPipelineRequest{
			PipelineName: "foo",
		},
	)
	require.NoError(t, err)
	require.Equal(t, expectedPipeline, getPipeline)
}

func TestBasicCreateAndStartJob(t *testing.T) {
	runTest(t, testBasicCreateAndStartJob)
}

func testBasicCreateAndStartJob(t *testing.T, jobAPIClient pps.JobAPIClient, pipelineAPIClient pps.PipelineAPIClient) {
	inputDir, err := ioutil.TempDir("/tmp/pachyderm-test", "")
	require.NoError(t, err)
	outputDir, err := ioutil.TempDir("/tmp/pachyderm-test", "")
	require.NoError(t, err)
	file, err := os.Create(filepath.Join(inputDir, "foo.txt"))
	require.NoError(t, err)
	_, err = file.Write([]byte("hello"))
	require.NoError(t, err)
	require.NoError(t, file.Close())
	job := &pps.Job{
		Spec: &pps.Job_Transform{
			Transform: &pps.Transform{
				Image: "ubuntu:14.04",
				Cmd: []string{
					fmt.Sprintf("for i in /var/lib/pps/host/%s/*; do cp $i /var/lib/pps/host/%s/; done", inputDir, outputDir),
				},
			},
		},
		JobInput: []*pps.JobInput{
			&pps.JobInput{
				Input: &pps.JobInput_HostDir{
					HostDir: inputDir,
				},
			},
		},
		JobOutput: []*pps.JobOutput{
			&pps.JobOutput{
				Output: &pps.JobOutput_HostDir{
					HostDir: outputDir,
				},
			},
		},
	}
	createJob, err := jobAPIClient.CreateJob(
		context.Background(),
		&pps.CreateJobRequest{
			Job: job,
		},
	)
	require.NoError(t, err)
	_, err = jobAPIClient.StartJob(
		context.Background(),
		&pps.StartJobRequest{
			JobId: createJob.Id,
		},
	)
	require.NoError(t, err)
	jobStatus, err := getFinalJobStatus(jobAPIClient, createJob.Id)
	require.NoError(t, err)
	require.Equal(t, pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS, jobStatus.Type)
	data, err := ioutil.ReadFile(filepath.Join(outputDir, "foo.txt"))
	require.NoError(t, err)
	require.Equal(t, []byte("hello"), data)
}

func getFinalJobStatus(jobAPIClient pps.JobAPIClient, jobID string) (*pps.JobStatus, error) {
	// TODO(pedge): not good
	ticker := time.NewTicker(time.Second)
	for i := 0; i < 20; i++ {
		<-ticker.C
		jobStatus, err := jobAPIClient.GetJobStatus(
			context.Background(),
			&pps.GetJobStatusRequest{
				JobId: jobID,
			},
		)
		if err != nil {
			return nil, err
		}
		protolog.Printf("status at tick %d: %v\n", i, jobStatus)
		switch jobStatus.Type {
		case pps.JobStatusType_JOB_STATUS_TYPE_ERROR:
			return jobStatus, nil
		case pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS:
			return jobStatus, nil
		}
	}
	return nil, fmt.Errorf("did not get final job status for %s", jobID)
}

func runTest(
	t *testing.T,
	f func(t *testing.T, jobAPIClient pps.JobAPIClient, pipelineAPIClient pps.PipelineAPIClient),
) {
	containerClient, err := getTestContainerClient()
	require.NoError(t, err)
	persistAPIServer, err := getTestRethinkAPIServer()
	require.NoError(t, err)
	persistAPIClient := persist.NewLocalAPIClient(persistAPIServer)
	pfstesting.RunTest(
		t,
		func(t *testing.T, apiClient pfs.APIClient, internalAPIClient pfs.InternalAPIClient, cluster pfstesting.Cluster) {
			prototest.RunT(
				t,
				testNumServers,
				func(servers map[string]*grpc.Server) {
					jobAPIServer := jobserver.NewAPIServer(persistAPIClient, containerClient)
					jobAPIClient := pps.NewLocalJobAPIClient(jobAPIServer)
					for _, server := range servers {
						pps.RegisterJobAPIServer(server, jobAPIServer)
						pps.RegisterPipelineAPIServer(server, pipelineserver.NewTestAPIServer(apiClient, jobAPIClient, persistAPIClient))
					}
				},
				func(t *testing.T, clientConns map[string]*grpc.ClientConn) {
					var clientConn *grpc.ClientConn
					for _, c := range clientConns {
						clientConn = c
						break
					}
					f(
						t,
						pps.NewJobAPIClient(
							clientConn,
						),
						pps.NewPipelineAPIClient(
							clientConn,
						),
					)
				},
			)
		},
	)
}

func getTestContainerClient() (container.Client, error) {
	client, err := docker.NewClientFromEnv()
	if err != nil {
		return nil, err
	}
	return container.NewDockerClient(client), nil
}

func getTestRethinkAPIServer() (persist.APIServer, error) {
	address, err := getTestRethinkAddress()
	if err != nil {
		return nil, err
	}
	databaseName := strings.Replace(uuid.NewV4().String(), "-", "", -1)
	if err := persistserver.InitDBs(address, databaseName); err != nil {
		return nil, err
	}
	return persistserver.NewRethinkAPIServer(address, databaseName)
}

func getTestRethinkAddress() (string, error) {
	rethinkAddr := os.Getenv("RETHINK_PORT_28015_TCP_ADDR")
	if rethinkAddr == "" {
		return "", errors.New("RETHINK_PORT_28015_TCP_ADDR not set")
	}
	return fmt.Sprintf("%s:28015", rethinkAddr), nil
}
<commit_msg>comment out everything but imports in src/pps/testing/testing_test.go<commit_after>package testing

/*
const (
	testNumServers = 1
)

func TestCreateAndGetPipeline(t *testing.T) {
	t.Skip()
	runTest(t, testCreateAndGetPipeline)
}

func testCreateAndGetPipeline(t *testing.T, jobAPIClient pps.JobAPIClient, pipelineAPIClient pps.PipelineAPIClient) {
	expectedPipeline := &pps.Pipeline{
		Name: "foo",
		Transform: &pps.Transform{
			Image: "ubuntu:14.04",
			Cmd: []string{
				"which",
				"bash",
			},
		},
		PipelineInput: []*pps.PipelineInput{
			&pps.PipelineInput{
				Input: &pps.PipelineInput_HostDir{
					HostDir: "/path/to/dir",
				},
			},
		},
	}
	pipeline, err := pipelineAPIClient.CreatePipeline(
		context.Background(),
		&pps.CreatePipelineRequest{
			Pipeline: expectedPipeline,
		},
	)
	require.NoError(t, err)
	require.Equal(t, expectedPipeline, pipeline)
	getPipeline, err := pipelineAPIClient.GetPipeline(
		context.Background(),
		&pps.GetPipelineRequest{
			PipelineName: "foo",
		},
	)
	require.NoError(t, err)
	require.Equal(t, expectedPipeline, getPipeline)
}

func TestBasicCreateAndStartJob(t *testing.T) {
	runTest(t, testBasicCreateAndStartJob)
}

func testBasicCreateAndStartJob(t *testing.T, jobAPIClient pps.JobAPIClient, pipelineAPIClient pps.PipelineAPIClient) {
	inputDir, err := ioutil.TempDir("/tmp/pachyderm-test", "")
	require.NoError(t, err)
	outputDir, err := ioutil.TempDir("/tmp/pachyderm-test", "")
	require.NoError(t, err)
	file, err := os.Create(filepath.Join(inputDir, "foo.txt"))
	require.NoError(t, err)
	_, err = file.Write([]byte("hello"))
	require.NoError(t, err)
	require.NoError(t, file.Close())
	job := &pps.Job{
		Spec: &pps.Job_Transform{
			Transform: &pps.Transform{
				Image: "ubuntu:14.04",
				Cmd: []string{
					fmt.Sprintf("for i in /var/lib/pps/host/%s/*; do cp $i /var/lib/pps/host/%s/; done", inputDir, outputDir),
				},
			},
		},
		JobInput: []*pps.JobInput{
			&pps.JobInput{
				Input: &pps.JobInput_HostDir{
					HostDir: inputDir,
				},
			},
		},
		JobOutput: []*pps.JobOutput{
			&pps.JobOutput{
				Output: &pps.JobOutput_HostDir{
					HostDir: outputDir,
				},
			},
		},
	}
	createJob, err := jobAPIClient.CreateJob(
		context.Background(),
		&pps.CreateJobRequest{
			Job: job,
		},
	)
	require.NoError(t, err)
	_, err = jobAPIClient.StartJob(
		context.Background(),
		&pps.StartJobRequest{
			JobId: createJob.Id,
		},
	)
	require.NoError(t, err)
	jobStatus, err := getFinalJobStatus(jobAPIClient, createJob.Id)
	require.NoError(t, err)
	require.Equal(t, pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS, jobStatus.Type)
	data, err := ioutil.ReadFile(filepath.Join(outputDir, "foo.txt"))
	require.NoError(t, err)
	require.Equal(t, []byte("hello"), data)
}

func getFinalJobStatus(jobAPIClient pps.JobAPIClient, jobID string) (*pps.JobStatus, error) {
	// TODO(pedge): not good
	ticker := time.NewTicker(time.Second)
	for i := 0; i < 20; i++ {
		<-ticker.C
		jobStatus, err := jobAPIClient.GetJobStatus(
			context.Background(),
			&pps.GetJobStatusRequest{
				JobId: jobID,
			},
		)
		if err != nil {
			return nil, err
		}
		protolog.Printf("status at tick %d: %v\n", i, jobStatus)
		switch jobStatus.Type {
		case pps.JobStatusType_JOB_STATUS_TYPE_ERROR:
			return jobStatus, nil
		case pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS:
			return jobStatus, nil
		}
	}
	return nil, fmt.Errorf("did not get final job status for %s", jobID)
}

func runTest(
	t *testing.T,
	f func(t *testing.T, jobAPIClient pps.JobAPIClient, pipelineAPIClient pps.PipelineAPIClient),
) {
	containerClient, err := getTestContainerClient()
	require.NoError(t, err)
	persistAPIServer, err := getTestRethinkAPIServer()
	require.NoError(t, err)
	persistAPIClient := persist.NewLocalAPIClient(persistAPIServer)
	pfstesting.RunTest(
		t,
		func(t *testing.T, apiClient pfs.APIClient, internalAPIClient pfs.InternalAPIClient, cluster pfstesting.Cluster) {
			prototest.RunT(
				t,
				testNumServers,
				func(servers map[string]*grpc.Server) {
					jobAPIServer := jobserver.NewAPIServer(persistAPIClient, containerClient)
					jobAPIClient := pps.NewLocalJobAPIClient(jobAPIServer)
					for _, server := range servers {
						pps.RegisterJobAPIServer(server, jobAPIServer)
						pps.RegisterPipelineAPIServer(server, pipelineserver.NewTestAPIServer(apiClient, jobAPIClient, persistAPIClient))
					}
				},
				func(t *testing.T, clientConns map[string]*grpc.ClientConn) {
					var clientConn *grpc.ClientConn
					for _, c := range clientConns {
						clientConn = c
						break
					}
					f(
						t,
						pps.NewJobAPIClient(
							clientConn,
						),
						pps.NewPipelineAPIClient(
							clientConn,
						),
					)
				},
			)
		},
	)
}

func getTestContainerClient() (container.Client, error) {
	client, err := docker.NewClientFromEnv()
	if err != nil {
		return nil, err
	}
	return container.NewDockerClient(client), nil
}

func getTestRethinkAPIServer() (persist.APIServer, error) {
	address, err := getTestRethinkAddress()
	if err != nil {
		return nil, err
	}
	databaseName := strings.Replace(uuid.NewV4().String(), "-", "", -1)
	if err := persistserver.InitDBs(address, databaseName); err != nil {
		return nil, err
	}
	return persistserver.NewRethinkAPIServer(address, databaseName)
}

func getTestRethinkAddress() (string, error) {
	rethinkAddr := os.Getenv("RETHINK_PORT_28015_TCP_ADDR")
	if rethinkAddr == "" {
		return "", errors.New("RETHINK_PORT_28015_TCP_ADDR not set")
	}
	return fmt.Sprintf("%s:28015", rethinkAddr), nil
}
*/
<|endoftext|>


<commit_before>package manager

/*
 * 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.
 */

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/common/log"
	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/common/poller"
	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/traffic_monitor/config"
	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/traffic_monitor/enum"
	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/traffic_monitor/peer"
	to "github.com/apache/incubator-trafficcontrol/traffic_ops/client"
)

// CopyTrafficMonitorConfigMap returns a deep copy of the given TrafficMonitorConfigMap
func CopyTrafficMonitorConfigMap(a *to.TrafficMonitorConfigMap) to.TrafficMonitorConfigMap {
	b := to.TrafficMonitorConfigMap{}
	b.TrafficServer = map[string]to.TrafficServer{}
	b.CacheGroup = map[string]to.TMCacheGroup{}
	b.Config = map[string]interface{}{}
	b.TrafficMonitor = map[string]to.TrafficMonitor{}
	b.DeliveryService = map[string]to.TMDeliveryService{}
	b.Profile = map[string]to.TMProfile{}
	for k, v := range a.TrafficServer {
		b.TrafficServer[k] = v
	}
	for k, v := range a.CacheGroup {
		b.CacheGroup[k] = v
	}
	for k, v := range a.Config {
		b.Config[k] = v
	}
	for k, v := range a.TrafficMonitor {
		b.TrafficMonitor[k] = v
	}
	for k, v := range a.DeliveryService {
		b.DeliveryService[k] = v
	}
	for k, v := range a.Profile {
		b.Profile[k] = v
	}
	return b
}

// TrafficMonitorConfigMapThreadsafe encapsulates a TrafficMonitorConfigMap safe for multiple readers and a single writer.
type TrafficMonitorConfigMapThreadsafe struct {
	monitorConfig *to.TrafficMonitorConfigMap
	m             *sync.RWMutex
}

// NewTrafficMonitorConfigMapThreadsafe returns an encapsulated TrafficMonitorConfigMap safe for multiple readers and a single writer.
func NewTrafficMonitorConfigMapThreadsafe() TrafficMonitorConfigMapThreadsafe {
	return TrafficMonitorConfigMapThreadsafe{monitorConfig: &to.TrafficMonitorConfigMap{}, m: &sync.RWMutex{}}
}

// Get returns the TrafficMonitorConfigMap. Callers MUST NOT modify, it is not threadsafe for mutation. If mutation is necessary, call CopyTrafficMonitorConfigMap().
func (t *TrafficMonitorConfigMapThreadsafe) Get() to.TrafficMonitorConfigMap {
	t.m.RLock()
	defer t.m.RUnlock()
	return *t.monitorConfig
}

// Set sets the TrafficMonitorConfigMap. This is only safe for one writer. This MUST NOT be called by multiple threads.
func (t *TrafficMonitorConfigMapThreadsafe) Set(c to.TrafficMonitorConfigMap) {
	t.m.Lock()
	*t.monitorConfig = c
	t.m.Unlock()
}

// StartMonitorConfigManager runs the monitor config manager goroutine, and returns the threadsafe data which it sets.
func StartMonitorConfigManager(
	monitorConfigPollChan <-chan to.TrafficMonitorConfigMap,
	localStates peer.CRStatesThreadsafe,
	statURLSubscriber chan<- poller.HttpPollerConfig,
	healthURLSubscriber chan<- poller.HttpPollerConfig,
	peerURLSubscriber chan<- poller.HttpPollerConfig,
	cachesChangeSubscriber chan<- struct{},
	cfg config.Config,
	staticAppData StaticAppData,
) TrafficMonitorConfigMapThreadsafe {
	monitorConfig := NewTrafficMonitorConfigMapThreadsafe()
	go monitorConfigListen(monitorConfig,
		monitorConfigPollChan,
		localStates,
		statURLSubscriber,
		healthURLSubscriber,
		peerURLSubscriber,
		cachesChangeSubscriber,
		cfg,
		staticAppData,
	)
	return monitorConfig
}

// trafficOpsHealthConnectionTimeoutToDuration takes the int from Traffic Ops, which is in milliseconds, and returns a time.Duration
// TODO change Traffic Ops Client API to a time.Duration
func trafficOpsHealthConnectionTimeoutToDuration(t int) time.Duration {
	return time.Duration(t) * time.Millisecond
}

// trafficOpsPeerPollIntervalToDuration takes the int from Traffic Ops, which is in milliseconds, and returns a time.Duration
// TODO change Traffic Ops Client API to a time.Duration
func trafficOpsPeerPollIntervalToDuration(t int) time.Duration {
	return time.Duration(t) * time.Millisecond
}

// trafficOpsStatPollIntervalToDuration takes the int from Traffic Ops, which is in milliseconds, and returns a time.Duration
// TODO change Traffic Ops Client API to a time.Duration
func trafficOpsStatPollIntervalToDuration(t int) time.Duration {
	return time.Duration(t) * time.Millisecond
}

// trafficOpsHealthPollIntervalToDuration takes the int from Traffic Ops, which is in milliseconds, and returns a time.Duration
// TODO change Traffic Ops Client API to a time.Duration
func trafficOpsHealthPollIntervalToDuration(t int) time.Duration {
	return time.Duration(t) * time.Millisecond
}

// getPollIntervals reads the Traffic Ops Client monitorConfig structure, and parses and returns the health, peer, and stat poll intervals
func getHealthPeerStatPollIntervals(monitorConfig to.TrafficMonitorConfigMap, cfg config.Config) (time.Duration, time.Duration, time.Duration, error) {
	peerPollIntervalI, peerPollIntervalExists := monitorConfig.Config["peers.polling.interval"]
	if !peerPollIntervalExists {
		return 0, 0, 0, fmt.Errorf("Traffic Ops Monitor config missing 'peers.polling.interval', not setting config changes.\n")
	}
	peerPollIntervalInt, peerPollIntervalIsInt := peerPollIntervalI.(float64)
	if !peerPollIntervalIsInt {
		return 0, 0, 0, fmt.Errorf("Traffic Ops Monitor config 'peers.polling.interval' value '%v' type %T is not an integer, not setting config changes.\n", peerPollIntervalI, peerPollIntervalI)
	}
	peerPollInterval := trafficOpsPeerPollIntervalToDuration(int(peerPollIntervalInt))

	statPollIntervalI, statPollIntervalExists := monitorConfig.Config["health.polling.interval"]
	if !statPollIntervalExists {
		return 0, 0, 0, fmt.Errorf("Traffic Ops Monitor config missing 'health.polling.interval', not setting config changes.\n")
	}
	statPollIntervalInt, statPollIntervalIsInt := statPollIntervalI.(float64)
	if !statPollIntervalIsInt {
		return 0, 0, 0, fmt.Errorf("Traffic Ops Monitor config 'health.polling.interval' value '%v' type %T is not an integer, not setting config changes.\n", statPollIntervalI, statPollIntervalI)
	}
	statPollInterval := trafficOpsStatPollIntervalToDuration(int(statPollIntervalInt))

	healthPollIntervalI, healthPollIntervalExists := monitorConfig.Config["heartbeat.polling.interval"]
	healthPollIntervalInt, healthPollIntervalIsInt := healthPollIntervalI.(float64)
	if !healthPollIntervalExists {
		log.Warnf("Traffic Ops Monitor config missing 'heartbeat.polling.interval', using health for heartbeat.\n")
		healthPollIntervalInt = statPollIntervalInt
	} else if !healthPollIntervalIsInt {
		log.Warnf("Traffic Ops Monitor config 'heartbeat.polling.interval' value '%v' type %T is not an integer, using health for heartbeat\n", statPollIntervalI, statPollIntervalI)
		healthPollIntervalInt = statPollIntervalInt
	}
	healthPollInterval := trafficOpsHealthPollIntervalToDuration(int(healthPollIntervalInt))

	return healthPollInterval, peerPollInterval, statPollInterval, nil
}

// TODO timing, and determine if the case, or its internal `for`, should be put in a goroutine
// TODO determine if subscribers take action on change, and change to mutexed objects if not.
func monitorConfigListen(
	monitorConfigTS TrafficMonitorConfigMapThreadsafe,
	monitorConfigPollChan <-chan to.TrafficMonitorConfigMap,
	localStates peer.CRStatesThreadsafe,
	statURLSubscriber chan<- poller.HttpPollerConfig,
	healthURLSubscriber chan<- poller.HttpPollerConfig,
	peerURLSubscriber chan<- poller.HttpPollerConfig,
	cachesChangeSubscriber chan<- struct{},
	cfg config.Config,
	staticAppData StaticAppData,
) {
	for monitorConfig := range monitorConfigPollChan {
		monitorConfigTS.Set(monitorConfig)
		healthURLs := map[string]poller.PollConfig{}
		statURLs := map[string]poller.PollConfig{}
		peerURLs := map[string]poller.PollConfig{}
		caches := map[string]string{}

		healthPollInterval, peerPollInterval, statPollInterval, err := getHealthPeerStatPollIntervals(monitorConfig, cfg)
		if err != nil {
			continue
		}

		for _, srv := range monitorConfig.TrafficServer {
			caches[srv.HostName] = srv.Status

			cacheName := enum.CacheName(srv.HostName)

			srvStatus := enum.CacheStatusFromString(srv.Status)
			if srvStatus == enum.CacheStatusOnline {
				localStates.SetCache(cacheName, peer.IsAvailable{IsAvailable: true})
				continue
			}
			if srvStatus == enum.CacheStatusOffline {
				continue
			}
			// seed states with available = false until our polling cycle picks up a result
			if _, exists := localStates.GetCache(cacheName); !exists {
				localStates.SetCache(cacheName, peer.IsAvailable{IsAvailable: false})
			}

			url := monitorConfig.Profile[srv.Profile].Parameters.HealthPollingURL
			if url == "" {
				log.Errorf("monitor config server %v profile %v has no polling URL; can't poll", srv.HostName, srv.Profile)
				continue
			}
			r := strings.NewReplacer(
				"${hostname}", srv.IP,
				"${interface_name}", srv.InterfaceName,
				"application=system", "application=plugin.remap",
				"application=", "application=plugin.remap",
			)
			url = r.Replace(url)

			connTimeout := trafficOpsHealthConnectionTimeoutToDuration(monitorConfig.Profile[srv.Profile].Parameters.HealthConnectionTimeout)
			healthURLs[srv.HostName] = poller.PollConfig{URL: url, Timeout: connTimeout}
			r = strings.NewReplacer("application=plugin.remap", "application=")
			statURL := r.Replace(url)
			statURLs[srv.HostName] = poller.PollConfig{URL: statURL, Timeout: connTimeout}
		}

		for _, srv := range monitorConfig.TrafficMonitor {
			if srv.HostName == staticAppData.Hostname {
				continue
			}
			if enum.CacheStatusFromString(srv.Status) != enum.CacheStatusOnline {
				continue
			}
			// TODO: the URL should be config driven. -jse
			url := fmt.Sprintf("http://%s:%d/publish/CrStates?raw", srv.IP, srv.Port)
			peerURLs[srv.HostName] = poller.PollConfig{URL: url} // TODO determine timeout.
		}

		statURLSubscriber <- poller.HttpPollerConfig{Urls: statURLs, Interval: statPollInterval}
		healthURLSubscriber <- poller.HttpPollerConfig{Urls: healthURLs, Interval: healthPollInterval}
		peerURLSubscriber <- poller.HttpPollerConfig{Urls: peerURLs, Interval: peerPollInterval}

		for cacheName := range localStates.GetCaches() {
			if _, exists := monitorConfig.TrafficServer[string(cacheName)]; !exists {
				log.Warnf("Removing %s from localStates", cacheName)
				localStates.DeleteCache(cacheName)
			}
		}

		cachesChangeSubscriber <- struct{}{}

		// TODO because there are multiple writers to localStates.DeliveryService, there is a race condition, where MonitorConfig (this func) and HealthResultManager could write at the same time, and the HealthResultManager could overwrite a delivery service addition or deletion here. Probably the simplest and most performant fix would be a lock-free algorithm using atomic compare-and-swaps.
		for _, ds := range monitorConfig.DeliveryService {
			// since caches default to unavailable, also default DS false
			if _, exists := localStates.GetDeliveryService(enum.DeliveryServiceName(ds.XMLID)); !exists {
				localStates.SetDeliveryService(enum.DeliveryServiceName(ds.XMLID), peer.Deliveryservice{IsAvailable: false, DisabledLocations: []enum.CacheName{}}) // important to initialize DisabledLocations, so JSON is `[]` not `null`
			}
		}
		for ds := range localStates.GetDeliveryServices() {
			if _, exists := monitorConfig.DeliveryService[string(ds)]; !exists {
				localStates.DeleteDeliveryService(ds)
			}
		}
	}
}
<commit_msg>Fix TM2 health poll to only get system stats<commit_after>package manager

/*
 * 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.
 */

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/common/log"
	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/common/poller"
	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/traffic_monitor/config"
	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/traffic_monitor/enum"
	"github.com/apache/incubator-trafficcontrol/traffic_monitor/experimental/traffic_monitor/peer"
	to "github.com/apache/incubator-trafficcontrol/traffic_ops/client"
)

// CopyTrafficMonitorConfigMap returns a deep copy of the given TrafficMonitorConfigMap
func CopyTrafficMonitorConfigMap(a *to.TrafficMonitorConfigMap) to.TrafficMonitorConfigMap {
	b := to.TrafficMonitorConfigMap{}
	b.TrafficServer = map[string]to.TrafficServer{}
	b.CacheGroup = map[string]to.TMCacheGroup{}
	b.Config = map[string]interface{}{}
	b.TrafficMonitor = map[string]to.TrafficMonitor{}
	b.DeliveryService = map[string]to.TMDeliveryService{}
	b.Profile = map[string]to.TMProfile{}
	for k, v := range a.TrafficServer {
		b.TrafficServer[k] = v
	}
	for k, v := range a.CacheGroup {
		b.CacheGroup[k] = v
	}
	for k, v := range a.Config {
		b.Config[k] = v
	}
	for k, v := range a.TrafficMonitor {
		b.TrafficMonitor[k] = v
	}
	for k, v := range a.DeliveryService {
		b.DeliveryService[k] = v
	}
	for k, v := range a.Profile {
		b.Profile[k] = v
	}
	return b
}

// TrafficMonitorConfigMapThreadsafe encapsulates a TrafficMonitorConfigMap safe for multiple readers and a single writer.
type TrafficMonitorConfigMapThreadsafe struct {
	monitorConfig *to.TrafficMonitorConfigMap
	m             *sync.RWMutex
}

// NewTrafficMonitorConfigMapThreadsafe returns an encapsulated TrafficMonitorConfigMap safe for multiple readers and a single writer.
func NewTrafficMonitorConfigMapThreadsafe() TrafficMonitorConfigMapThreadsafe {
	return TrafficMonitorConfigMapThreadsafe{monitorConfig: &to.TrafficMonitorConfigMap{}, m: &sync.RWMutex{}}
}

// Get returns the TrafficMonitorConfigMap. Callers MUST NOT modify, it is not threadsafe for mutation. If mutation is necessary, call CopyTrafficMonitorConfigMap().
func (t *TrafficMonitorConfigMapThreadsafe) Get() to.TrafficMonitorConfigMap {
	t.m.RLock()
	defer t.m.RUnlock()
	return *t.monitorConfig
}

// Set sets the TrafficMonitorConfigMap. This is only safe for one writer. This MUST NOT be called by multiple threads.
func (t *TrafficMonitorConfigMapThreadsafe) Set(c to.TrafficMonitorConfigMap) {
	t.m.Lock()
	*t.monitorConfig = c
	t.m.Unlock()
}

// StartMonitorConfigManager runs the monitor config manager goroutine, and returns the threadsafe data which it sets.
func StartMonitorConfigManager(
	monitorConfigPollChan <-chan to.TrafficMonitorConfigMap,
	localStates peer.CRStatesThreadsafe,
	statURLSubscriber chan<- poller.HttpPollerConfig,
	healthURLSubscriber chan<- poller.HttpPollerConfig,
	peerURLSubscriber chan<- poller.HttpPollerConfig,
	cachesChangeSubscriber chan<- struct{},
	cfg config.Config,
	staticAppData StaticAppData,
) TrafficMonitorConfigMapThreadsafe {
	monitorConfig := NewTrafficMonitorConfigMapThreadsafe()
	go monitorConfigListen(monitorConfig,
		monitorConfigPollChan,
		localStates,
		statURLSubscriber,
		healthURLSubscriber,
		peerURLSubscriber,
		cachesChangeSubscriber,
		cfg,
		staticAppData,
	)
	return monitorConfig
}

// trafficOpsHealthConnectionTimeoutToDuration takes the int from Traffic Ops, which is in milliseconds, and returns a time.Duration
// TODO change Traffic Ops Client API to a time.Duration
func trafficOpsHealthConnectionTimeoutToDuration(t int) time.Duration {
	return time.Duration(t) * time.Millisecond
}

// trafficOpsPeerPollIntervalToDuration takes the int from Traffic Ops, which is in milliseconds, and returns a time.Duration
// TODO change Traffic Ops Client API to a time.Duration
func trafficOpsPeerPollIntervalToDuration(t int) time.Duration {
	return time.Duration(t) * time.Millisecond
}

// trafficOpsStatPollIntervalToDuration takes the int from Traffic Ops, which is in milliseconds, and returns a time.Duration
// TODO change Traffic Ops Client API to a time.Duration
func trafficOpsStatPollIntervalToDuration(t int) time.Duration {
	return time.Duration(t) * time.Millisecond
}

// trafficOpsHealthPollIntervalToDuration takes the int from Traffic Ops, which is in milliseconds, and returns a time.Duration
// TODO change Traffic Ops Client API to a time.Duration
func trafficOpsHealthPollIntervalToDuration(t int) time.Duration {
	return time.Duration(t) * time.Millisecond
}

// getPollIntervals reads the Traffic Ops Client monitorConfig structure, and parses and returns the health, peer, and stat poll intervals
func getHealthPeerStatPollIntervals(monitorConfig to.TrafficMonitorConfigMap, cfg config.Config) (time.Duration, time.Duration, time.Duration, error) {
	peerPollIntervalI, peerPollIntervalExists := monitorConfig.Config["peers.polling.interval"]
	if !peerPollIntervalExists {
		return 0, 0, 0, fmt.Errorf("Traffic Ops Monitor config missing 'peers.polling.interval', not setting config changes.\n")
	}
	peerPollIntervalInt, peerPollIntervalIsInt := peerPollIntervalI.(float64)
	if !peerPollIntervalIsInt {
		return 0, 0, 0, fmt.Errorf("Traffic Ops Monitor config 'peers.polling.interval' value '%v' type %T is not an integer, not setting config changes.\n", peerPollIntervalI, peerPollIntervalI)
	}
	peerPollInterval := trafficOpsPeerPollIntervalToDuration(int(peerPollIntervalInt))

	statPollIntervalI, statPollIntervalExists := monitorConfig.Config["health.polling.interval"]
	if !statPollIntervalExists {
		return 0, 0, 0, fmt.Errorf("Traffic Ops Monitor config missing 'health.polling.interval', not setting config changes.\n")
	}
	statPollIntervalInt, statPollIntervalIsInt := statPollIntervalI.(float64)
	if !statPollIntervalIsInt {
		return 0, 0, 0, fmt.Errorf("Traffic Ops Monitor config 'health.polling.interval' value '%v' type %T is not an integer, not setting config changes.\n", statPollIntervalI, statPollIntervalI)
	}
	statPollInterval := trafficOpsStatPollIntervalToDuration(int(statPollIntervalInt))

	healthPollIntervalI, healthPollIntervalExists := monitorConfig.Config["heartbeat.polling.interval"]
	healthPollIntervalInt, healthPollIntervalIsInt := healthPollIntervalI.(float64)
	if !healthPollIntervalExists {
		log.Warnf("Traffic Ops Monitor config missing 'heartbeat.polling.interval', using health for heartbeat.\n")
		healthPollIntervalInt = statPollIntervalInt
	} else if !healthPollIntervalIsInt {
		log.Warnf("Traffic Ops Monitor config 'heartbeat.polling.interval' value '%v' type %T is not an integer, using health for heartbeat\n", statPollIntervalI, statPollIntervalI)
		healthPollIntervalInt = statPollIntervalInt
	}
	healthPollInterval := trafficOpsHealthPollIntervalToDuration(int(healthPollIntervalInt))

	return healthPollInterval, peerPollInterval, statPollInterval, nil
}

// TODO timing, and determine if the case, or its internal `for`, should be put in a goroutine
// TODO determine if subscribers take action on change, and change to mutexed objects if not.
func monitorConfigListen(
	monitorConfigTS TrafficMonitorConfigMapThreadsafe,
	monitorConfigPollChan <-chan to.TrafficMonitorConfigMap,
	localStates peer.CRStatesThreadsafe,
	statURLSubscriber chan<- poller.HttpPollerConfig,
	healthURLSubscriber chan<- poller.HttpPollerConfig,
	peerURLSubscriber chan<- poller.HttpPollerConfig,
	cachesChangeSubscriber chan<- struct{},
	cfg config.Config,
	staticAppData StaticAppData,
) {
	for monitorConfig := range monitorConfigPollChan {
		monitorConfigTS.Set(monitorConfig)
		healthURLs := map[string]poller.PollConfig{}
		statURLs := map[string]poller.PollConfig{}
		peerURLs := map[string]poller.PollConfig{}
		caches := map[string]string{}

		healthPollInterval, peerPollInterval, statPollInterval, err := getHealthPeerStatPollIntervals(monitorConfig, cfg)
		if err != nil {
			continue
		}

		for _, srv := range monitorConfig.TrafficServer {
			caches[srv.HostName] = srv.Status

			cacheName := enum.CacheName(srv.HostName)

			srvStatus := enum.CacheStatusFromString(srv.Status)
			if srvStatus == enum.CacheStatusOnline {
				localStates.SetCache(cacheName, peer.IsAvailable{IsAvailable: true})
				continue
			}
			if srvStatus == enum.CacheStatusOffline {
				continue
			}
			// seed states with available = false until our polling cycle picks up a result
			if _, exists := localStates.GetCache(cacheName); !exists {
				localStates.SetCache(cacheName, peer.IsAvailable{IsAvailable: false})
			}

			url := monitorConfig.Profile[srv.Profile].Parameters.HealthPollingURL
			if url == "" {
				log.Errorf("monitor config server %v profile %v has no polling URL; can't poll", srv.HostName, srv.Profile)
				continue
			}
			r := strings.NewReplacer(
				"${hostname}", srv.IP,
				"${interface_name}", srv.InterfaceName,
				"application=plugin.remap", "application=system",
				"application=", "application=system",
			)
			url = r.Replace(url)

			connTimeout := trafficOpsHealthConnectionTimeoutToDuration(monitorConfig.Profile[srv.Profile].Parameters.HealthConnectionTimeout)
			healthURLs[srv.HostName] = poller.PollConfig{URL: url, Timeout: connTimeout}
			r = strings.NewReplacer("application=system", "application=")
			statURL := r.Replace(url)
			statURLs[srv.HostName] = poller.PollConfig{URL: statURL, Timeout: connTimeout}
		}

		for _, srv := range monitorConfig.TrafficMonitor {
			if srv.HostName == staticAppData.Hostname {
				continue
			}
			if enum.CacheStatusFromString(srv.Status) != enum.CacheStatusOnline {
				continue
			}
			// TODO: the URL should be config driven. -jse
			url := fmt.Sprintf("http://%s:%d/publish/CrStates?raw", srv.IP, srv.Port)
			peerURLs[srv.HostName] = poller.PollConfig{URL: url} // TODO determine timeout.
		}

		statURLSubscriber <- poller.HttpPollerConfig{Urls: statURLs, Interval: statPollInterval}
		healthURLSubscriber <- poller.HttpPollerConfig{Urls: healthURLs, Interval: healthPollInterval}
		peerURLSubscriber <- poller.HttpPollerConfig{Urls: peerURLs, Interval: peerPollInterval}

		for cacheName := range localStates.GetCaches() {
			if _, exists := monitorConfig.TrafficServer[string(cacheName)]; !exists {
				log.Warnf("Removing %s from localStates", cacheName)
				localStates.DeleteCache(cacheName)
			}
		}

		cachesChangeSubscriber <- struct{}{}

		// TODO because there are multiple writers to localStates.DeliveryService, there is a race condition, where MonitorConfig (this func) and HealthResultManager could write at the same time, and the HealthResultManager could overwrite a delivery service addition or deletion here. Probably the simplest and most performant fix would be a lock-free algorithm using atomic compare-and-swaps.
		for _, ds := range monitorConfig.DeliveryService {
			// since caches default to unavailable, also default DS false
			if _, exists := localStates.GetDeliveryService(enum.DeliveryServiceName(ds.XMLID)); !exists {
				localStates.SetDeliveryService(enum.DeliveryServiceName(ds.XMLID), peer.Deliveryservice{IsAvailable: false, DisabledLocations: []enum.CacheName{}}) // important to initialize DisabledLocations, so JSON is `[]` not `null`
			}
		}
		for ds := range localStates.GetDeliveryServices() {
			if _, exists := monitorConfig.DeliveryService[string(ds)]; !exists {
				localStates.DeleteDeliveryService(ds)
			}
		}
	}
}
<|endoftext|>


<commit_before>package redisserver

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
)

type redisapi struct {
	redisconn redis.Conn
}

func (re *redisapi) Connect(network string, addredd string) {
	conn, err := redis.Dial(network, addredd)
	if err != nil {
		fmt.Println("connect error:", err.Error())
	} else {
		fmt.Println("connect redis success")
	}
	re.redisconn = conn
	//c:=NewConnection(conn)
	//conn.Close()
}
func (re *redisapi) ping() {
	err := redis.PubSubConn{}.Ping("test")
	if err != nil {
		fmt.Println(err.Error())
	}
}
func (re *redisapi) Get(args ...interface{}){
	reply,err:=re.redisconn.Do("GET",args)
	if err!=nil{
		fmt.Println(err.Error())
	}else {
		fmt.Println(reply)
	}

}
<commit_msg>add redis commad<commit_after>package redisserver

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
)

type redisapi struct {
	redisconn redis.Conn
}

func (re *redisapi) Connect(network string, addredd string) {
	conn, err := redis.Dial(network, addredd)
	if err != nil {
		fmt.Println("connect error:", err.Error())
	} else {
		fmt.Println("connect redis success")
	}
	re.redisconn = conn
	//c:=NewConnection(conn)
	//conn.Close()
}
func (re *redisapi) Ping() {
	reply,err:=re.redisconn.Do("PING","test")
	if err!=nil{
		fmt.Println(err.Error())
	}else {
		fmt.Println(redis.String(reply,err))
	}
}
func (re *redisapi)Set(args ...interface{}){
	reply,err:=re.redisconn.Do("SET",args)
	redis.String(reply,err)
	if err!=nil {
		fmt.Println(err.Error())
	}else {
		fmt.Println(reply)
	}
}

func (re *redisapi) Get(args interface{}){
	reply,err:=re.redisconn.Do("GET",args)
	if err!=nil{
		fmt.Println(err.Error())
	}else {
		fmt.Println(reply)
	}

}

func (re *redisapi) Delete(args ...interface{}){
	reply,err:=re.redisconn.Do("DEL",args)
	if err!=nil{
		fmt.Println(err.Error())
	}else {
		fmt.Println(reply)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"net/http/httputil"

	"github.com/getlantern/flashlight/client"
)

var (
	// https://github.com/getlantern/lantern/issues/3147
	redirectSites = []string{
		// sites that redirect to http when accessed through https
		"http://www.bbc.co.uk",
		"http://www.speedtest.net",
		// sites that redirect to itself if the "Host:" HTTP header contains port
		"http://lowendbox.com",
		"http://sourceforge.net",
		"http://krypted.com/mac-security/manage-profiles-from-the-command-line-in-os-x-10-9/",
	}

	atsDefaultBody = []byte(`

<TITLE>Not Found on Accelerator</TITLE>

Not Found on Accelerator

Description: Your request on the specified host was not found.
Check the location and try again.

`)
)

func verifyFallback(fb *client.ChainedServerInfo, c *http.Client) {
	verifyMimicWhenNoAuthToken(fb, c)
	verifyMimicWithInvalidAuthToken(fb, c)
	for _, s := range redirectSites {
		verifyRedirectSites(fb, c, s)
	}
}

func verifyMimicWhenNoAuthToken(fb *client.ChainedServerInfo, c *http.Client) {
	req, err := http.NewRequest("GET", "http://www.google.com/humans.txt", nil)
	if err != nil {
		log.Errorf("%v: NewRequest() error : %v", fb.Addr, err)
		return
	}
	// intentionally not set auth token
	doVerifyMimic(fb.Addr, req, c)
}

func verifyMimicWithInvalidAuthToken(fb *client.ChainedServerInfo, c *http.Client) {
	req, err := http.NewRequest("GET", "http://www.google.com/humans.txt", nil)
	if err != nil {
		log.Errorf("%v: NewRequest() error : %v", fb.Addr, err)
		return
	}
	req.Header.Set("X-LANTERN-AUTH-TOKEN", "invalid")
	doVerifyMimic(fb.Addr, req, c)
}

func doVerifyMimic(addr string, req *http.Request, c *http.Client) {
	resp, err := c.Do(req)
	if err != nil {
		log.Errorf("%v: requesting humans.txt failed: %v", addr, err)
		if *verbose {
			reqStr, _ := httputil.DumpRequestOut(req, true)
			log.Debug(string(reqStr))
		}
		return
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			log.Debugf("Unable to close response body: %v", err)
		}
	}()
	if resp.StatusCode != 404 {
		log.Errorf("%v: should get 404 if auth failed", addr)
		return
	}
	if resp.Header.Get("Content-Type") != "text/html" ||
		resp.Header.Get("Cache-Control") != "no-store" ||
		resp.Header.Get("Connection") != "keep-alive" ||
		resp.Header.Get("Content-Language") != "en" ||
		resp.Header.Get("Content-Length") != "297" {
		log.Errorf("%v: should have correct headers present", addr)
		return
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Errorf("%v: error reading response body: %v", addr, err)
		return
	}
	if !bytes.Equal(b, atsDefaultBody) {
		log.Errorf("%v: response body mismatch for invalid request", addr)
		log.Debugf("Body expected: %v", string(atsDefaultBody))
		log.Debugf("Body got: %v", string(b))
		return
	}
	log.Debugf("%v: OK.", addr)
}

func verifyRedirectSites(fb *client.ChainedServerInfo, c *http.Client, url string) {
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("X-LANTERN-AUTH-TOKEN", fb.AuthToken)
	resp, err := c.Do(req)
	if err != nil {
		log.Errorf("%v: requesting %s failed: %v", fb.Addr, url, err)
		if *verbose {
			reqStr, _ := httputil.DumpRequestOut(req, true)
			log.Debug(string(reqStr))
		}
		return
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			log.Debugf("Unable to close response body: %v", err)
		}
	}()
	if resp.StatusCode != 200 {
		log.Errorf("%v: bad status code %v for %s", fb.Addr, resp.StatusCode, url)
		if *verbose {
			respStr, _ := httputil.DumpResponse(resp, true)
			log.Debug(string(respStr))
		}
		return
	}
	log.Debugf("%v via %s: OK.", fb.Addr, url)
}
<commit_msg>code review updates<commit_after>package main

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"net/http/httputil"

	"github.com/getlantern/flashlight/client"
)

var (
	// https://github.com/getlantern/lantern/issues/3147
	redirectSites = []string{
		// sites that redirect to http when accessed through https
		"http://www.bbc.co.uk",
		"http://www.speedtest.net",
		// sites that redirect to itself if the "Host:" HTTP header contains port
		"http://lowendbox.com",
		"http://sourceforge.net",
		"http://krypted.com/mac-security/manage-profiles-from-the-command-line-in-os-x-10-9/",
	}

	atsDefaultBody = []byte(`

<TITLE>Not Found on Accelerator</TITLE>

Not Found on Accelerator

Description: Your request on the specified host was not found.
Check the location and try again.

`)
)

func verifyFallback(fb *client.ChainedServerInfo, c *http.Client) {
	verifyMimicWhenNoAuthToken(fb, c)
	verifyMimicWithInvalidAuthToken(fb, c)
	for _, s := range redirectSites {
		verifyRedirectSites(fb, c, s)
	}
}

func verifyMimicWhenNoAuthToken(fb *client.ChainedServerInfo, c *http.Client) {
	req, err := http.NewRequest("GET", "http://www.google.com/humans.txt", nil)
	if err != nil {
		log.Errorf("%v: NewRequest() error : %v", fb.Addr, err)
		return
	}
	// intentionally not set auth token
	doVerifyMimic(fb.Addr, req, c)
}

func verifyMimicWithInvalidAuthToken(fb *client.ChainedServerInfo, c *http.Client) {
	req, err := http.NewRequest("GET", "http://www.google.com/humans.txt", nil)
	if err != nil {
		log.Errorf("%v: NewRequest() error : %v", fb.Addr, err)
		return
	}
	req.Header.Set("X-LANTERN-AUTH-TOKEN", "invalid")
	doVerifyMimic(fb.Addr, req, c)
}

func doVerifyMimic(addr string, req *http.Request, c *http.Client) {
	resp, err := c.Do(req)
	if err != nil {
		log.Errorf("%v: requesting humans.txt failed: %v", addr, err)
		return
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			log.Debugf("Unable to close response body: %v", err)
		}
	}()
	if resp.StatusCode != 404 {
		log.Errorf("%v: should get 404 if auth failed", addr)
		if *verbose {
			respStr, _ := httputil.DumpResponse(resp, true)
			log.Debug(string(respStr))
		}
		return
	}
	if resp.Header.Get("Content-Type") != "text/html" ||
		resp.Header.Get("Cache-Control") != "no-store" ||
		resp.Header.Get("Connection") != "keep-alive" ||
		resp.Header.Get("Content-Language") != "en" ||
		resp.Header.Get("Content-Length") != "297" {
		log.Errorf("%v: should have correct headers present", addr)
		if *verbose {
			respStr, _ := httputil.DumpResponse(resp, true)
			log.Debug(string(respStr))
		}
		return
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Errorf("%v: error reading response body: %v", addr, err)
		return
	}
	if !bytes.Equal(b, atsDefaultBody) {
		log.Errorf("%v: response body mismatch for invalid request", addr)
		log.Debugf("Body expected: %v", string(atsDefaultBody))
		log.Debugf("Body got: %v", string(b))
		return
	}
	log.Debugf("%v: OK.", addr)
}

func verifyRedirectSites(fb *client.ChainedServerInfo, c *http.Client, url string) {
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("X-LANTERN-AUTH-TOKEN", fb.AuthToken)
	if err != nil {
		log.Errorf("error make request to %s: %v", url, err)
		return
	}
	resp, err := c.Do(req)
	if err != nil {
		log.Errorf("%v: requesting %s failed: %v", fb.Addr, url, err)
		if *verbose {
			reqStr, _ := httputil.DumpRequestOut(req, true)
			log.Debug(string(reqStr))
		}
		return
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			log.Debugf("Unable to close response body: %v", err)
		}
	}()
	if resp.StatusCode != 200 {
		log.Errorf("%v: bad status code %v for %s", fb.Addr, resp.StatusCode, url)
		if *verbose {
			respStr, _ := httputil.DumpResponse(resp, true)
			log.Debug(string(respStr))
		}
		return
	}
	log.Debugf("%v via %s: OK.", fb.Addr, url)
}
<|endoftext|>


<commit_before>/*
Given a fixed length array arr of integers, duplicate each occurrence of zero, shifting the remaining elements to the right.

Note that elements beyond the length of the original array are not written.

Do the above modifications to the input array in place, do not return anything from your function.

Example 1:

Input: [1,0,2,3,0,4,5,0]
Output: null
Explanation: After calling your function, the input array is modified to: [1,0,0,2,3,0,0,4]
Example 2:

Input: [1,2,3]
Output: null
Explanation: After calling your function, the input array is modified to: [1,2,3]
*/

package main

import (
	"log"
)

func main() {
	tests := [][]int{{1,0,2,3,0,4,5,0}, {1,2,3}}

	for testNo, test := range tests {
		log.Printf("#%d duplicateZeros(%v) == ", testNo, test)
		duplicateZeros(test)
		log.Printf("%v\n", test)
	}
}

func duplicateZeros(x []int)  {
	for i := 0; i < len(x); i++ {
		if x[i] == 0 {
			for j := len(x)-2; j > i; j-- {
				x[j]= x[j+1]
			}

			if i+1 < len(x) {
				x[i+1] = 0
				i++
			}
		}
	}
}
<commit_msg>Fix upshift bug<commit_after>/*
Given a fixed length array arr of integers, duplicate each occurrence of zero, shifting the remaining elements to the right.

Note that elements beyond the length of the original array are not written.

Do the above modifications to the input array in place, do not return anything from your function.

Example 1:

Input: [1,0,2,3,0,4,5,0]
Output: null
Explanation: After calling your function, the input array is modified to: [1,0,0,2,3,0,0,4]
Example 2:

Input: [1,2,3]
Output: null
Explanation: After calling your function, the input array is modified to: [1,2,3]
*/

package main

import (
	"log"
)

func main() {
	tests := [][]int{{1, 0, 2, 3}, {1, 0, 2, 3, 0, 4, 5, 0}, {1, 2, 3}}

	for testNo, test := range tests {
		log.Printf("#%d duplicateZeros(%v) == ", testNo, test)
		duplicateZeros(test)
		log.Printf("%v\n", test)
	}
}

func duplicateZeros(x []int) {
	for i := 0; i < len(x); i++ {
		if x[i] == 0 {
			// found a zero

			// shift everything up one position
			for j := len(x) - 2; j > i; j-- {
				x[j+1] = x[j]
			}

			// set next element to zero, and perform extra increment
			if i+1 < len(x) {
				x[i+1] = 0
				i++
			}
		}
	}
}
<|endoftext|>


<commit_before>package scaling_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"lib/models"
	"lib/policy_client"
	"lib/testsupport"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	pusherConfig "cf-pusher/config"

	"code.cloudfoundry.org/lager/lagertest"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-cf-experimental/rainmaker"
)

const Timeout_Check = 20 * time.Minute

// 2 * poll cycle time (5s)
const Policy_Update_Wait = 10 * time.Second

var _ = Describe("how the container network performs at scale", func() {
	Describe("scaling tests", func() {
		var (
			proxyApps   []string
			tickApps    []string
			registryApp string
			ports       []int
			testConfig  pusherConfig.Config
		)
		BeforeEach(func() {
			testConfig = pushConfig
			registryApp = pushConfig.Prefix + "registry"
			By("checking that destination app instances have registered themselves")
			checkRegistry(registryApp, 10*time.Second, 500*time.Millisecond, pushConfig.Applications*pushConfig.AppInstances)
		})
		JustBeforeEach(func() {
			proxyApps = []string{}
			for i := 0; i < testConfig.ProxyApplications; i++ {
				proxyApps = append(proxyApps, fmt.Sprintf(testConfig.Prefix+"proxy-%d", i))
			}

			tickApps = []string{}
			for i := 0; i < testConfig.Applications; i++ {
				tickApps = append(tickApps, fmt.Sprintf(testConfig.Prefix+"tick-%d", i))
			}

			ports = []int{8080}
			for i := 0; i < testConfig.ExtraListenPorts; i++ {
				ports = append(ports, 7000+i)
			}
		})
		runScalingTest := func() {
			It("allows the user to configure policies", func(done Done) {
				By(fmt.Sprintf("testing with %d source apps and %d destination apps", testConfig.ProxyApplications, testConfig.Applications))
				appIPs := getAppIPs(registryApp, tickApps)
				sample := sampleIPs(appIPs, testConfig.SampleSize)

				By(fmt.Sprintf("checking that the connection fails sampling %d out of %d IPs on %d ports", len(sample), len(appIPs), len(ports)))
				runWithTimeout("check connection failures", Timeout_Check, func() {
					assertConnectionFails(proxyApps, sample, ports, testConfig.ProxyInstances)
				})

				By(fmt.Sprintf("creating %d policies", len(proxyApps)*len(tickApps)*len(ports)))
				for _, proxyApp := range proxyApps {
					doAllPolicies("create", proxyApp, tickApps, ports)
				}

				By(fmt.Sprintf("waiting %s for policies to be updated on cells", Policy_Update_Wait))
				time.Sleep(Policy_Update_Wait)

				sample = sampleIPs(appIPs, testConfig.SampleSize)
				By(fmt.Sprintf("checking that the connection succeeds sampling %d out of %d IPs on %d ports to proxy", len(sample), len(appIPs), len(ports)))
				runWithTimeout("check connection success", Timeout_Check, func() {
					assertConnectionSucceeds(proxyApps, sample, ports, testConfig.ProxyInstances)
				})

				for _, proxyApp := range proxyApps {
					By(fmt.Sprintf("dumping stats to commit to stats repo for %s", proxyApp))
					dumpStats(proxyApp, testConfig.AppsDomain)
				}

				By("sleeping for 30 seconds while policies exist")
				time.Sleep(30 * time.Second)

				By(fmt.Sprintf("deleting %d policies", len(proxyApps)*len(tickApps)*len(ports)))
				for _, proxyApp := range proxyApps {
					doAllPolicies("delete", proxyApp, tickApps, ports)
				}

				By(fmt.Sprintf("waiting %s for policies to be updated on cells", Policy_Update_Wait))
				time.Sleep(Policy_Update_Wait)

				sample = sampleIPs(appIPs, testConfig.SampleSize)
				By(fmt.Sprintf("checking that the connection fails sampling %d out of %d IPs on %d ports", len(sample), len(appIPs), len(ports)))
				runWithTimeout("check connection failures, again", Timeout_Check, func() {
					assertConnectionFails(proxyApps, sample, ports, testConfig.ProxyInstances)
				})
				close(done)
			}, 30*60) // 30 minutes
		}
		Context("when one client with many backends", func() {
			BeforeEach(func() {
				testConfig.ProxyApplications = 1
			})
			runScalingTest()
		})
		Context("when one server with many clients", func() {
			BeforeEach(func() {
				testConfig.Applications = 1
			})
			runScalingTest()
		})
	})

	Describe("sampleIPs", func() {
		var population []string
		BeforeEach(func() {
			population = []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}
		})

		It("returns a sample of unique choices from the population", func() {
			sample := sampleIPs(population, 9)
			Expect(len(sample)).To(Equal(9))
			for i := 0; i < len(sample); i++ {
				for j := i + 1; j < len(sample); j++ {
					Expect(sample[i]).NotTo(Equal(sample[j]))
				}
			}
		})
		Context("when the sample size is larger than the population", func() {
			It("returns the whole population", func() {
				sample := sampleIPs(population, 999)
				Expect(sample).To(Equal(population))
			})
		})
		Context("when the sample size is equal to the population size", func() {
			It("returns the whole population", func() {
				sample := sampleIPs(population, len(population))
				Expect(sample).To(Equal(population))
			})
		})
		Context("when the sample size is zero", func() {
			It("returns the whole population", func() {
				sample := sampleIPs(population, 0)
				Expect(sample).To(Equal(population))
			})
		})
		Context("when the sample size is negative", func() {
			It("returns the whole population", func() {
				sample := sampleIPs(population, -1)
				Expect(sample).To(Equal(population))
			})
		})
	})
})

func getToken() string {
	cmd := exec.Command("cf", "oauth-token")
	session, err := gexec.Start(cmd, nil, nil)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session.Wait(Timeout_Short)).Should(gexec.Exit(0))
	rawOutput := string(session.Out.Contents())
	return strings.TrimSpace(strings.TrimPrefix(rawOutput, "bearer "))
}

func getGuids(sourceAppName string, dstAppNames []string) (string, []string) {
	dstGuids := []string{}
	sourceGuid := ""
	token := getToken()
	appsClient := rainmaker.NewApplicationsService(rainmaker.Config{Host: "http://" + config.ApiEndpoint})

	appsList, err := appsClient.List(token)
	Expect(err).NotTo(HaveOccurred())

	for {
		for _, app := range appsList.Applications {
			if app.Name == sourceAppName {
				sourceGuid = app.GUID
				continue
			}
			for _, tickAppName := range dstAppNames {
				if app.Name == tickAppName {
					dstGuids = append(dstGuids, app.GUID)
					break
				}
			}
		}
		if appsList.HasNextPage() {
			appsList, err = appsList.Next(token)
			Expect(err).NotTo(HaveOccurred())
		} else {
			break
		}
	}

	Expect(sourceGuid).NotTo(BeEmpty())
	Expect(dstGuids).To(HaveLen(len(dstAppNames)))

	return sourceGuid, dstGuids
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func doAllPolicies(action string, source string, dstList []string, dstPorts []int) {
	policyClient := policy_client.NewExternal(lagertest.NewTestLogger("test"), &http.Client{}, "http://"+config.ApiEndpoint)
	sourceGuid, dstGuids := getGuids(source, dstList)
	policies := []models.Policy{}
	for _, dstGuid := range dstGuids {
		for _, port := range dstPorts {
			policies = append(policies, models.Policy{
				Source: models.Source{
					ID: sourceGuid,
				},
				Destination: models.Destination{
					ID:       dstGuid,
					Port:     port,
					Protocol: "tcp",
				},
			})
		}
	}
	token := getToken()
	if action == "create" {
		Expect(policyClient.AddPolicies(token, policies)).To(Succeed())
	} else if action == "delete" {
		for i := 0; i < len(policies); i += 100 {
			Expect(policyClient.DeletePolicies(token, policies[i:min(i+100, len(policies))])).To(Succeed())
		}
	}
}

func runWithTimeout(operation string, timeout time.Duration, work func()) {
	done := make(chan bool)
	go func() {
		defer func() { close(done) }()
		defer GinkgoRecover()

		By(fmt.Sprintf("starting %s\n", operation))
		work()
		By(fmt.Sprintf("completed %s\n", operation))
		done <- true
	}()

	select {
	case ok := <-done:
		if !ok {
			Fail("failure during " + operation)
		}
	case <-time.After(timeout):
		Fail("timeout on " + operation)
	}
}

func dumpStats(host, domain string) {
	resp, err := httpGetBytes(fmt.Sprintf("http://%s.%s/stats", host, domain))
	Expect(err).NotTo(HaveOccurred())

	netStatsFile := os.Getenv("NETWORK_STATS_FILE")
	if netStatsFile != "" {
		Expect(ioutil.WriteFile(netStatsFile, resp.Body, 0600)).To(Succeed())
	}
}

type RegistryInstancesResponse struct {
	Instances []struct {
		ServiceName string `json:"service_name"`
		Endpoint    struct {
			Value string `json:"value"`
		} `json:"endpoint"`
	} `json:"instances"`
}

func getInstancesFromA8(registry string) (*RegistryInstancesResponse, error) {
	resp, err := httpGetBytes(fmt.Sprintf("http://%s.%s/api/v1/instances", registry, config.AppsDomain))
	if err != nil {
		return nil, err
	}

	Expect(resp.StatusCode).To(Equal(http.StatusOK))

	var instancesResponse RegistryInstancesResponse
	err = json.Unmarshal(resp.Body, &instancesResponse)
	if err != nil {
		return nil, err
	}
	return &instancesResponse, nil
}

func checkRegistry(registry string, timeout, pollingInterval time.Duration, totalInstances int) {
	registeredApps := func() (int, error) {
		instancesResponse, err := getInstancesFromA8(registry)
		if err != nil {
			return 0, err
		}
		return len(instancesResponse.Instances), nil
	}

	Eventually(registeredApps, timeout, pollingInterval).Should(Equal(totalInstances))
}

func getAppIPs(registry string, appNames []string) []string {
	instancesResponse, err := getInstancesFromA8(registry)
	Expect(err).NotTo(HaveOccurred())

	ips := []string{}
	for _, instance := range instancesResponse.Instances {
		for _, name := range appNames {
			if name == instance.ServiceName {
				ip, _, err := net.SplitHostPort(instance.Endpoint.Value)
				Expect(err).NotTo(HaveOccurred())
				ips = append(ips, ip)
			}
		}
	}
	return ips
}

func sampleIPs(population []string, sampleSize int) []string {
	populationSize := len(population)
	if len(population) <= sampleSize || sampleSize < 1 {
		return population
	}
	var sample = []string{}
	for i := 0; i < sampleSize; i++ {
		j := rand.Intn(populationSize)
		sample = append(sample, population[j])
		population = append(population[:j], population[j+1:]...)
		populationSize--
	}
	return sample
}

type SrcDstPair struct {
	Source string
	Dest   string
}

func assertConnectionSucceeds(sourceApps, destApps []string, ports []int, nProxies int) {
	parallelRunner := &testsupport.ParallelRunner{
		NumWorkers: 50 * nProxies,
	}
	pairs := []interface{}{}
	for _, s := range sourceApps {
		for _, d := range destApps {
			pairs = append(pairs, SrcDstPair{Source: s, Dest: d})
		}
	}
	parallelRunner.RunOnSlice(pairs, func(obj interface{}) {
		pair := obj.(SrcDstPair)
		for _, port := range ports {
			assertSingleConnection(pair.Dest, port, pair.Source, true)
		}
	})
}

func assertConnectionFails(sourceApps, destApps []string, ports []int, nProxies int) {
	parallelRunner := &testsupport.ParallelRunner{
		NumWorkers: 50 * nProxies,
	}
	pairs := []interface{}{}
	for _, s := range sourceApps {
		for _, d := range destApps {
			pairs = append(pairs, SrcDstPair{Source: s, Dest: d})
		}
	}
	parallelRunner.RunOnSlice(pairs, func(obj interface{}) {
		pair := obj.(SrcDstPair)
		for _, port := range ports {
			assertSingleConnection(pair.Dest, port, pair.Source, false)
		}
	})
}

func assertSingleConnection(destIP string, port int, sourceAppName string, shouldSucceed bool) {
	if shouldSucceed {
		assertResponseContains(destIP, port, sourceAppName, "application_name")
	} else {
		assertResponseContains(destIP, port, sourceAppName, "request failed")
	}
}

func assertResponseContains(destIP string, port int, sourceAppName string, desiredResponse string) {
	proxyTest := func() (string, error) {
		resp, err := httpGetBytes(fmt.Sprintf("http://%s.%s/proxy/%s:%d", sourceAppName, config.AppsDomain, destIP, port))
		if err != nil {
			return "", err
		}
		return string(resp.Body), nil
	}
	Eventually(proxyTest, 10*time.Second, 500*time.Millisecond).Should(ContainSubstring(desiredResponse))
}

var httpClient = &http.Client{
	Transport: &http.Transport{
		DisableKeepAlives: true,
		Dial: (&net.Dialer{
			Timeout:   4 * time.Second,
			KeepAlive: 0,
		}).Dial,
	},
}

type httpResp struct {
	StatusCode int
	Body       []byte
}

func httpGetBytes(url string) (httpResp, error) {
	resp, err := httpClient.Get(url)
	if err != nil {
		return httpResp{}, err
	}
	defer resp.Body.Close()

	respBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return httpResp{}, err
	}

	return httpResp{resp.StatusCode, respBytes}, nil
}
<commit_msg>scaling test prints timestamp in each step<commit_after>package scaling_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"lib/models"
	"lib/policy_client"
	"lib/testsupport"
	"math/rand"
	"net"
	"net/http"
	"os/exec"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"

	pusherConfig "cf-pusher/config"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-cf-experimental/rainmaker"
)

const Timeout_Check = 20 * time.Minute

// 2 * poll cycle time (5s)
const Policy_Update_Wait = 10 * time.Second

const Time_Format = "15:04:05"

var _ = Describe("how the container network performs at scale", func() {
	Describe("scaling tests", func() {
		var (
			proxyApps   []string
			tickApps    []string
			registryApp string
			ports       []int
			testConfig  pusherConfig.Config
		)
		BeforeEach(func() {
			testConfig = pushConfig
			registryApp = pushConfig.Prefix + "registry"
			By(fmt.Sprintf("%s checking that destination app instances have registered themselves", ts()))
			checkRegistry(registryApp, 10*time.Second, 500*time.Millisecond, pushConfig.Applications*pushConfig.AppInstances)
		})
		JustBeforeEach(func() {
			proxyApps = []string{}
			for i := 0; i < testConfig.ProxyApplications; i++ {
				proxyApps = append(proxyApps, fmt.Sprintf(testConfig.Prefix+"proxy-%d", i))
			}

			tickApps = []string{}
			for i := 0; i < testConfig.Applications; i++ {
				tickApps = append(tickApps, fmt.Sprintf(testConfig.Prefix+"tick-%d", i))
			}

			ports = []int{8080}
			for i := 0; i < testConfig.ExtraListenPorts; i++ {
				ports = append(ports, 7000+i)
			}
		})
		runScalingTest := func() {
			It("allows the user to configure policies", func(done Done) {
				By(fmt.Sprintf("%s testing with %d source apps and %d destination apps", ts(), testConfig.ProxyApplications, testConfig.Applications))
				appIPs := getAppIPs(registryApp, tickApps)
				sample := sampleIPs(appIPs, testConfig.SampleSize)

				By(fmt.Sprintf("%s checking that the connection fails sampling %d out of %d IPs on %d ports", ts(), len(sample), len(appIPs), len(ports)))
				runWithTimeout("check connection failures", Timeout_Check, func() {
					assertConnectionFails(proxyApps, sample, ports, testConfig.ProxyInstances)
				})

				By(fmt.Sprintf("%s creating %d policies", ts(), len(proxyApps)*len(tickApps)*len(ports)))
				for _, proxyApp := range proxyApps {
					doAllPolicies("create", proxyApp, tickApps, ports)
				}

				By(fmt.Sprintf("%s waiting %s for policies to be updated on cells", ts(), Policy_Update_Wait))
				time.Sleep(Policy_Update_Wait)

				sample = sampleIPs(appIPs, testConfig.SampleSize)
				By(fmt.Sprintf("%s checking that the connection succeeds sampling %d out of %d IPs on %d ports to proxy", ts(), len(sample), len(appIPs), len(ports)))
				runWithTimeout("check connection success", Timeout_Check, func() {
					assertConnectionSucceeds(proxyApps, sample, ports, testConfig.ProxyInstances)
				})

				By(fmt.Sprintf("%s sleeping for 30 seconds while policies exist", ts()))
				time.Sleep(30 * time.Second)

				By(fmt.Sprintf("%s deleting %d policies", ts(), len(proxyApps)*len(tickApps)*len(ports)))
				for _, proxyApp := range proxyApps {
					doAllPolicies("delete", proxyApp, tickApps, ports)
				}

				By(fmt.Sprintf("%s waiting %s for policies to be updated on cells", ts(), Policy_Update_Wait))
				time.Sleep(Policy_Update_Wait)

				sample = sampleIPs(appIPs, testConfig.SampleSize)
				By(fmt.Sprintf("%s checking that the connection fails sampling %d out of %d IPs on %d ports", ts(), len(sample), len(appIPs), len(ports)))
				runWithTimeout("check connection failures, again", Timeout_Check, func() {
					assertConnectionFails(proxyApps, sample, ports, testConfig.ProxyInstances)
				})
				close(done)
			}, 30*60) // 30 minutes
		}
		Context("when one client with many backends", func() {
			BeforeEach(func() {
				testConfig.ProxyApplications = 1
			})
			runScalingTest()
		})
		Context("when one server with many clients", func() {
			BeforeEach(func() {
				testConfig.Applications = 1
			})
			runScalingTest()
		})
	})

	Describe("sampleIPs", func() {
		var population []string
		BeforeEach(func() {
			population = []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}
		})

		It("returns a sample of unique choices from the population", func() {
			sample := sampleIPs(population, 9)
			Expect(len(sample)).To(Equal(9))
			for i := 0; i < len(sample); i++ {
				for j := i + 1; j < len(sample); j++ {
					Expect(sample[i]).NotTo(Equal(sample[j]))
				}
			}
		})
		Context("when the sample size is larger than the population", func() {
			It("returns the whole population", func() {
				sample := sampleIPs(population, 999)
				Expect(sample).To(Equal(population))
			})
		})
		Context("when the sample size is equal to the population size", func() {
			It("returns the whole population", func() {
				sample := sampleIPs(population, len(population))
				Expect(sample).To(Equal(population))
			})
		})
		Context("when the sample size is zero", func() {
			It("returns the whole population", func() {
				sample := sampleIPs(population, 0)
				Expect(sample).To(Equal(population))
			})
		})
		Context("when the sample size is negative", func() {
			It("returns the whole population", func() {
				sample := sampleIPs(population, -1)
				Expect(sample).To(Equal(population))
			})
		})
	})
})

func getToken() string {
	cmd := exec.Command("cf", "oauth-token")
	session, err := gexec.Start(cmd, nil, nil)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session.Wait(Timeout_Short)).Should(gexec.Exit(0))
	rawOutput := string(session.Out.Contents())
	return strings.TrimSpace(strings.TrimPrefix(rawOutput, "bearer "))
}

func getGuids(sourceAppName string, dstAppNames []string) (string, []string) {
	dstGuids := []string{}
	sourceGuid := ""
	token := getToken()
	appsClient := rainmaker.NewApplicationsService(rainmaker.Config{Host: "http://" + config.ApiEndpoint})

	appsList, err := appsClient.List(token)
	Expect(err).NotTo(HaveOccurred())

	for {
		for _, app := range appsList.Applications {
			if app.Name == sourceAppName {
				sourceGuid = app.GUID
				continue
			}
			for _, tickAppName := range dstAppNames {
				if app.Name == tickAppName {
					dstGuids = append(dstGuids, app.GUID)
					break
				}
			}
		}
		if appsList.HasNextPage() {
			appsList, err = appsList.Next(token)
			Expect(err).NotTo(HaveOccurred())
		} else {
			break
		}
	}

	Expect(sourceGuid).NotTo(BeEmpty())
	Expect(dstGuids).To(HaveLen(len(dstAppNames)))

	return sourceGuid, dstGuids
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func doAllPolicies(action string, source string, dstList []string, dstPorts []int) {

	logger := lager.NewLogger("test")
	logger.RegisterSink(lager.NewWriterSink(GinkgoWriter, lager.INFO))

	policyClient := policy_client.NewExternal(logger, &http.Client{}, "http://"+config.ApiEndpoint)
	sourceGuid, dstGuids := getGuids(source, dstList)
	policies := []models.Policy{}
	for _, dstGuid := range dstGuids {
		for _, port := range dstPorts {
			policies = append(policies, models.Policy{
				Source: models.Source{
					ID: sourceGuid,
				},
				Destination: models.Destination{
					ID:       dstGuid,
					Port:     port,
					Protocol: "tcp",
				},
			})
		}
	}
	token := getToken()
	if action == "create" {
		Expect(policyClient.AddPolicies(token, policies)).To(Succeed())
	} else if action == "delete" {
		for i := 0; i < len(policies); i += 100 {
			Expect(policyClient.DeletePolicies(token, policies[i:min(i+100, len(policies))])).To(Succeed())
		}
	}
}

func runWithTimeout(operation string, timeout time.Duration, work func()) {
	done := make(chan bool)
	go func() {
		defer func() { close(done) }()
		defer GinkgoRecover()

		By(fmt.Sprintf("%s starting %s\n", ts(), operation))
		work()
		By(fmt.Sprintf("%s completed %s\n", ts(), operation))
		done <- true
	}()

	select {
	case ok := <-done:
		if !ok {
			Fail("failure during " + operation)
		}
	case <-time.After(timeout):
		Fail("timeout on " + operation)
	}
}

type RegistryInstancesResponse struct {
	Instances []struct {
		ServiceName string `json:"service_name"`
		Endpoint    struct {
			Value string `json:"value"`
		} `json:"endpoint"`
	} `json:"instances"`
}

func getInstancesFromA8(registry string) (*RegistryInstancesResponse, error) {
	resp, err := httpGetBytes(fmt.Sprintf("http://%s.%s/api/v1/instances", registry, config.AppsDomain))
	if err != nil {
		return nil, err
	}

	Expect(resp.StatusCode).To(Equal(http.StatusOK))

	var instancesResponse RegistryInstancesResponse
	err = json.Unmarshal(resp.Body, &instancesResponse)
	if err != nil {
		return nil, err
	}
	return &instancesResponse, nil
}

func checkRegistry(registry string, timeout, pollingInterval time.Duration, totalInstances int) {
	registeredApps := func() (int, error) {
		instancesResponse, err := getInstancesFromA8(registry)
		if err != nil {
			return 0, err
		}
		return len(instancesResponse.Instances), nil
	}

	Eventually(registeredApps, timeout, pollingInterval).Should(Equal(totalInstances))
}

func getAppIPs(registry string, appNames []string) []string {
	instancesResponse, err := getInstancesFromA8(registry)
	Expect(err).NotTo(HaveOccurred())

	ips := []string{}
	for _, instance := range instancesResponse.Instances {
		for _, name := range appNames {
			if name == instance.ServiceName {
				ip, _, err := net.SplitHostPort(instance.Endpoint.Value)
				Expect(err).NotTo(HaveOccurred())
				ips = append(ips, ip)
			}
		}
	}
	return ips
}

func sampleIPs(population []string, sampleSize int) []string {
	populationSize := len(population)
	if len(population) <= sampleSize || sampleSize < 1 {
		return population
	}
	var sample = []string{}
	for i := 0; i < sampleSize; i++ {
		j := rand.Intn(populationSize)
		sample = append(sample, population[j])
		population = append(population[:j], population[j+1:]...)
		populationSize--
	}
	return sample
}

type SrcDstPair struct {
	Source string
	Dest   string
}

func assertConnectionSucceeds(sourceApps, destApps []string, ports []int, nProxies int) {
	parallelRunner := &testsupport.ParallelRunner{
		NumWorkers: 50 * nProxies,
	}
	pairs := []interface{}{}
	for _, s := range sourceApps {
		for _, d := range destApps {
			pairs = append(pairs, SrcDstPair{Source: s, Dest: d})
		}
	}
	parallelRunner.RunOnSlice(pairs, func(obj interface{}) {
		pair := obj.(SrcDstPair)
		for _, port := range ports {
			assertSingleConnection(pair.Dest, port, pair.Source, true)
		}
	})
}

func assertConnectionFails(sourceApps, destApps []string, ports []int, nProxies int) {
	parallelRunner := &testsupport.ParallelRunner{
		NumWorkers: 50 * nProxies,
	}
	pairs := []interface{}{}
	for _, s := range sourceApps {
		for _, d := range destApps {
			pairs = append(pairs, SrcDstPair{Source: s, Dest: d})
		}
	}
	parallelRunner.RunOnSlice(pairs, func(obj interface{}) {
		pair := obj.(SrcDstPair)
		for _, port := range ports {
			assertSingleConnection(pair.Dest, port, pair.Source, false)
		}
	})
}

func assertSingleConnection(destIP string, port int, sourceAppName string, shouldSucceed bool) {
	if shouldSucceed {
		assertResponseContains(destIP, port, sourceAppName, "application_name")
	} else {
		assertResponseContains(destIP, port, sourceAppName, "request failed")
	}
}

func assertResponseContains(destIP string, port int, sourceAppName string, desiredResponse string) {
	proxyTest := func() (string, error) {
		resp, err := httpGetBytes(fmt.Sprintf("http://%s.%s/proxy/%s:%d", sourceAppName, config.AppsDomain, destIP, port))
		if err != nil {
			return "", err
		}
		return string(resp.Body), nil
	}
	Eventually(proxyTest, 10*time.Second, 500*time.Millisecond).Should(ContainSubstring(desiredResponse))
}

var httpClient = &http.Client{
	Transport: &http.Transport{
		DisableKeepAlives: true,
		Dial: (&net.Dialer{
			Timeout:   4 * time.Second,
			KeepAlive: 0,
		}).Dial,
	},
}

type httpResp struct {
	StatusCode int
	Body       []byte
}

func httpGetBytes(url string) (httpResp, error) {
	resp, err := httpClient.Get(url)
	if err != nil {
		return httpResp{}, err
	}
	defer resp.Body.Close()

	respBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return httpResp{}, err
	}

	return httpResp{resp.StatusCode, respBytes}, nil
}

func ts() string {
	return time.Now().Format(Time_Format)
}
<|endoftext|>