<commit_before>package dondeestas

import (
	"bytes"
	"encoding/json"
	"log"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
)

func getExpectedPersonRequest() (*http.Request, *Person, *httptest.ResponseRecorder) {
	expectedPerson, _ := createRandomPerson()
	expectedPersonJson, _ := json.Marshal(expectedPerson)
	expectedPersonStr := string(expectedPersonJson)
	req := httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))

	return req, expectedPerson, httptest.NewRecorder()
}

func TestRouting_GetJson(t *testing.T) {
	expectedPerson, _ := createRandomPerson()
	expectedPersonJson, _ := json.Marshal(expectedPerson)
	expectedPersonStr := string(expectedPersonJson)
	req := httptest.NewRequest("GET", "http://blah.com/foo", bytes.NewBufferString(expectedPersonStr))

	var person Person
	if err := getJson(req.Body, person); err != nil {
		t.Fatalf("Encountered error when retrieving json from string: %s", err)
	}

	// Test forcing the function to read a closed stream
	req = httptest.NewRequest("GET", "http://blah.com/foo", bytes.NewBufferString(expectedPersonStr))
	if err := req.Body.Close(); err != nil {
		t.Fatalf("Could not close test request body!")
	}
	if err := getJson(req.Body, person); err == nil {
		t.Error("Did not receive expected error when reading closed stream")
	}

	// Incorrect JSON object
	if err := getJson(bytes.NewBufferString(`{"id":"foo"}`), person); err == nil  {
		t.Error("Did not receive expected error on bad JSON unmarshalling")
	}
}

func TestRouting_PostJson(t *testing.T) {
	expectedPerson, _ := createRandomPerson()
	expectedStatus := http.StatusOK
	response := httptest.NewRecorder()
	postJson(response, expectedStatus, expectedPerson)

	if response.Code != expectedStatus {
		t.Fatalf("Expected http status %d but found %d", expectedStatus, response.Code)
	}

	var person Person
	if err := json.Unmarshal(response.Body.Bytes(), &person); err != nil {
		t.Fatalf("Encountered error when retrieving json from string: %s", err)
	} else if !arePersonEqual(expectedPerson, &person) {
		t.Fatal("Did not receive expected person")
	}

	// TODO: Can the send be nil? Can response be nil?
}

func TestRouting_UpdatePersonHandler(t *testing.T) {
	log := log.New(os.Stdout, "", 0)
	response := httptest.NewRecorder()
	db, server, _ := createRandomDbClient()

	// Build a person
	expectedPerson, _ := createRandomPerson()
	expectedPersonJson, _ := json.Marshal(expectedPerson)
	expectedPersonStr := string(expectedPersonJson)

	// Test creating a new person
	req := httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))
	UpdatePersonHandler(log, db, response, req)
	if response.Code != http.StatusCreated {
		t.Errorf("Did not get expected HTTP status code. Instead got: %d", response.Code)
	}

	if person, err := (*db).Get(expectedPerson.Id); err != nil {
		t.Fatalf("Encountered error when retrieving person: %s", err)
	} else if !arePersonEqual(expectedPerson, person) {
		t.Fatal("Retrieved Person is not equivalent to the expected Person")
	}

	// Test updating the same person
	expectedPerson.Name = createRandomString()
	expectedPersonJson, _ = json.Marshal(expectedPerson)
	expectedPersonStr = string(expectedPersonJson)
	req = httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))
	UpdatePersonHandler(log, db, response, req)
	if response.Code != http.StatusCreated {
		t.Errorf("Did not get expected HTTP status code. Instead got: %d", response.Code)
	}

	if person, err := (*db).Get(expectedPerson.Id); err != nil {
		t.Fatalf("Encountered error when retrieving person: %s", err)
	} else if !arePersonEqual(expectedPerson, person) {
		t.Fatal("Retrieved Person is not equivalent to the expected Person")
	}

	// Test unable to process the body
	req = httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))
	req.Body.Close()
	UpdatePersonHandler(log, db, response, req)

	// Test when the database is unable to comply with the request
	req = httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))
	server.Close()
	UpdatePersonHandler(log, db, response, req)
	if response.Code != http.StatusInternalServerError {
		t.Errorf("Did not get expected HTTP status code of %d. Instead got: %d", http.StatusInternalServerError, response.Code)
	}
}

func TestRouting_PersonRequestHandler(t *testing.T) {
	log := log.New(os.Stdout, "", 0)
	response := httptest.NewRecorder()

	expectedPerson, _ := createRandomPerson()

	var dataReq PersonDataRequest
	dataReq.Ids = make([]string, 1)
	dataReq.Ids[0] = expectedPerson.Id
	personDataRequestJson, _ := json.Marshal(dataReq)
	personDataRequestStr := string(personDataRequestJson)
	req := httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(personDataRequestStr))

	db, server, _ := createRandomDbClient()
	defer server.Close()

	(*db).Create(*expectedPerson)

	PersonRequestHandler(log, db, response, req)

	// TODO: check status code
	// TODO: check number of people returned

	var person PersonDataResponse
	if err := json.Unmarshal(response.Body.Bytes(), &person); err != nil {
		t.Fatalf("Encountered error when retrieving json from string: %s", err)
	} else if !arePersonEqual(expectedPerson, &person.People[0]) {
		t.Fatal("Did not receive expected person")
	}

	// Test that ... the ... dummy db works...
	/* TODO: this mostly just tests the mock db anyway
	req = httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(createRandomString()))
	PersonRequestHandler(log, db, response, req)
	if response.Code != http.StatusUnprocessableEntity {
		t.Errorf("Did not receive expected failure HTTP status on bad request body: %d", response.Code)
	}
	*/
}

func TestRouting_ListenAndServe(t *testing.T) {
	t.Skip("TODO")
	port := rand.Int()
	log := log.New(os.Stdout, "", 0)
	db, server, _ := createRandomDbClient()
	defer server.Close()

	// TODO: use a goroutine to do some tests
	go ListenAndServe(log, port, db)
}
<commit_msg>I might also need a NopCloser, but we shall see if a Reader is enough<commit_after>package dondeestas

import (
	"bytes"
	"encoding/json"
	"log"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
)

func getExpectedPersonRequest() (*http.Request, *Person, *httptest.ResponseRecorder) {
	expectedPerson, _ := createRandomPerson()
	expectedPersonJson, _ := json.Marshal(expectedPerson)
	expectedPersonStr := string(expectedPersonJson)
	req := httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))

	return req, expectedPerson, httptest.NewRecorder()
}

func TestRouting_GetJson(t *testing.T) {
	expectedPerson, _ := createRandomPerson()
	expectedPersonJson, _ := json.Marshal(expectedPerson)
	expectedPersonStr := string(expectedPersonJson)
	req := httptest.NewRequest("GET", "http://blah.com/foo", bytes.NewBufferString(expectedPersonStr))

	var person Person
	if err := getJson(req.Body, person); err != nil {
		t.Fatalf("Encountered error when retrieving json from string: %s", err)
	}

	// Test forcing the function to read a closed stream
	req = httptest.NewRequest("GET", "http://blah.com/foo", bytes.NewBufferString(expectedPersonStr))
	if err := req.Body.Close(); err != nil {
		t.Fatalf("Could not close test request body!")
	}
	if err := getJson(req.Body, person); err == nil {
		t.Error("Did not receive expected error when reading closed stream")
	}

	// Incorrect JSON object
	if err := getJson(bytes.NewReader(bytes.NewBufferString(`{"id":"foo"}`).Bytes()), person); err == nil  {
		t.Error("Did not receive expected error on bad JSON unmarshalling")
	}
}

func TestRouting_PostJson(t *testing.T) {
	expectedPerson, _ := createRandomPerson()
	expectedStatus := http.StatusOK
	response := httptest.NewRecorder()
	postJson(response, expectedStatus, expectedPerson)

	if response.Code != expectedStatus {
		t.Fatalf("Expected http status %d but found %d", expectedStatus, response.Code)
	}

	var person Person
	if err := json.Unmarshal(response.Body.Bytes(), &person); err != nil {
		t.Fatalf("Encountered error when retrieving json from string: %s", err)
	} else if !arePersonEqual(expectedPerson, &person) {
		t.Fatal("Did not receive expected person")
	}

	// TODO: Can the send be nil? Can response be nil?
}

func TestRouting_UpdatePersonHandler(t *testing.T) {
	log := log.New(os.Stdout, "", 0)
	response := httptest.NewRecorder()
	db, server, _ := createRandomDbClient()

	// Build a person
	expectedPerson, _ := createRandomPerson()
	expectedPersonJson, _ := json.Marshal(expectedPerson)
	expectedPersonStr := string(expectedPersonJson)

	// Test creating a new person
	req := httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))
	UpdatePersonHandler(log, db, response, req)
	if response.Code != http.StatusCreated {
		t.Errorf("Did not get expected HTTP status code. Instead got: %d", response.Code)
	}

	if person, err := (*db).Get(expectedPerson.Id); err != nil {
		t.Fatalf("Encountered error when retrieving person: %s", err)
	} else if !arePersonEqual(expectedPerson, person) {
		t.Fatal("Retrieved Person is not equivalent to the expected Person")
	}

	// Test updating the same person
	expectedPerson.Name = createRandomString()
	expectedPersonJson, _ = json.Marshal(expectedPerson)
	expectedPersonStr = string(expectedPersonJson)
	req = httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))
	UpdatePersonHandler(log, db, response, req)
	if response.Code != http.StatusCreated {
		t.Errorf("Did not get expected HTTP status code. Instead got: %d", response.Code)
	}

	if person, err := (*db).Get(expectedPerson.Id); err != nil {
		t.Fatalf("Encountered error when retrieving person: %s", err)
	} else if !arePersonEqual(expectedPerson, person) {
		t.Fatal("Retrieved Person is not equivalent to the expected Person")
	}

	// Test unable to process the body
	req = httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))
	req.Body.Close()
	UpdatePersonHandler(log, db, response, req)

	// Test when the database is unable to comply with the request
	req = httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(expectedPersonStr))
	server.Close()
	UpdatePersonHandler(log, db, response, req)
	if response.Code != http.StatusInternalServerError {
		t.Errorf("Did not get expected HTTP status code of %d. Instead got: %d", http.StatusInternalServerError, response.Code)
	}
}

func TestRouting_PersonRequestHandler(t *testing.T) {
	log := log.New(os.Stdout, "", 0)
	response := httptest.NewRecorder()

	expectedPerson, _ := createRandomPerson()

	var dataReq PersonDataRequest
	dataReq.Ids = make([]string, 1)
	dataReq.Ids[0] = expectedPerson.Id
	personDataRequestJson, _ := json.Marshal(dataReq)
	personDataRequestStr := string(personDataRequestJson)
	req := httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(personDataRequestStr))

	db, server, _ := createRandomDbClient()
	defer server.Close()

	(*db).Create(*expectedPerson)

	PersonRequestHandler(log, db, response, req)

	// TODO: check status code
	// TODO: check number of people returned

	var person PersonDataResponse
	if err := json.Unmarshal(response.Body.Bytes(), &person); err != nil {
		t.Fatalf("Encountered error when retrieving json from string: %s", err)
	} else if !arePersonEqual(expectedPerson, &person.People[0]) {
		t.Fatal("Did not receive expected person")
	}

	// Test that ... the ... dummy db works...
	/* TODO: this mostly just tests the mock db anyway
	req = httptest.NewRequest("GET", "http://"+createRandomString(), bytes.NewBufferString(createRandomString()))
	PersonRequestHandler(log, db, response, req)
	if response.Code != http.StatusUnprocessableEntity {
		t.Errorf("Did not receive expected failure HTTP status on bad request body: %d", response.Code)
	}
	*/
}

func TestRouting_ListenAndServe(t *testing.T) {
	t.Skip("TODO")
	port := rand.Int()
	log := log.New(os.Stdout, "", 0)
	db, server, _ := createRandomDbClient()
	defer server.Close()

	// TODO: use a goroutine to do some tests
	go ListenAndServe(log, port, db)
}
<|endoftext|>


<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 artifact

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	pb "github.com/apache/beam/sdks/go/pkg/beam/model/jobmanagement_v1"
	pipepb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/util/grpcx"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// TestRetrieve tests that we can successfully retrieve fresh files.
func TestRetrieve(t *testing.T) {
	cc := startServer(t)
	defer cc.Close()

	ctx := grpcx.WriteWorkerID(context.Background(), "idA")
	keys := []string{"foo", "bar", "baz/baz/baz"}
	st := "whatever"
	rt, artifacts := populate(ctx, cc, t, keys, 300, st)

	dst := makeTempDir(t)
	defer os.RemoveAll(dst)

	client := pb.NewLegacyArtifactRetrievalServiceClient(cc)
	for _, a := range artifacts {
		filename := makeFilename(dst, a.Name)
		if err := Retrieve(ctx, client, a, rt, dst); err != nil {
			t.Errorf("failed to retrieve %v: %v", a.Name, err)
			continue
		}
		verifySHA256(t, filename, a.Sha256)
	}
}

// TestMultiRetrieve tests that we can successfully retrieve fresh files
// concurrently.
func TestMultiRetrieve(t *testing.T) {
	cc := startServer(t)
	defer cc.Close()

	ctx := grpcx.WriteWorkerID(context.Background(), "idB")
	keys := []string{"1", "2", "3", "4", "a/5", "a/6", "a/7", "a/8", "a/a/9", "a/a/10", "a/b/11", "a/b/12"}
	st := "whatever"
	rt, artifacts := populate(ctx, cc, t, keys, 300, st)

	dst := makeTempDir(t)
	defer os.RemoveAll(dst)

	client := pb.NewLegacyArtifactRetrievalServiceClient(cc)
	if err := LegacyMultiRetrieve(ctx, client, 10, artifacts, rt, dst); err != nil {
		t.Errorf("failed to retrieve: %v", err)
	}

	for _, a := range artifacts {
		verifySHA256(t, makeFilename(dst, a.Name), a.Sha256)
	}
}

// populate stages a set of artifacts with the given keys, each with
// slightly different sizes and chucksizes.
func populate(ctx context.Context, cc *grpc.ClientConn, t *testing.T, keys []string, size int, st string) (string, []*pb.ArtifactMetadata) {
	scl := pb.NewLegacyArtifactStagingServiceClient(cc)

	var artifacts []*pb.ArtifactMetadata
	for i, key := range keys {
		a := stage(ctx, scl, t, key, size+7*i, 97+i, st)
		artifacts = append(artifacts, a)
	}
	token, err := Commit(ctx, scl, artifacts, st)
	if err != nil {
		t.Fatalf("failed to commit manifest: %v", err)
		return "", nil
	}
	return token, artifacts
}

// stage stages an artifact with the given key, size and chuck size. The content is
// always 'z's.
func stage(ctx context.Context, scl pb.LegacyArtifactStagingServiceClient, t *testing.T, key string, size, chunkSize int, st string) *pb.ArtifactMetadata {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = 'z'
	}

	sha256W := sha256.New()
	sha256W.Write(data)
	hash := hex.EncodeToString(sha256W.Sum(nil))
	md := makeArtifact(key, hash)
	pmd := &jobpb.PutArtifactMetadata{
		Metadata:            md,
		StagingSessionToken: st,
	}

	stream, err := scl.PutArtifact(ctx)
	if err != nil {
		t.Fatalf("put failed: %v", err)
	}
	header := &jobpb.PutArtifactRequest{
		Content: &jobpb.PutArtifactRequest_Metadata{
			Metadata: pmd,
		},
	}
	if err := stream.Send(header); err != nil {
		t.Fatalf("send header failed: %v", err)
	}

	for i := 0; i < size; i += chunkSize {
		end := i + chunkSize
		if size < end {
			end = size
		}

		chunk := &jobpb.PutArtifactRequest{
			Content: &jobpb.PutArtifactRequest_Data{
				Data: &pb.ArtifactChunk{
					Data: data[i:end],
				},
			},
		}
		if err := stream.Send(chunk); err != nil {
			t.Fatalf("send chunk[%v:%v] failed: %v", i, end, err)
		}
	}
	if _, err := stream.CloseAndRecv(); err != nil {
		t.Fatalf("close failed: %v", err)
	}
	return md
}

// Test for new artifact retrieval.

func TestNewRetrieveWithManyFiles(t *testing.T) {
	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}

	client := &fakeRetrievalService{
		artifacts: expected,
	}

	dest := makeTempDir(t)
	defer os.RemoveAll(dest)
	ctx := grpcx.WriteWorkerID(context.Background(), "worker")

	mds, err := newMaterializeWithClient(ctx, client, client.resolvedArtifacts(), dest)
	if err != nil {
		t.Fatalf("materialize failed: %v", err)
	}

	checkStagedFiles(mds, dest, expected, t)
}

func TestNewRetrieveWithResolution(t *testing.T) {
	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}

	client := &fakeRetrievalService{
		artifacts: expected,
	}

	dest := makeTempDir(t)
	defer os.RemoveAll(dest)
	ctx := grpcx.WriteWorkerID(context.Background(), "worker")

	mds, err := newMaterializeWithClient(ctx, client, client.unresolvedArtifacts(), dest)
	if err != nil {
		t.Fatalf("materialize failed: %v", err)
	}

	checkStagedFiles(mds, dest, expected, t)
}

func checkStagedFiles(mds []*jobpb.ArtifactMetadata, dest string, expected map[string]string, t *testing.T) {
	if len(mds) != len(expected) {
		t.Errorf("wrong number of artifacts staged %v vs %v", len(mds), len(expected))
	}
	for _, md := range mds {
		filename := filepath.Join(dest, filepath.FromSlash(md.Name))
		fd, err := os.Open(filename)
		if err != nil {
			t.Errorf("error opening file %v", err)
		}
		defer fd.Close()

		data := make([]byte, 1<<20)
		n, err := fd.Read(data)
		if err != nil {
			t.Errorf("error reading file %v", err)
		}

		if string(data[:n]) != expected[md.Name] {
			t.Errorf("missmatched contents for %v: '%s' vs '%s'", md.Name, string(data[:n]), expected[md.Name])
		}
	}
}

type fakeRetrievalService struct {
	artifacts map[string]string // name -> content
}

func (fake *fakeRetrievalService) resolvedArtifacts() []*pipepb.ArtifactInformation {
	var artifacts []*pipepb.ArtifactInformation
	for name, contents := range fake.artifacts {
		payload, _ := proto.Marshal(&pipepb.ArtifactStagingToRolePayload{
			StagedName: name})
		artifacts = append(artifacts, &pipepb.ArtifactInformation{
			TypeUrn:     "resolved",
			TypePayload: []byte(contents),
			RoleUrn:     URNStagingTo,
			RolePayload: payload,
		})
	}
	return artifacts
}

func (fake *fakeRetrievalService) unresolvedArtifacts() []*pipepb.ArtifactInformation {
	return []*pipepb.ArtifactInformation{
		&pipepb.ArtifactInformation{
			TypeUrn: "unresolved",
		},
	}
}

func (fake *fakeRetrievalService) ResolveArtifacts(ctx context.Context, request *jobpb.ResolveArtifactsRequest, opts ...grpc.CallOption) (*jobpb.ResolveArtifactsResponse, error) {
	response := jobpb.ResolveArtifactsResponse{}
	for _, dep := range request.Artifacts {
		if dep.TypeUrn == "unresolved" {
			response.Replacements = append(response.Replacements, fake.resolvedArtifacts()...)
		} else {
			response.Replacements = append(response.Replacements, dep)
		}
	}
	return &response, nil
}

func (fake *fakeRetrievalService) GetArtifact(ctx context.Context, request *jobpb.GetArtifactRequest, opts ...grpc.CallOption) (jobpb.ArtifactRetrievalService_GetArtifactClient, error) {
	if request.Artifact.TypeUrn == "resolved" {
		return &fakeGetArtifactResponseStream{data: request.Artifact.TypePayload}, nil
	}
	return nil, errors.Errorf("Unsupported artifact %v", request.Artifact)
}

type fakeGetArtifactResponseStream struct {
	data  []byte
	index int
}

func (fake *fakeGetArtifactResponseStream) Recv() (*jobpb.GetArtifactResponse, error) {
	if fake.index < len(fake.data) {
		fake.index++
		return &jobpb.GetArtifactResponse{Data: fake.data[fake.index-1 : fake.index]}, nil
	}
	return nil, io.EOF
}

func (fake *fakeGetArtifactResponseStream) RecvMsg(interface{}) error {
	return nil
}

func (fake *fakeGetArtifactResponseStream) SendMsg(interface{}) error {
	return nil
}

func (fake *fakeGetArtifactResponseStream) Header() (metadata.MD, error) {
	return nil, nil
}

func (fake *fakeGetArtifactResponseStream) Trailer() metadata.MD {
	return nil
}

func (fake *fakeGetArtifactResponseStream) Context() context.Context {
	return context.Background()
}

func (fake *fakeGetArtifactResponseStream) CloseSend() error {
	return nil
}

func verifySHA256(t *testing.T, filename, hash string) {
	actual, err := computeSHA256(filename)
	if err != nil {
		t.Errorf("failed to compute hash for %v: %v", filename, err)
		return
	}
	if actual != hash {
		t.Errorf("file %v has bad SHA256: %v, want %v", filename, actual, hash)
	}
}

func makeTempDir(t *testing.T) string {
	dir, err := ioutil.TempDir("", "artifact_test_")
	if err != nil {
		t.Errorf("Test failure: cannot create temporary directory: %+v", err)
	}
	return dir
}

func makeTempFiles(t *testing.T, dir string, keys []string, size int) []string {
	var sha256s []string
	for i, key := range keys {
		hash := makeTempFile(t, makeFilename(dir, key), size+i)
		sha256s = append(sha256s, hash)
	}
	return sha256s
}

func makeTempFile(t *testing.T, filename string, size int) string {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = 'z'
	}

	if err := os.MkdirAll(filepath.Dir(filename), 0755); err != nil {
		t.Fatalf("cannot create directory for %s: %v", filename, err)
	}
	if err := ioutil.WriteFile(filename, data, 0644); err != nil {
		t.Fatalf("cannot create file %s: %v", filename, err)
	}

	sha256W := sha256.New()
	sha256W.Write(data)
	return hex.EncodeToString(sha256W.Sum(nil))
}

func makeArtifact(key, hash string) *jobpb.ArtifactMetadata {
	return &jobpb.ArtifactMetadata{
		Name:        key,
		Sha256:      hash,
		Permissions: 0644,
	}
}

func makeFilename(dir, key string) string {
	return filepath.Join(dir, filepath.FromSlash(key))
}
<commit_msg>Update materialize_test.go<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 artifact

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	jobpb "github.com/apache/beam/sdks/go/pkg/beam/model/jobmanagement_v1"
	pipepb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/util/grpcx"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// TestRetrieve tests that we can successfully retrieve fresh files.
func TestRetrieve(t *testing.T) {
	cc := startServer(t)
	defer cc.Close()

	ctx := grpcx.WriteWorkerID(context.Background(), "idA")
	keys := []string{"foo", "bar", "baz/baz/baz"}
	st := "whatever"
	rt, artifacts := populate(ctx, cc, t, keys, 300, st)

	dst := makeTempDir(t)
	defer os.RemoveAll(dst)

	client := jobpb.NewLegacyArtifactRetrievalServiceClient(cc)
	for _, a := range artifacts {
		filename := makeFilename(dst, a.Name)
		if err := Retrieve(ctx, client, a, rt, dst); err != nil {
			t.Errorf("failed to retrieve %v: %v", a.Name, err)
			continue
		}
		verifySHA256(t, filename, a.Sha256)
	}
}

// TestMultiRetrieve tests that we can successfully retrieve fresh files
// concurrently.
func TestMultiRetrieve(t *testing.T) {
	cc := startServer(t)
	defer cc.Close()

	ctx := grpcx.WriteWorkerID(context.Background(), "idB")
	keys := []string{"1", "2", "3", "4", "a/5", "a/6", "a/7", "a/8", "a/a/9", "a/a/10", "a/b/11", "a/b/12"}
	st := "whatever"
	rt, artifacts := populate(ctx, cc, t, keys, 300, st)

	dst := makeTempDir(t)
	defer os.RemoveAll(dst)

	client := jobpb.NewLegacyArtifactRetrievalServiceClient(cc)
	if err := LegacyMultiRetrieve(ctx, client, 10, artifacts, rt, dst); err != nil {
		t.Errorf("failed to retrieve: %v", err)
	}

	for _, a := range artifacts {
		verifySHA256(t, makeFilename(dst, a.Name), a.Sha256)
	}
}

// populate stages a set of artifacts with the given keys, each with
// slightly different sizes and chucksizes.
func populate(ctx context.Context, cc *grpc.ClientConn, t *testing.T, keys []string, size int, st string) (string, []*jobpb.ArtifactMetadata) {
	scl := jobpb.NewLegacyArtifactStagingServiceClient(cc)

	var artifacts []*jobpb.ArtifactMetadata
	for i, key := range keys {
		a := stage(ctx, scl, t, key, size+7*i, 97+i, st)
		artifacts = append(artifacts, a)
	}
	token, err := Commit(ctx, scl, artifacts, st)
	if err != nil {
		t.Fatalf("failed to commit manifest: %v", err)
		return "", nil
	}
	return token, artifacts
}

// stage stages an artifact with the given key, size and chuck size. The content is
// always 'z's.
func stage(ctx context.Context, scl jobpb.LegacyArtifactStagingServiceClient, t *testing.T, key string, size, chunkSize int, st string) *jobpb.ArtifactMetadata {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = 'z'
	}

	sha256W := sha256.New()
	sha256W.Write(data)
	hash := hex.EncodeToString(sha256W.Sum(nil))
	md := makeArtifact(key, hash)
	pmd := &jobpb.PutArtifactMetadata{
		Metadata:            md,
		StagingSessionToken: st,
	}

	stream, err := scl.PutArtifact(ctx)
	if err != nil {
		t.Fatalf("put failed: %v", err)
	}
	header := &jobpb.PutArtifactRequest{
		Content: &jobpb.PutArtifactRequest_Metadata{
			Metadata: pmd,
		},
	}
	if err := stream.Send(header); err != nil {
		t.Fatalf("send header failed: %v", err)
	}

	for i := 0; i < size; i += chunkSize {
		end := i + chunkSize
		if size < end {
			end = size
		}

		chunk := &jobpb.PutArtifactRequest{
			Content: &jobpb.PutArtifactRequest_Data{
				Data: &jobpb.ArtifactChunk{
					Data: data[i:end],
				},
			},
		}
		if err := stream.Send(chunk); err != nil {
			t.Fatalf("send chunk[%v:%v] failed: %v", i, end, err)
		}
	}
	if _, err := stream.CloseAndRecv(); err != nil {
		t.Fatalf("close failed: %v", err)
	}
	return md
}

// Test for new artifact retrieval.

func TestNewRetrieveWithManyFiles(t *testing.T) {
	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}

	client := &fakeRetrievalService{
		artifacts: expected,
	}

	dest := makeTempDir(t)
	defer os.RemoveAll(dest)
	ctx := grpcx.WriteWorkerID(context.Background(), "worker")

	mds, err := newMaterializeWithClient(ctx, client, client.resolvedArtifacts(), dest)
	if err != nil {
		t.Fatalf("materialize failed: %v", err)
	}

	checkStagedFiles(mds, dest, expected, t)
}

func TestNewRetrieveWithResolution(t *testing.T) {
	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}

	client := &fakeRetrievalService{
		artifacts: expected,
	}

	dest := makeTempDir(t)
	defer os.RemoveAll(dest)
	ctx := grpcx.WriteWorkerID(context.Background(), "worker")

	mds, err := newMaterializeWithClient(ctx, client, client.unresolvedArtifacts(), dest)
	if err != nil {
		t.Fatalf("materialize failed: %v", err)
	}

	checkStagedFiles(mds, dest, expected, t)
}

func checkStagedFiles(mds []*jobpb.ArtifactMetadata, dest string, expected map[string]string, t *testing.T) {
	if len(mds) != len(expected) {
		t.Errorf("wrong number of artifacts staged %v vs %v", len(mds), len(expected))
	}
	for _, md := range mds {
		filename := filepath.Join(dest, filepath.FromSlash(md.Name))
		fd, err := os.Open(filename)
		if err != nil {
			t.Errorf("error opening file %v", err)
		}
		defer fd.Close()

		data := make([]byte, 1<<20)
		n, err := fd.Read(data)
		if err != nil {
			t.Errorf("error reading file %v", err)
		}

		if string(data[:n]) != expected[md.Name] {
			t.Errorf("missmatched contents for %v: '%s' vs '%s'", md.Name, string(data[:n]), expected[md.Name])
		}
	}
}

type fakeRetrievalService struct {
	artifacts map[string]string // name -> content
}

func (fake *fakeRetrievalService) resolvedArtifacts() []*pipepb.ArtifactInformation {
	var artifacts []*pipepb.ArtifactInformation
	for name, contents := range fake.artifacts {
		payload, _ := proto.Marshal(&pipepb.ArtifactStagingToRolePayload{
			StagedName: name})
		artifacts = append(artifacts, &pipepb.ArtifactInformation{
			TypeUrn:     "resolved",
			TypePayload: []byte(contents),
			RoleUrn:     URNStagingTo,
			RolePayload: payload,
		})
	}
	return artifacts
}

func (fake *fakeRetrievalService) unresolvedArtifacts() []*pipepb.ArtifactInformation {
	return []*pipepb.ArtifactInformation{
		&pipepb.ArtifactInformation{
			TypeUrn: "unresolved",
		},
	}
}

func (fake *fakeRetrievalService) ResolveArtifacts(ctx context.Context, request *jobpb.ResolveArtifactsRequest, opts ...grpc.CallOption) (*jobpb.ResolveArtifactsResponse, error) {
	response := jobpb.ResolveArtifactsResponse{}
	for _, dep := range request.Artifacts {
		if dep.TypeUrn == "unresolved" {
			response.Replacements = append(response.Replacements, fake.resolvedArtifacts()...)
		} else {
			response.Replacements = append(response.Replacements, dep)
		}
	}
	return &response, nil
}

func (fake *fakeRetrievalService) GetArtifact(ctx context.Context, request *jobpb.GetArtifactRequest, opts ...grpc.CallOption) (jobpb.ArtifactRetrievalService_GetArtifactClient, error) {
	if request.Artifact.TypeUrn == "resolved" {
		return &fakeGetArtifactResponseStream{data: request.Artifact.TypePayload}, nil
	}
	return nil, errors.Errorf("Unsupported artifact %v", request.Artifact)
}

type fakeGetArtifactResponseStream struct {
	data  []byte
	index int
}

func (fake *fakeGetArtifactResponseStream) Recv() (*jobpb.GetArtifactResponse, error) {
	if fake.index < len(fake.data) {
		fake.index++
		return &jobpb.GetArtifactResponse{Data: fake.data[fake.index-1 : fake.index]}, nil
	}
	return nil, io.EOF
}

func (fake *fakeGetArtifactResponseStream) RecvMsg(interface{}) error {
	return nil
}

func (fake *fakeGetArtifactResponseStream) SendMsg(interface{}) error {
	return nil
}

func (fake *fakeGetArtifactResponseStream) Header() (metadata.MD, error) {
	return nil, nil
}

func (fake *fakeGetArtifactResponseStream) Trailer() metadata.MD {
	return nil
}

func (fake *fakeGetArtifactResponseStream) Context() context.Context {
	return context.Background()
}

func (fake *fakeGetArtifactResponseStream) CloseSend() error {
	return nil
}

func verifySHA256(t *testing.T, filename, hash string) {
	actual, err := computeSHA256(filename)
	if err != nil {
		t.Errorf("failed to compute hash for %v: %v", filename, err)
		return
	}
	if actual != hash {
		t.Errorf("file %v has bad SHA256: %v, want %v", filename, actual, hash)
	}
}

func makeTempDir(t *testing.T) string {
	dir, err := ioutil.TempDir("", "artifact_test_")
	if err != nil {
		t.Errorf("Test failure: cannot create temporary directory: %+v", err)
	}
	return dir
}

func makeTempFiles(t *testing.T, dir string, keys []string, size int) []string {
	var sha256s []string
	for i, key := range keys {
		hash := makeTempFile(t, makeFilename(dir, key), size+i)
		sha256s = append(sha256s, hash)
	}
	return sha256s
}

func makeTempFile(t *testing.T, filename string, size int) string {
	data := make([]byte, size)
	for i := 0; i < size; i++ {
		data[i] = 'z'
	}

	if err := os.MkdirAll(filepath.Dir(filename), 0755); err != nil {
		t.Fatalf("cannot create directory for %s: %v", filename, err)
	}
	if err := ioutil.WriteFile(filename, data, 0644); err != nil {
		t.Fatalf("cannot create file %s: %v", filename, err)
	}

	sha256W := sha256.New()
	sha256W.Write(data)
	return hex.EncodeToString(sha256W.Sum(nil))
}

func makeArtifact(key, hash string) *jobpb.ArtifactMetadata {
	return &jobpb.ArtifactMetadata{
		Name:        key,
		Sha256:      hash,
		Permissions: 0644,
	}
}

func makeFilename(dir, key string) string {
	return filepath.Join(dir, filepath.FromSlash(key))
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package app does all of the work necessary to create a Kubernetes
// APIServer by binding together the API, master and APIServer infrastructure.
// It can be configured and called directly or via the hyperkube cache.
package app

import (
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/pborman/uuid"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	"k8s.io/kubernetes/federation/cmd/federation-apiserver/app/options"
	"k8s.io/kubernetes/pkg/admission"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/apiserver/authenticator"
	authorizerunion "k8s.io/kubernetes/pkg/auth/authorizer/union"
	"k8s.io/kubernetes/pkg/auth/user"
	"k8s.io/kubernetes/pkg/controller/informers"
	"k8s.io/kubernetes/pkg/generated/openapi"
	"k8s.io/kubernetes/pkg/genericapiserver"
	"k8s.io/kubernetes/pkg/genericapiserver/authorizer"
	genericvalidation "k8s.io/kubernetes/pkg/genericapiserver/validation"
	"k8s.io/kubernetes/pkg/registry/cachesize"
	"k8s.io/kubernetes/pkg/registry/generic"
	"k8s.io/kubernetes/pkg/registry/generic/registry"
	"k8s.io/kubernetes/pkg/routes"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/pkg/version"
	authenticatorunion "k8s.io/kubernetes/plugin/pkg/auth/authenticator/request/union"
)

// NewAPIServerCommand creates a *cobra.Command object with default parameters
func NewAPIServerCommand() *cobra.Command {
	s := options.NewServerRunOptions()
	s.AddFlags(pflag.CommandLine)
	cmd := &cobra.Command{
		Use: "federation-apiserver",
		Long: `The Kubernetes federation API server validates and configures data
for the api objects which include pods, services, replicationcontrollers, and
others. The API Server services REST operations and provides the frontend to the
cluster's shared state through which all other components interact.`,
		Run: func(cmd *cobra.Command, args []string) {
		},
	}
	return cmd
}

// Run runs the specified APIServer.  This should never exit.
func Run(s *options.ServerRunOptions) error {
	genericvalidation.VerifyEtcdServersList(s.GenericServerRunOptions)
	genericapiserver.DefaultAndValidateRunOptions(s.GenericServerRunOptions)
	genericConfig := genericapiserver.NewConfig(). // create the new config
							ApplyOptions(s.GenericServerRunOptions). // apply the options selected
							Complete()                               // set default values based on the known values

	if err := genericConfig.MaybeGenerateServingCerts(); err != nil {
		glog.Fatalf("Failed to generate service certificate: %v", err)
	}

	// TODO: register cluster federation resources here.
	resourceConfig := genericapiserver.NewResourceConfig()

	if s.GenericServerRunOptions.StorageConfig.DeserializationCacheSize == 0 {
		// When size of cache is not explicitly set, set it to 50000
		s.GenericServerRunOptions.StorageConfig.DeserializationCacheSize = 50000
	}
	storageGroupsToEncodingVersion, err := s.GenericServerRunOptions.StorageGroupsToEncodingVersion()
	if err != nil {
		glog.Fatalf("error generating storage version map: %s", err)
	}
	storageFactory, err := genericapiserver.BuildDefaultStorageFactory(
		s.GenericServerRunOptions.StorageConfig, s.GenericServerRunOptions.DefaultStorageMediaType, api.Codecs,
		genericapiserver.NewDefaultResourceEncodingConfig(), storageGroupsToEncodingVersion,
		[]unversioned.GroupVersionResource{}, resourceConfig, s.GenericServerRunOptions.RuntimeConfig)
	if err != nil {
		glog.Fatalf("error in initializing storage factory: %s", err)
	}

	for _, override := range s.GenericServerRunOptions.EtcdServersOverrides {
		tokens := strings.Split(override, "#")
		if len(tokens) != 2 {
			glog.Errorf("invalid value of etcd server overrides: %s", override)
			continue
		}

		apiresource := strings.Split(tokens[0], "/")
		if len(apiresource) != 2 {
			glog.Errorf("invalid resource definition: %s", tokens[0])
			continue
		}
		group := apiresource[0]
		resource := apiresource[1]
		groupResource := unversioned.GroupResource{Group: group, Resource: resource}

		servers := strings.Split(tokens[1], ";")
		storageFactory.SetEtcdLocation(groupResource, servers)
	}

	apiAuthenticator, securityDefinitions, err := authenticator.New(authenticator.AuthenticatorConfig{
		Anonymous:           s.GenericServerRunOptions.AnonymousAuth,
		AnyToken:            s.GenericServerRunOptions.EnableAnyToken,
		BasicAuthFile:       s.GenericServerRunOptions.BasicAuthFile,
		ClientCAFile:        s.GenericServerRunOptions.ClientCAFile,
		TokenAuthFile:       s.GenericServerRunOptions.TokenAuthFile,
		OIDCIssuerURL:       s.GenericServerRunOptions.OIDCIssuerURL,
		OIDCClientID:        s.GenericServerRunOptions.OIDCClientID,
		OIDCCAFile:          s.GenericServerRunOptions.OIDCCAFile,
		OIDCUsernameClaim:   s.GenericServerRunOptions.OIDCUsernameClaim,
		OIDCGroupsClaim:     s.GenericServerRunOptions.OIDCGroupsClaim,
		KeystoneURL:         s.GenericServerRunOptions.KeystoneURL,
		RequestHeaderConfig: s.GenericServerRunOptions.AuthenticationRequestHeaderConfig(),
	})
	if err != nil {
		glog.Fatalf("Invalid Authentication Config: %v", err)
	}

	privilegedLoopbackToken := uuid.NewRandom().String()
	selfClientConfig, err := s.GenericServerRunOptions.NewSelfClientConfig(privilegedLoopbackToken)
	if err != nil {
		glog.Fatalf("Failed to create clientset: %v", err)
	}
	client, err := s.GenericServerRunOptions.NewSelfClient(privilegedLoopbackToken)
	if err != nil {
		glog.Errorf("Failed to create clientset: %v", err)
	}
	sharedInformers := informers.NewSharedInformerFactory(client, 10*time.Minute)

	authorizationConfig := authorizer.AuthorizationConfig{
		PolicyFile:                  s.GenericServerRunOptions.AuthorizationPolicyFile,
		WebhookConfigFile:           s.GenericServerRunOptions.AuthorizationWebhookConfigFile,
		WebhookCacheAuthorizedTTL:   s.GenericServerRunOptions.AuthorizationWebhookCacheAuthorizedTTL,
		WebhookCacheUnauthorizedTTL: s.GenericServerRunOptions.AuthorizationWebhookCacheUnauthorizedTTL,
		RBACSuperUser:               s.GenericServerRunOptions.AuthorizationRBACSuperUser,
		InformerFactory:             sharedInformers,
	}
	authorizationModeNames := strings.Split(s.GenericServerRunOptions.AuthorizationMode, ",")
	apiAuthorizer, err := authorizer.NewAuthorizerFromAuthorizationConfig(authorizationModeNames, authorizationConfig)
	if err != nil {
		glog.Fatalf("Invalid Authorization Config: %v", err)
	}

	admissionControlPluginNames := strings.Split(s.GenericServerRunOptions.AdmissionControl, ",")

	// TODO(dims): We probably need to add an option "EnableLoopbackToken"
	if apiAuthenticator != nil {
		var uid = uuid.NewRandom().String()
		tokens := make(map[string]*user.DefaultInfo)
		tokens[privilegedLoopbackToken] = &user.DefaultInfo{
			Name:   user.APIServerUser,
			UID:    uid,
			Groups: []string{user.SystemPrivilegedGroup},
		}

		tokenAuthenticator := authenticator.NewAuthenticatorFromTokens(tokens)
		apiAuthenticator = authenticatorunion.New(tokenAuthenticator, apiAuthenticator)

		tokenAuthorizer := authorizer.NewPrivilegedGroups(user.SystemPrivilegedGroup)
		apiAuthorizer = authorizerunion.New(tokenAuthorizer, apiAuthorizer)
	}

	pluginInitializer := admission.NewPluginInitializer(sharedInformers, apiAuthorizer)

	admissionController, err := admission.NewFromPlugins(client, admissionControlPluginNames, s.GenericServerRunOptions.AdmissionControlConfigFile, pluginInitializer)
	if err != nil {
		glog.Fatalf("Failed to initialize plugins: %v", err)
	}

	kubeVersion := version.Get()
	genericConfig.Version = &kubeVersion
	genericConfig.LoopbackClientConfig = selfClientConfig
	genericConfig.Authenticator = apiAuthenticator
	genericConfig.Authorizer = apiAuthorizer
	genericConfig.AdmissionControl = admissionController
	genericConfig.APIResourceConfigSource = storageFactory.APIResourceConfigSource
	genericConfig.OpenAPIConfig.Definitions = openapi.OpenAPIDefinitions
	genericConfig.EnableOpenAPISupport = true
	genericConfig.OpenAPIConfig.SecurityDefinitions = securityDefinitions

	// TODO: Move this to generic api server (Need to move the command line flag).
	if s.GenericServerRunOptions.EnableWatchCache {
		cachesize.InitializeWatchCacheSizes(s.GenericServerRunOptions.TargetRAMMB)
		cachesize.SetWatchCacheSizes(s.GenericServerRunOptions.WatchCacheSizes)
	}

	m, err := genericConfig.New()
	if err != nil {
		return err
	}

	routes.UIRedirect{}.Install(m.HandlerContainer)
	routes.Logs{}.Install(m.HandlerContainer)

	restOptionsFactory := restOptionsFactory{
		storageFactory:          storageFactory,
		deleteCollectionWorkers: s.GenericServerRunOptions.DeleteCollectionWorkers,
	}
	if s.GenericServerRunOptions.EnableWatchCache {
		restOptionsFactory.storageDecorator = registry.StorageWithCacher
	} else {
		restOptionsFactory.storageDecorator = generic.UndecoratedStorage
	}

	installFederationAPIs(m, restOptionsFactory)
	installCoreAPIs(s, m, restOptionsFactory)
	installExtensionsAPIs(m, restOptionsFactory)

	sharedInformers.Start(wait.NeverStop)
	m.PrepareRun().Run()
	return nil
}

type restOptionsFactory struct {
	storageFactory          genericapiserver.StorageFactory
	storageDecorator        generic.StorageDecorator
	deleteCollectionWorkers int
}

func (f restOptionsFactory) NewFor(resource unversioned.GroupResource) generic.RESTOptions {
	config, err := f.storageFactory.NewConfig(resource)
	if err != nil {
		glog.Fatalf("Unable to find storage config for %v, due to %v", resource, err.Error())
	}
	return generic.RESTOptions{
		StorageConfig:           config,
		Decorator:               f.storageDecorator,
		DeleteCollectionWorkers: f.deleteCollectionWorkers,
		ResourcePrefix:          f.storageFactory.ResourcePrefix(resource),
	}
}
<commit_msg>Make GenericApiServer.Run interruptable and fail on first listen<commit_after>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package app does all of the work necessary to create a Kubernetes
// APIServer by binding together the API, master and APIServer infrastructure.
// It can be configured and called directly or via the hyperkube cache.
package app

import (
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/pborman/uuid"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	"k8s.io/kubernetes/federation/cmd/federation-apiserver/app/options"
	"k8s.io/kubernetes/pkg/admission"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/apiserver/authenticator"
	authorizerunion "k8s.io/kubernetes/pkg/auth/authorizer/union"
	"k8s.io/kubernetes/pkg/auth/user"
	"k8s.io/kubernetes/pkg/controller/informers"
	"k8s.io/kubernetes/pkg/generated/openapi"
	"k8s.io/kubernetes/pkg/genericapiserver"
	"k8s.io/kubernetes/pkg/genericapiserver/authorizer"
	genericvalidation "k8s.io/kubernetes/pkg/genericapiserver/validation"
	"k8s.io/kubernetes/pkg/registry/cachesize"
	"k8s.io/kubernetes/pkg/registry/generic"
	"k8s.io/kubernetes/pkg/registry/generic/registry"
	"k8s.io/kubernetes/pkg/routes"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/pkg/version"
	authenticatorunion "k8s.io/kubernetes/plugin/pkg/auth/authenticator/request/union"
)

// NewAPIServerCommand creates a *cobra.Command object with default parameters
func NewAPIServerCommand() *cobra.Command {
	s := options.NewServerRunOptions()
	s.AddFlags(pflag.CommandLine)
	cmd := &cobra.Command{
		Use: "federation-apiserver",
		Long: `The Kubernetes federation API server validates and configures data
for the api objects which include pods, services, replicationcontrollers, and
others. The API Server services REST operations and provides the frontend to the
cluster's shared state through which all other components interact.`,
		Run: func(cmd *cobra.Command, args []string) {
		},
	}
	return cmd
}

// Run runs the specified APIServer.  This should never exit.
func Run(s *options.ServerRunOptions) error {
	genericvalidation.VerifyEtcdServersList(s.GenericServerRunOptions)
	genericapiserver.DefaultAndValidateRunOptions(s.GenericServerRunOptions)
	genericConfig := genericapiserver.NewConfig(). // create the new config
							ApplyOptions(s.GenericServerRunOptions). // apply the options selected
							Complete()                               // set default values based on the known values

	if err := genericConfig.MaybeGenerateServingCerts(); err != nil {
		glog.Fatalf("Failed to generate service certificate: %v", err)
	}

	// TODO: register cluster federation resources here.
	resourceConfig := genericapiserver.NewResourceConfig()

	if s.GenericServerRunOptions.StorageConfig.DeserializationCacheSize == 0 {
		// When size of cache is not explicitly set, set it to 50000
		s.GenericServerRunOptions.StorageConfig.DeserializationCacheSize = 50000
	}
	storageGroupsToEncodingVersion, err := s.GenericServerRunOptions.StorageGroupsToEncodingVersion()
	if err != nil {
		glog.Fatalf("error generating storage version map: %s", err)
	}
	storageFactory, err := genericapiserver.BuildDefaultStorageFactory(
		s.GenericServerRunOptions.StorageConfig, s.GenericServerRunOptions.DefaultStorageMediaType, api.Codecs,
		genericapiserver.NewDefaultResourceEncodingConfig(), storageGroupsToEncodingVersion,
		[]unversioned.GroupVersionResource{}, resourceConfig, s.GenericServerRunOptions.RuntimeConfig)
	if err != nil {
		glog.Fatalf("error in initializing storage factory: %s", err)
	}

	for _, override := range s.GenericServerRunOptions.EtcdServersOverrides {
		tokens := strings.Split(override, "#")
		if len(tokens) != 2 {
			glog.Errorf("invalid value of etcd server overrides: %s", override)
			continue
		}

		apiresource := strings.Split(tokens[0], "/")
		if len(apiresource) != 2 {
			glog.Errorf("invalid resource definition: %s", tokens[0])
			continue
		}
		group := apiresource[0]
		resource := apiresource[1]
		groupResource := unversioned.GroupResource{Group: group, Resource: resource}

		servers := strings.Split(tokens[1], ";")
		storageFactory.SetEtcdLocation(groupResource, servers)
	}

	apiAuthenticator, securityDefinitions, err := authenticator.New(authenticator.AuthenticatorConfig{
		Anonymous:           s.GenericServerRunOptions.AnonymousAuth,
		AnyToken:            s.GenericServerRunOptions.EnableAnyToken,
		BasicAuthFile:       s.GenericServerRunOptions.BasicAuthFile,
		ClientCAFile:        s.GenericServerRunOptions.ClientCAFile,
		TokenAuthFile:       s.GenericServerRunOptions.TokenAuthFile,
		OIDCIssuerURL:       s.GenericServerRunOptions.OIDCIssuerURL,
		OIDCClientID:        s.GenericServerRunOptions.OIDCClientID,
		OIDCCAFile:          s.GenericServerRunOptions.OIDCCAFile,
		OIDCUsernameClaim:   s.GenericServerRunOptions.OIDCUsernameClaim,
		OIDCGroupsClaim:     s.GenericServerRunOptions.OIDCGroupsClaim,
		KeystoneURL:         s.GenericServerRunOptions.KeystoneURL,
		RequestHeaderConfig: s.GenericServerRunOptions.AuthenticationRequestHeaderConfig(),
	})
	if err != nil {
		glog.Fatalf("Invalid Authentication Config: %v", err)
	}

	privilegedLoopbackToken := uuid.NewRandom().String()
	selfClientConfig, err := s.GenericServerRunOptions.NewSelfClientConfig(privilegedLoopbackToken)
	if err != nil {
		glog.Fatalf("Failed to create clientset: %v", err)
	}
	client, err := s.GenericServerRunOptions.NewSelfClient(privilegedLoopbackToken)
	if err != nil {
		glog.Errorf("Failed to create clientset: %v", err)
	}
	sharedInformers := informers.NewSharedInformerFactory(client, 10*time.Minute)

	authorizationConfig := authorizer.AuthorizationConfig{
		PolicyFile:                  s.GenericServerRunOptions.AuthorizationPolicyFile,
		WebhookConfigFile:           s.GenericServerRunOptions.AuthorizationWebhookConfigFile,
		WebhookCacheAuthorizedTTL:   s.GenericServerRunOptions.AuthorizationWebhookCacheAuthorizedTTL,
		WebhookCacheUnauthorizedTTL: s.GenericServerRunOptions.AuthorizationWebhookCacheUnauthorizedTTL,
		RBACSuperUser:               s.GenericServerRunOptions.AuthorizationRBACSuperUser,
		InformerFactory:             sharedInformers,
	}
	authorizationModeNames := strings.Split(s.GenericServerRunOptions.AuthorizationMode, ",")
	apiAuthorizer, err := authorizer.NewAuthorizerFromAuthorizationConfig(authorizationModeNames, authorizationConfig)
	if err != nil {
		glog.Fatalf("Invalid Authorization Config: %v", err)
	}

	admissionControlPluginNames := strings.Split(s.GenericServerRunOptions.AdmissionControl, ",")

	// TODO(dims): We probably need to add an option "EnableLoopbackToken"
	if apiAuthenticator != nil {
		var uid = uuid.NewRandom().String()
		tokens := make(map[string]*user.DefaultInfo)
		tokens[privilegedLoopbackToken] = &user.DefaultInfo{
			Name:   user.APIServerUser,
			UID:    uid,
			Groups: []string{user.SystemPrivilegedGroup},
		}

		tokenAuthenticator := authenticator.NewAuthenticatorFromTokens(tokens)
		apiAuthenticator = authenticatorunion.New(tokenAuthenticator, apiAuthenticator)

		tokenAuthorizer := authorizer.NewPrivilegedGroups(user.SystemPrivilegedGroup)
		apiAuthorizer = authorizerunion.New(tokenAuthorizer, apiAuthorizer)
	}

	pluginInitializer := admission.NewPluginInitializer(sharedInformers, apiAuthorizer)

	admissionController, err := admission.NewFromPlugins(client, admissionControlPluginNames, s.GenericServerRunOptions.AdmissionControlConfigFile, pluginInitializer)
	if err != nil {
		glog.Fatalf("Failed to initialize plugins: %v", err)
	}

	kubeVersion := version.Get()
	genericConfig.Version = &kubeVersion
	genericConfig.LoopbackClientConfig = selfClientConfig
	genericConfig.Authenticator = apiAuthenticator
	genericConfig.Authorizer = apiAuthorizer
	genericConfig.AdmissionControl = admissionController
	genericConfig.APIResourceConfigSource = storageFactory.APIResourceConfigSource
	genericConfig.OpenAPIConfig.Definitions = openapi.OpenAPIDefinitions
	genericConfig.EnableOpenAPISupport = true
	genericConfig.OpenAPIConfig.SecurityDefinitions = securityDefinitions

	// TODO: Move this to generic api server (Need to move the command line flag).
	if s.GenericServerRunOptions.EnableWatchCache {
		cachesize.InitializeWatchCacheSizes(s.GenericServerRunOptions.TargetRAMMB)
		cachesize.SetWatchCacheSizes(s.GenericServerRunOptions.WatchCacheSizes)
	}

	m, err := genericConfig.New()
	if err != nil {
		return err
	}

	routes.UIRedirect{}.Install(m.HandlerContainer)
	routes.Logs{}.Install(m.HandlerContainer)

	restOptionsFactory := restOptionsFactory{
		storageFactory:          storageFactory,
		deleteCollectionWorkers: s.GenericServerRunOptions.DeleteCollectionWorkers,
	}
	if s.GenericServerRunOptions.EnableWatchCache {
		restOptionsFactory.storageDecorator = registry.StorageWithCacher
	} else {
		restOptionsFactory.storageDecorator = generic.UndecoratedStorage
	}

	installFederationAPIs(m, restOptionsFactory)
	installCoreAPIs(s, m, restOptionsFactory)
	installExtensionsAPIs(m, restOptionsFactory)

	sharedInformers.Start(wait.NeverStop)
	m.PrepareRun().Run(wait.NeverStop)
	return nil
}

type restOptionsFactory struct {
	storageFactory          genericapiserver.StorageFactory
	storageDecorator        generic.StorageDecorator
	deleteCollectionWorkers int
}

func (f restOptionsFactory) NewFor(resource unversioned.GroupResource) generic.RESTOptions {
	config, err := f.storageFactory.NewConfig(resource)
	if err != nil {
		glog.Fatalf("Unable to find storage config for %v, due to %v", resource, err.Error())
	}
	return generic.RESTOptions{
		StorageConfig:           config,
		Decorator:               f.storageDecorator,
		DeleteCollectionWorkers: f.deleteCollectionWorkers,
		ResourcePrefix:          f.storageFactory.ResourcePrefix(resource),
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package client

import (
	"github.com/keybase/cli"
	"github.com/keybase/client/go/libcmdline"
	"github.com/keybase/client/go/libkb"
	"golang.org/x/net/context"
)

type mode int

const (
	modeNormal mode = iota
	modeShort
	modeVerbose
)

type CmdVersion struct {
	mode mode
	svc  bool
	libkb.Contextified
}

func NewCmdVersionRunner(g *libkb.GlobalContext) *CmdVersion {
	return &CmdVersion{
		mode:         modeNormal,
		svc:          true,
		Contextified: libkb.NewContextified(g),
	}
}

func NewCmdVersion(cl *libcmdline.CommandLine, g *libkb.GlobalContext) cli.Command {
	return cli.Command{
		Name:  "version",
		Usage: "Print out version and build information",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "f, format",
				Usage: "Alternate format for version output. Specify 's' for simple (1.2.3) or 'v' for verbose. Default (blank) includes build number (1.2.3-400).",
			},
			cli.BoolFlag{
				Name:  "S, no-service",
				Usage: "Don't report on the service's build information",
			},
		},
		Action: func(c *cli.Context) {
			cl.ChooseCommand(NewCmdVersionRunner(g), "verison", c)
			cl.SetForkCmd(libcmdline.NoFork)
			cl.SetLogForward(libcmdline.LogForwardNone)
		},
	}
}

func (v *CmdVersion) ParseArgv(c *cli.Context) error {
	switch c.String("format") {
	case "s":
		v.mode = modeShort
	case "v":
		v.mode = modeVerbose
	}
	v.svc = !c.Bool("S")
	return nil
}

func (v *CmdVersion) Run() error {
	var err error
	v.runLocal()
	if v.svc {
		err = v.runService()
	}
	return err
}

func (v *CmdVersion) runService() error {
	cli, err := GetConfigClient(v.G())
	if err != nil {
		v.G().Log.Debug("no service running: %v", err)
		return nil
	}
	res, err := cli.GetConfig(context.TODO(), 0)
	if err != nil {
		return err
	}

	dui := v.G().UI.GetDumbOutputUI()

	switch v.mode {
	case modeShort:
		dui.Printf("Service: %s\n", res.VersionShort)
	case modeNormal:
		dui.Printf("Service: %s\n", res.Version)
	case modeVerbose:
		dui.Printf("\n--------- Service Version Information -------\n")
		dui.Printf("%s\n", res.VersionFull)
	}
	return nil
}

func (v *CmdVersion) runLocal() {
	dui := v.G().UI.GetDumbOutputUI()
	prfx := ""
	if v.svc {
		prfx = "Client:  "
	}
	switch v.mode {
	case modeShort:
		dui.Printf("%s%s\n", prfx, libkb.Version)
	case modeNormal:
		dui.Printf("%s%s\n", prfx, libkb.VersionString())
	case modeVerbose:
		libkb.VersionMessage(func(s string) { dui.Printf("%s\n", s) })
	}
}

func (v *CmdVersion) GetUsage() libkb.Usage {
	return libkb.Usage{}
}
<commit_msg>Fix typo<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package client

import (
	"github.com/keybase/cli"
	"github.com/keybase/client/go/libcmdline"
	"github.com/keybase/client/go/libkb"
	"golang.org/x/net/context"
)

type mode int

const (
	modeNormal mode = iota
	modeShort
	modeVerbose
)

type CmdVersion struct {
	mode mode
	svc  bool
	libkb.Contextified
}

func NewCmdVersionRunner(g *libkb.GlobalContext) *CmdVersion {
	return &CmdVersion{
		mode:         modeNormal,
		svc:          true,
		Contextified: libkb.NewContextified(g),
	}
}

func NewCmdVersion(cl *libcmdline.CommandLine, g *libkb.GlobalContext) cli.Command {
	return cli.Command{
		Name:  "version",
		Usage: "Print out version and build information",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "f, format",
				Usage: "Alternate format for version output. Specify 's' for simple (1.2.3) or 'v' for verbose. Default (blank) includes build number (1.2.3-400).",
			},
			cli.BoolFlag{
				Name:  "S, no-service",
				Usage: "Don't report on the service's build information",
			},
		},
		Action: func(c *cli.Context) {
			cl.ChooseCommand(NewCmdVersionRunner(g), "version", c)
			cl.SetForkCmd(libcmdline.NoFork)
			cl.SetLogForward(libcmdline.LogForwardNone)
		},
	}
}

func (v *CmdVersion) ParseArgv(c *cli.Context) error {
	switch c.String("format") {
	case "s":
		v.mode = modeShort
	case "v":
		v.mode = modeVerbose
	}
	v.svc = !c.Bool("S")
	return nil
}

func (v *CmdVersion) Run() error {
	var err error
	v.runLocal()
	if v.svc {
		err = v.runService()
	}
	return err
}

func (v *CmdVersion) runService() error {
	cli, err := GetConfigClient(v.G())
	if err != nil {
		v.G().Log.Debug("no service running: %v", err)
		return nil
	}
	res, err := cli.GetConfig(context.TODO(), 0)
	if err != nil {
		return err
	}

	dui := v.G().UI.GetDumbOutputUI()

	switch v.mode {
	case modeShort:
		dui.Printf("Service: %s\n", res.VersionShort)
	case modeNormal:
		dui.Printf("Service: %s\n", res.Version)
	case modeVerbose:
		dui.Printf("\n--------- Service Version Information -------\n")
		dui.Printf("%s\n", res.VersionFull)
	}
	return nil
}

func (v *CmdVersion) runLocal() {
	dui := v.G().UI.GetDumbOutputUI()
	prfx := ""
	if v.svc {
		prfx = "Client:  "
	}
	switch v.mode {
	case modeShort:
		dui.Printf("%s%s\n", prfx, libkb.Version)
	case modeNormal:
		dui.Printf("%s%s\n", prfx, libkb.VersionString())
	case modeVerbose:
		libkb.VersionMessage(func(s string) { dui.Printf("%s\n", s) })
	}
}

func (v *CmdVersion) GetUsage() libkb.Usage {
	return libkb.Usage{}
}
<|endoftext|>


<commit_before><commit_msg>BUG FIX: goroutine leakage when kafka has problem<commit_after><|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package logger

import (
	"fmt"
	"runtime"
	"strings"
	"time"

	logging "github.com/keybase/go-logging"

	"golang.org/x/net/context"
)

// TestLogBackend is an interface for logging to a test object (i.e.,
// a *testing.T).  We define this in order to avoid pulling in the
// "testing" package in exported code.
type TestLogBackend interface {
	Error(args ...interface{})
	Errorf(format string, args ...interface{})
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Log(args ...interface{})
	Logf(format string, args ...interface{})
	Failed() bool
}

// TestLogger is a Logger that writes to a TestLogBackend.  All
// messages except Fatal are printed using Logf, to avoid failing a
// test that is trying to test an error condition.  No context tags
// are logged.
type TestLogger struct {
	log        TestLogBackend
	extraDepth int
}

func NewTestLogger(log TestLogBackend) *TestLogger {
	return &TestLogger{log: log}
}

// Verify TestLogger fully implements the Logger interface.
var _ Logger = (*TestLogger)(nil)

// ctx can be `nil`
func (log *TestLogger) common(ctx context.Context, lvl logging.Level, useFatal bool, fmts string, arg ...interface{}) {
	if ctx != nil {
		if useFatal {
			log.log.Fatalf(prepareString(ctx,
				log.prefixCaller(log.extraDepth, lvl, fmts)), arg...)
		} else {
			log.log.Logf(prepareString(ctx,
				log.prefixCaller(log.extraDepth, lvl, fmts)), arg...)
		}
	} else {
		if useFatal {
			log.log.Fatalf(log.prefixCaller(log.extraDepth, lvl, fmts), arg...)
		} else {
			log.log.Logf(log.prefixCaller(log.extraDepth, lvl, fmts), arg...)
		}
	}
}

func (log *TestLogger) prefixCaller(extraDepth int, lvl logging.Level, fmts string) string {
	// The testing library doesn't let us control the stack depth,
	// and it always prints out its own prefix, so use \r to clear
	// it out (at least on a terminal) and do our own formatting.
	_, file, line, _ := runtime.Caller(3 + extraDepth)
	elements := strings.Split(file, "/")
	failed := ""
	if log.log.Failed() {
		failed = "[X] "
	}

	return fmt.Sprintf("\r%s %s%s:%d: [%.1s] %s", time.Now().Format("2006-01-02 15:04:05.00000"),
		failed, elements[len(elements)-1], line, lvl, fmts)
}

func (log *TestLogger) Debug(fmts string, arg ...interface{}) {
	log.common(nil, logging.INFO, false, fmts, arg...)
}

func (log *TestLogger) CDebugf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.DEBUG, false, fmts, arg...)
}

func (log *TestLogger) Info(fmts string, arg ...interface{}) {
	log.common(nil, logging.INFO, false, fmts, arg...)
}

func (log *TestLogger) CInfof(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.INFO, false, fmts, arg...)
}

func (log *TestLogger) Notice(fmts string, arg ...interface{}) {
	log.common(nil, logging.NOTICE, false, fmts, arg...)
}

func (log *TestLogger) CNoticef(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.NOTICE, false, fmts, arg...)
}

func (log *TestLogger) Warning(fmts string, arg ...interface{}) {
	log.common(nil, logging.WARNING, false, fmts, arg...)
}

func (log *TestLogger) CWarningf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.WARNING, false, fmts, arg...)
}

func (log *TestLogger) Error(fmts string, arg ...interface{}) {
	log.common(nil, logging.ERROR, false, fmts, arg...)
}

func (log *TestLogger) Errorf(fmts string, arg ...interface{}) {
	log.common(nil, logging.ERROR, false, fmts, arg...)
}

func (log *TestLogger) CErrorf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.ERROR, false, fmts, arg...)
}

func (log *TestLogger) Critical(fmts string, arg ...interface{}) {
	log.common(nil, logging.CRITICAL, false, fmts, arg...)
}

func (log *TestLogger) CCriticalf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.CRITICAL, false, fmts, arg...)
}

func (log *TestLogger) Fatalf(fmts string, arg ...interface{}) {
	log.common(nil, logging.CRITICAL, true, fmts, arg...)
}

func (log *TestLogger) CFatalf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.CRITICAL, true, fmts, arg...)
}

func (log *TestLogger) Profile(fmts string, arg ...interface{}) {
	log.common(nil, logging.CRITICAL, false, fmts, arg...)
}

func (log *TestLogger) Configure(style string, debug bool, filename string) {
	// no-op
}

func (log *TestLogger) RotateLogFile() error {
	// no-op
	return nil
}

func (log *TestLogger) CloneWithAddedDepth(depth int) Logger {
	clone := *log
	clone.extraDepth += depth
	return &clone
}

// no-op stubs to fulfill the Logger interface
func (log *TestLogger) SetExternalHandler(_ ExternalHandler) {}
<commit_msg>log on test failure<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package logger

import (
	"fmt"
	"runtime"
	"strings"
	"sync/atomic"
	"time"

	logging "github.com/keybase/go-logging"

	"golang.org/x/net/context"
)

// TestLogBackend is an interface for logging to a test object (i.e.,
// a *testing.T).  We define this in order to avoid pulling in the
// "testing" package in exported code.
type TestLogBackend interface {
	Error(args ...interface{})
	Errorf(format string, args ...interface{})
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Log(args ...interface{})
	Logf(format string, args ...interface{})
	Failed() bool
	Name() string
}

// TestLogger is a Logger that writes to a TestLogBackend.  All
// messages except Fatal are printed using Logf, to avoid failing a
// test that is trying to test an error condition.  No context tags
// are logged.
type TestLogger struct {
	log          TestLogBackend
	extraDepth   int
	failReported int32 // 1 is true
}

func NewTestLogger(log TestLogBackend) *TestLogger {
	return &TestLogger{log: log}
}

// Verify TestLogger fully implements the Logger interface.
var _ Logger = (*TestLogger)(nil)

// ctx can be `nil`
func (log *TestLogger) common(ctx context.Context, lvl logging.Level, useFatal bool, fmts string, arg ...interface{}) {
	if log.log.Failed() {
		failReported := atomic.SwapInt32(&(log.failReported), 1)
		if failReported == 0 {
			log.log.Logf("TEST FAILED: %s", log.log.Name())
		}
	}

	if ctx != nil {
		if useFatal {
			log.log.Fatalf(prepareString(ctx,
				log.prefixCaller(log.extraDepth, lvl, fmts)), arg...)
		} else {
			log.log.Logf(prepareString(ctx,
				log.prefixCaller(log.extraDepth, lvl, fmts)), arg...)
		}
	} else {
		if useFatal {
			log.log.Fatalf(log.prefixCaller(log.extraDepth, lvl, fmts), arg...)
		} else {
			log.log.Logf(log.prefixCaller(log.extraDepth, lvl, fmts), arg...)
		}
	}
}

func (log *TestLogger) prefixCaller(extraDepth int, lvl logging.Level, fmts string) string {
	// The testing library doesn't let us control the stack depth,
	// and it always prints out its own prefix, so use \r to clear
	// it out (at least on a terminal) and do our own formatting.
	_, file, line, _ := runtime.Caller(3 + extraDepth)
	elements := strings.Split(file, "/")
	failed := ""
	if log.log.Failed() {
		failed = "[X] "
	}

	return fmt.Sprintf("\r%s %s%s:%d: [%.1s] %s", time.Now().Format("2006-01-02 15:04:05.00000"),
		failed, elements[len(elements)-1], line, lvl, fmts)
}

func (log *TestLogger) Debug(fmts string, arg ...interface{}) {
	log.common(nil, logging.INFO, false, fmts, arg...)
}

func (log *TestLogger) CDebugf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.DEBUG, false, fmts, arg...)
}

func (log *TestLogger) Info(fmts string, arg ...interface{}) {
	log.common(nil, logging.INFO, false, fmts, arg...)
}

func (log *TestLogger) CInfof(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.INFO, false, fmts, arg...)
}

func (log *TestLogger) Notice(fmts string, arg ...interface{}) {
	log.common(nil, logging.NOTICE, false, fmts, arg...)
}

func (log *TestLogger) CNoticef(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.NOTICE, false, fmts, arg...)
}

func (log *TestLogger) Warning(fmts string, arg ...interface{}) {
	log.common(nil, logging.WARNING, false, fmts, arg...)
}

func (log *TestLogger) CWarningf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.WARNING, false, fmts, arg...)
}

func (log *TestLogger) Error(fmts string, arg ...interface{}) {
	log.common(nil, logging.ERROR, false, fmts, arg...)
}

func (log *TestLogger) Errorf(fmts string, arg ...interface{}) {
	log.common(nil, logging.ERROR, false, fmts, arg...)
}

func (log *TestLogger) CErrorf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.ERROR, false, fmts, arg...)
}

func (log *TestLogger) Critical(fmts string, arg ...interface{}) {
	log.common(nil, logging.CRITICAL, false, fmts, arg...)
}

func (log *TestLogger) CCriticalf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.CRITICAL, false, fmts, arg...)
}

func (log *TestLogger) Fatalf(fmts string, arg ...interface{}) {
	log.common(nil, logging.CRITICAL, true, fmts, arg...)
}

func (log *TestLogger) CFatalf(ctx context.Context, fmts string,
	arg ...interface{}) {
	log.common(ctx, logging.CRITICAL, true, fmts, arg...)
}

func (log *TestLogger) Profile(fmts string, arg ...interface{}) {
	log.common(nil, logging.CRITICAL, false, fmts, arg...)
}

func (log *TestLogger) Configure(style string, debug bool, filename string) {
	// no-op
}

func (log *TestLogger) RotateLogFile() error {
	// no-op
	return nil
}

func (log *TestLogger) CloneWithAddedDepth(depth int) Logger {
	clone := *log
	clone.extraDepth += depth
	return &clone
}

// no-op stubs to fulfill the Logger interface
func (log *TestLogger) SetExternalHandler(_ ExternalHandler) {}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package servenv contains functionality that is common for all
// Vitess server programs.  It defines and initializes command line
// flags that control the runtime environment.
//
// After a server program has called flag.Parse, it needs to call
// env.Init to make env use the command line variables to initialize
// the environment. It also needs to call env.Close before exiting.
//
// Note: If you need to plug in any custom initialization/cleanup for
// a vitess distribution, register them using OnInit and onClose. A
// clean way of achieving that is adding to this package a file with
// an init() function that registers the hooks.
package servenv

import (
	"flag"
	"net/url"
	"os"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	// register the HTTP handlers for profiling
	_ "net/http/pprof"

	"vitess.io/vitess/go/event"
	"vitess.io/vitess/go/netutil"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/vt/log"

	// register the proper init and shutdown hooks for logging
	_ "vitess.io/vitess/go/vt/logutil"
)

var (
	// Port is part of the flags used when calling RegisterDefaultFlags.
	Port *int

	// Flags to alter the behavior of the library.
	lameduckPeriod       = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping")
	onTermTimeout        = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping")
	memProfileRate       = flag.Int("mem-profile-rate", 512*1024, "profile every n bytes allocated")
	mutexProfileFraction = flag.Int("mutex-profile-fraction", 0, "profile every n mutex contention events (see runtime.SetMutexProfileFraction)")

	// mutex used to protect the Init function
	mu sync.Mutex

	onInitHooks     event.Hooks
	onTermHooks     event.Hooks
	onTermSyncHooks event.Hooks
	onRunHooks      event.Hooks
	inited          bool

	// ListeningURL is filled in when calling Run, contains the server URL.
	ListeningURL url.URL
)

// Init is the first phase of the server startup.
func Init() {
	mu.Lock()
	defer mu.Unlock()
	if inited {
		log.Fatal("servenv.Init called second time")
	}
	inited = true

	// Once you run as root, you pretty much destroy the chances of a
	// non-privileged user starting the program correctly.
	if uid := os.Getuid(); uid == 0 {
		log.Exitf("servenv.Init: running this as root makes no sense")
	}

	runtime.MemProfileRate = *memProfileRate

	if *mutexProfileFraction != 0 {
		log.Infof("setting mutex profile fraction to %v", *mutexProfileFraction)
		runtime.SetMutexProfileFraction(*mutexProfileFraction)
	}

	// We used to set this limit directly, but you pretty much have to
	// use a root account to allow increasing a limit reliably. Dropping
	// privileges is also tricky. The best strategy is to make a shell
	// script set up the limits as root and switch users before starting
	// the server.
	fdLimit := &syscall.Rlimit{}
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, fdLimit); err != nil {
		log.Errorf("max-open-fds failed: %v", err)
	}
	fdl := stats.NewGauge("MaxFds", "File descriptor limit")
	fdl.Set(int64(fdLimit.Cur))

	onInitHooks.Fire()
}

func populateListeningURL(port int32) {
	host, err := netutil.FullyQualifiedHostname()
	if err != nil {
		host, err = os.Hostname()
		if err != nil {
			log.Exitf("os.Hostname() failed: %v", err)
		}
	}
	ListeningURL = url.URL{
		Scheme: "http",
		Host:   netutil.JoinHostPort(host, port),
		Path:   "/",
	}
}

// OnInit registers f to be run at the beginning of the app
// lifecycle. It should be called in an init() function.
func OnInit(f func()) {
	onInitHooks.Add(f)
}

// OnTerm registers a function to be run when the process receives a SIGTERM.
// This allows the program to change its behavior during the lameduck period.
//
// All hooks are run in parallel, and there is no guarantee that the process
// will wait for them to finish before dying when the lameduck period expires.
//
// See also: OnTermSync
func OnTerm(f func()) {
	onTermHooks.Add(f)
}

// OnTermSync registers a function to be run when the process receives SIGTERM.
// This allows the program to change its behavior during the lameduck period.
//
// All hooks are run in parallel, and the process will do its best to wait
// (up to -onterm_timeout) for all of them to finish before dying.
//
// See also: OnTerm
func OnTermSync(f func()) {
	onTermSyncHooks.Add(f)
}

// fireOnTermSyncHooks returns true iff all the hooks finish before the timeout.
func fireOnTermSyncHooks(timeout time.Duration) bool {
	log.Infof("Firing synchronous OnTermSync hooks and waiting up to %v for them", timeout)

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	done := make(chan struct{})
	go func() {
		onTermSyncHooks.Fire()
		close(done)
	}()

	select {
	case <-done:
		log.Infof("OnTermSync hooks finished")
		return true
	case <-timer.C:
		log.Infof("OnTermSync hooks timed out")
		return false
	}
}

// OnRun registers f to be run right at the beginning of Run. All
// hooks are run in parallel.
func OnRun(f func()) {
	onRunHooks.Add(f)
}

// FireRunHooks fires the hooks registered by OnHook.
// Use this in a non-server to run the hooks registered
// by servenv.OnRun().
func FireRunHooks() {
	onRunHooks.Fire()
}

// RegisterDefaultFlags registers the default flags for
// listening to a given port for standard connections.
// If calling this, then call RunDefault()
func RegisterDefaultFlags() {
	Port = flag.Int("port", 0, "port for the server")
}

// RunDefault calls Run() with the parameters from the flags.
func RunDefault() {
	Run(*Port)
}

// ParseFlags initializes flags and handles the common case when no positional
// arguments are expected.
func ParseFlags(cmd string) {
	flag.Parse()

	if *Version {
		AppVersion.Print()
		os.Exit(0)
	}

	args := flag.Args()
	if len(args) > 0 {
		flag.Usage()
		log.Exitf("%s doesn't take any positional arguments, got '%s'", cmd, strings.Join(args, " "))
	}
}

// ParseFlagsWithArgs initializes flags and returns the positional arguments
func ParseFlagsWithArgs(cmd string) []string {
	flag.Parse()

	if *Version {
		AppVersion.Print()
		os.Exit(0)
	}

	args := flag.Args()
	if len(args) == 0 {
		log.Exitf("%s expected at least one positional argument", cmd)
	}

	return args
}
<commit_msg>go/vt/servenv/servenv.go:  ParseFlags():  Always print version; when `-version` is passed, exit immediately after<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package servenv contains functionality that is common for all
// Vitess server programs.  It defines and initializes command line
// flags that control the runtime environment.
//
// After a server program has called flag.Parse, it needs to call
// env.Init to make env use the command line variables to initialize
// the environment. It also needs to call env.Close before exiting.
//
// Note: If you need to plug in any custom initialization/cleanup for
// a vitess distribution, register them using OnInit and onClose. A
// clean way of achieving that is adding to this package a file with
// an init() function that registers the hooks.
package servenv

import (
	"flag"
	"net/url"
	"os"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	// register the HTTP handlers for profiling
	_ "net/http/pprof"

	"vitess.io/vitess/go/event"
	"vitess.io/vitess/go/netutil"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/vt/log"

	// register the proper init and shutdown hooks for logging
	_ "vitess.io/vitess/go/vt/logutil"
)

var (
	// Port is part of the flags used when calling RegisterDefaultFlags.
	Port *int

	// Flags to alter the behavior of the library.
	lameduckPeriod       = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping")
	onTermTimeout        = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping")
	memProfileRate       = flag.Int("mem-profile-rate", 512*1024, "profile every n bytes allocated")
	mutexProfileFraction = flag.Int("mutex-profile-fraction", 0, "profile every n mutex contention events (see runtime.SetMutexProfileFraction)")

	// mutex used to protect the Init function
	mu sync.Mutex

	onInitHooks     event.Hooks
	onTermHooks     event.Hooks
	onTermSyncHooks event.Hooks
	onRunHooks      event.Hooks
	inited          bool

	// ListeningURL is filled in when calling Run, contains the server URL.
	ListeningURL url.URL
)

// Init is the first phase of the server startup.
func Init() {
	mu.Lock()
	defer mu.Unlock()
	if inited {
		log.Fatal("servenv.Init called second time")
	}
	inited = true

	// Once you run as root, you pretty much destroy the chances of a
	// non-privileged user starting the program correctly.
	if uid := os.Getuid(); uid == 0 {
		log.Exitf("servenv.Init: running this as root makes no sense")
	}

	runtime.MemProfileRate = *memProfileRate

	if *mutexProfileFraction != 0 {
		log.Infof("setting mutex profile fraction to %v", *mutexProfileFraction)
		runtime.SetMutexProfileFraction(*mutexProfileFraction)
	}

	// We used to set this limit directly, but you pretty much have to
	// use a root account to allow increasing a limit reliably. Dropping
	// privileges is also tricky. The best strategy is to make a shell
	// script set up the limits as root and switch users before starting
	// the server.
	fdLimit := &syscall.Rlimit{}
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, fdLimit); err != nil {
		log.Errorf("max-open-fds failed: %v", err)
	}
	fdl := stats.NewGauge("MaxFds", "File descriptor limit")
	fdl.Set(int64(fdLimit.Cur))

	onInitHooks.Fire()
}

func populateListeningURL(port int32) {
	host, err := netutil.FullyQualifiedHostname()
	if err != nil {
		host, err = os.Hostname()
		if err != nil {
			log.Exitf("os.Hostname() failed: %v", err)
		}
	}
	ListeningURL = url.URL{
		Scheme: "http",
		Host:   netutil.JoinHostPort(host, port),
		Path:   "/",
	}
}

// OnInit registers f to be run at the beginning of the app
// lifecycle. It should be called in an init() function.
func OnInit(f func()) {
	onInitHooks.Add(f)
}

// OnTerm registers a function to be run when the process receives a SIGTERM.
// This allows the program to change its behavior during the lameduck period.
//
// All hooks are run in parallel, and there is no guarantee that the process
// will wait for them to finish before dying when the lameduck period expires.
//
// See also: OnTermSync
func OnTerm(f func()) {
	onTermHooks.Add(f)
}

// OnTermSync registers a function to be run when the process receives SIGTERM.
// This allows the program to change its behavior during the lameduck period.
//
// All hooks are run in parallel, and the process will do its best to wait
// (up to -onterm_timeout) for all of them to finish before dying.
//
// See also: OnTerm
func OnTermSync(f func()) {
	onTermSyncHooks.Add(f)
}

// fireOnTermSyncHooks returns true iff all the hooks finish before the timeout.
func fireOnTermSyncHooks(timeout time.Duration) bool {
	log.Infof("Firing synchronous OnTermSync hooks and waiting up to %v for them", timeout)

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	done := make(chan struct{})
	go func() {
		onTermSyncHooks.Fire()
		close(done)
	}()

	select {
	case <-done:
		log.Infof("OnTermSync hooks finished")
		return true
	case <-timer.C:
		log.Infof("OnTermSync hooks timed out")
		return false
	}
}

// OnRun registers f to be run right at the beginning of Run. All
// hooks are run in parallel.
func OnRun(f func()) {
	onRunHooks.Add(f)
}

// FireRunHooks fires the hooks registered by OnHook.
// Use this in a non-server to run the hooks registered
// by servenv.OnRun().
func FireRunHooks() {
	onRunHooks.Fire()
}

// RegisterDefaultFlags registers the default flags for
// listening to a given port for standard connections.
// If calling this, then call RunDefault()
func RegisterDefaultFlags() {
	Port = flag.Int("port", 0, "port for the server")
}

// RunDefault calls Run() with the parameters from the flags.
func RunDefault() {
	Run(*Port)
}

// ParseFlags initializes flags and handles the common case when no positional
// arguments are expected.
func ParseFlags(cmd string) {
	flag.Parse()

	AppVersion.Print()
	if *Version {
		os.Exit(0)
	}

	args := flag.Args()
	if len(args) > 0 {
		flag.Usage()
		log.Exitf("%s doesn't take any positional arguments, got '%s'", cmd, strings.Join(args, " "))
	}
}

// ParseFlagsWithArgs initializes flags and returns the positional arguments
func ParseFlagsWithArgs(cmd string) []string {
	flag.Parse()

	if *Version {
		AppVersion.Print()
		os.Exit(0)
	}

	args := flag.Args()
	if len(args) == 0 {
		log.Exitf("%s expected at least one positional argument", cmd)
	}

	return args
}
<|endoftext|>


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

package loader_test

// This file enumerates all packages beneath $GOROOT, loads them, plus
// their external tests if any, runs the type checker on them, and
// prints some summary information.
//
// Run test with GOMAXPROCS=8.

import (
	"fmt"
	"go/ast"
	"go/token"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"code.google.com/p/go.tools/go/loader"
	"code.google.com/p/go.tools/go/types"
)

func allPackages() []string {
	var pkgs []string
	root := filepath.Join(runtime.GOROOT(), "src/pkg") + string(os.PathSeparator)
	filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		// Prune the search if we encounter any of these names:
		switch filepath.Base(path) {
		case "testdata", ".hg":
			return filepath.SkipDir
		}
		if info.IsDir() {
			pkg := strings.TrimPrefix(path, root)
			switch pkg {
			case "builtin", "pkg":
				return filepath.SkipDir // skip these subtrees
			case "":
				return nil // ignore root of tree
			}
			pkgs = append(pkgs, pkg)
		}

		return nil
	})
	return pkgs
}

func TestStdlib(t *testing.T) {
	runtime.GC()
	t0 := time.Now()
	var memstats runtime.MemStats
	runtime.ReadMemStats(&memstats)
	alloc := memstats.Alloc

	// Load, parse and type-check the program.
	var conf loader.Config
	for _, path := range allPackages() {
		if err := conf.ImportWithTests(path); err != nil {
			t.Error(err)
		}
	}

	prog, err := conf.Load()
	if err != nil {
		t.Fatalf("Load failed: %v", err)
	}

	t1 := time.Now()
	runtime.GC()
	runtime.ReadMemStats(&memstats)

	numPkgs := len(prog.AllPackages)
	if want := 206; numPkgs < want {
		t.Errorf("Loaded only %d packages, want at least %d", numPkgs, want)
	}

	// Dump package members.
	if false {
		for pkg := range prog.AllPackages {
			fmt.Printf("Package %s:\n", pkg.Path())
			scope := pkg.Scope()
			for _, name := range scope.Names() {
				if ast.IsExported(name) {
					fmt.Printf("\t%s\n", types.ObjectString(pkg, scope.Lookup(name)))
				}
			}
			fmt.Println()
		}
	}

	// Check that Test functions for io/ioutil, regexp and
	// compress/bzip2 are all simultaneously present.
	// (The apparent cycle formed when augmenting all three of
	// these packages by their tests was the original motivation
	// for reporting b/7114.)
	//
	// compress/bzip2.TestBitReader in bzip2_test.go    imports io/ioutil
	// io/ioutil.TestTempFile       in tempfile_test.go imports regexp
	// regexp.TestRE2Search         in exec_test.go     imports compress/bzip2
	for _, test := range []struct{ pkg, fn string }{
		{"io/ioutil", "TestTempFile"},
		{"regexp", "TestRE2Search"},
		{"compress/bzip2", "TestBitReader"},
	} {
		info := prog.Imported[test.pkg]
		if info == nil {
			t.Errorf("failed to load package %q", test.pkg)
			continue
		}
		obj, _ := info.Pkg.Scope().Lookup(test.fn).(*types.Func)
		if obj == nil {
			t.Errorf("package %q has no func %q", test.pkg, test.fn)
			continue
		}
	}

	// Dump some statistics.

	// determine line count
	var lineCount int
	prog.Fset.Iterate(func(f *token.File) bool {
		lineCount += f.LineCount()
		return true
	})

	t.Log("GOMAXPROCS:           ", runtime.GOMAXPROCS(0))
	t.Log("#Source lines:        ", lineCount)
	t.Log("Load/parse/typecheck: ", t1.Sub(t0))
	t.Log("#MB:                  ", int64(memstats.Alloc-alloc)/1000000)
}
<commit_msg>go.tools/go/loader: fix broken test.<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package loader_test

// This file enumerates all packages beneath $GOROOT, loads them, plus
// their external tests if any, runs the type checker on them, and
// prints some summary information.
//
// Run test with GOMAXPROCS=8.

import (
	"fmt"
	"go/ast"
	"go/token"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"code.google.com/p/go.tools/go/loader"
	"code.google.com/p/go.tools/go/types"
)

func allPackages() []string {
	var pkgs []string
	root := filepath.Join(runtime.GOROOT(), "src/pkg") + string(os.PathSeparator)
	filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		// Prune the search if we encounter any of these names:
		switch filepath.Base(path) {
		case "testdata", ".hg":
			return filepath.SkipDir
		}
		if info.IsDir() {
			pkg := strings.TrimPrefix(path, root)
			switch pkg {
			case "builtin", "pkg":
				return filepath.SkipDir // skip these subtrees
			case "":
				return nil // ignore root of tree
			}
			pkgs = append(pkgs, pkg)
		}

		return nil
	})
	return pkgs
}

func TestStdlib(t *testing.T) {
	runtime.GC()
	t0 := time.Now()
	var memstats runtime.MemStats
	runtime.ReadMemStats(&memstats)
	alloc := memstats.Alloc

	// Load, parse and type-check the program.
	var conf loader.Config
	for _, path := range allPackages() {
		if err := conf.ImportWithTests(path); err != nil {
			t.Error(err)
		}
	}

	prog, err := conf.Load()
	if err != nil {
		t.Fatalf("Load failed: %v", err)
	}

	t1 := time.Now()
	runtime.GC()
	runtime.ReadMemStats(&memstats)

	numPkgs := len(prog.AllPackages)
	if want := 205; numPkgs < want {
		t.Errorf("Loaded only %d packages, want at least %d", numPkgs, want)
	}

	// Dump package members.
	if false {
		for pkg := range prog.AllPackages {
			fmt.Printf("Package %s:\n", pkg.Path())
			scope := pkg.Scope()
			for _, name := range scope.Names() {
				if ast.IsExported(name) {
					fmt.Printf("\t%s\n", types.ObjectString(pkg, scope.Lookup(name)))
				}
			}
			fmt.Println()
		}
	}

	// Check that Test functions for io/ioutil, regexp and
	// compress/bzip2 are all simultaneously present.
	// (The apparent cycle formed when augmenting all three of
	// these packages by their tests was the original motivation
	// for reporting b/7114.)
	//
	// compress/bzip2.TestBitReader in bzip2_test.go    imports io/ioutil
	// io/ioutil.TestTempFile       in tempfile_test.go imports regexp
	// regexp.TestRE2Search         in exec_test.go     imports compress/bzip2
	for _, test := range []struct{ pkg, fn string }{
		{"io/ioutil", "TestTempFile"},
		{"regexp", "TestRE2Search"},
		{"compress/bzip2", "TestBitReader"},
	} {
		info := prog.Imported[test.pkg]
		if info == nil {
			t.Errorf("failed to load package %q", test.pkg)
			continue
		}
		obj, _ := info.Pkg.Scope().Lookup(test.fn).(*types.Func)
		if obj == nil {
			t.Errorf("package %q has no func %q", test.pkg, test.fn)
			continue
		}
	}

	// Dump some statistics.

	// determine line count
	var lineCount int
	prog.Fset.Iterate(func(f *token.File) bool {
		lineCount += f.LineCount()
		return true
	})

	t.Log("GOMAXPROCS:           ", runtime.GOMAXPROCS(0))
	t.Log("#Source lines:        ", lineCount)
	t.Log("Load/parse/typecheck: ", t1.Sub(t0))
	t.Log("#MB:                  ", int64(memstats.Alloc-alloc)/1000000)
}
<|endoftext|>