Golang gomock.NewController函数代码示例

func ProtobufDecoderSpec(c gospec.Context) {
	t := &ts.SimpleT{}
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	msg := ts.GetTestMessage()
	config := NewPipelineConfig(nil) // Initializes globals.

	c.Specify("A ProtobufDecoder", func() {
		encoded, err := proto.Marshal(msg)
		c.Assume(err, gs.IsNil)
		pack := NewPipelinePack(config.inputRecycleChan)
		decoder := new(ProtobufDecoder)
		decoder.sampleDenominator = 1000 // Since we don't call decoder.Init().

		c.Specify("decodes a protobuf message", func() {
			pack.MsgBytes = encoded
			_, err := decoder.Decode(pack)
			c.Expect(err, gs.IsNil)
			c.Expect(pack.Message, gs.Equals, msg)
			v, ok := pack.Message.GetFieldValue("foo")
			c.Expect(ok, gs.IsTrue)
			c.Expect(v, gs.Equals, "bar")
		})

		c.Specify("returns an error for bunk encoding", func() {
			bunk := append([]byte{0, 0, 0}, encoded...)
			pack.MsgBytes = bunk
			_, err := decoder.Decode(pack)
			c.Expect(err, gs.Not(gs.IsNil))
		})
	})
}

func TestRun(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	mockOR := pipelinemock.NewMockOutputRunner(mockCtrl)
	mockPH := pipelinemock.NewMockPluginHelper(mockCtrl)
	mockFirehose := NewMockRecordPutter(mockCtrl)

	firehoseOutput := FirehoseOutput{
		client: mockFirehose,
	}

	testChan := make(chan *pipeline.PipelinePack)
	mockOR.EXPECT().InChan().Return(testChan)

	// Send test input through the channel
	input := `{"key":"value"}`
	go func() {
		testPack := pipeline.PipelinePack{
			Message: &message.Message{
				Payload: &input,
			},
		}

		testChan <- &testPack
		close(testChan)
	}()

	mockFirehose.EXPECT().PutRecord([]byte(input)).Return(nil)

	err := firehoseOutput.Run(mockOR, mockPH)
	assert.NoError(t, err, "did not expect err for valid Run()")
}

func StatsdInputSpec(c gs.Context) {
	t := &pipeline_ts.SimpleT{}
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	pConfig := NewPipelineConfig(nil)
	ith := new(plugins_ts.InputTestHelper)
	ith.Msg = pipeline_ts.GetTestMessage()
	ith.Pack = NewPipelinePack(pConfig.InputRecycleChan())
	ith.PackSupply = make(chan *PipelinePack, 1)

	// Specify localhost, but we're not really going to use the network
	ith.AddrStr = "localhost:55565"
	ith.ResolvedAddrStr = "127.0.0.1:55565"

	// set up mock helper, input runner, and stat accumulator
	ith.MockHelper = NewMockPluginHelper(ctrl)
	ith.MockInputRunner = NewMockInputRunner(ctrl)
	mockStatAccum := NewMockStatAccumulator(ctrl)

	c.Specify("A StatsdInput", func() {
		statsdInput := StatsdInput{}
		config := statsdInput.ConfigStruct().(*StatsdInputConfig)

		config.Address = ith.AddrStr
		err := statsdInput.Init(config)
		c.Assume(err, gs.IsNil)
		realListener := statsdInput.listener
		c.Expect(realListener.LocalAddr().String(), gs.Equals, ith.ResolvedAddrStr)
		realListener.Close()
		mockListener := pipeline_ts.NewMockConn(ctrl)
		statsdInput.listener = mockListener

		ith.MockHelper.EXPECT().StatAccumulator("StatAccumInput").Return(mockStatAccum, nil)
		mockListener.EXPECT().Close()
		mockListener.EXPECT().SetReadDeadline(gomock.Any())

		c.Specify("sends a Stat to the StatAccumulator", func() {
			statName := "sample.count"
			statVal := 303
			msg := fmt.Sprintf("%s:%d|c\n", statName, statVal)
			expected := Stat{statName, strconv.Itoa(statVal), "c", float32(1)}
			mockStatAccum.EXPECT().DropStat(expected).Return(true)
			readCall := mockListener.EXPECT().Read(make([]byte, 512))
			readCall.Return(len(msg), nil)
			readCall.Do(func(msgBytes []byte) {
				copy(msgBytes, []byte(msg))
				statsdInput.Stop()
			})
			var wg sync.WaitGroup
			wg.Add(1)
			go func() {
				err = statsdInput.Run(ith.MockInputRunner, ith.MockHelper)
				c.Expect(err, gs.IsNil)
				wg.Done()
			}()
			wg.Wait()
		})
	})
}

func createFixtures(t *testing.T) (reporter *ErrorReporter, ctrl *gomock.Controller) {
	// reporter acts as a testing.T-like object that we pass to the
	// Controller. We use it to test that the mock considered tests
	// successful or failed.
	reporter = NewErrorReporter(t)
	ctrl = gomock.NewController(reporter)
	return
}

func TestPanicOverridesExpectationChecks(t *testing.T) {
	ctrl := gomock.NewController(t)
	reporter := NewErrorReporter(t)

	reporter.assertFatal(func() {
		ctrl.RecordCall(new(Subject), "FooMethod", "1")
		defer ctrl.Finish()
		reporter.Fatalf("Intentional panic")
	})
}

func TestCreateEndpoint(t *testing.T) {
	useMockFuncs()
	defer useStdFuncs()

	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	mckEndHandler := NewMockHandler(mockCtrl)

	Convey("Should allocate update endpoints", t, func() {
		app := NewApplication()
		app.endpointTemplate = testEndpointTemplate

		Convey("Should not encrypt endpoints without a key", func() {
			app.SetTokenKey("")
			app.SetEndpointHandler(mckEndHandler)

			mckEndHandler.EXPECT().URL().Return("https://example.com")

			endpoint, err := app.CreateEndpoint("123")
			So(err, ShouldBeNil)
			So(endpoint, ShouldEqual, "https://example.com/123")
		})

		Convey("Should encrypt endpoints with a key", func() {
			app.SetTokenKey("HVozKz_n-DPopP5W877DpRKQOW_dylVf")
			app.SetEndpointHandler(mckEndHandler)

			mckEndHandler.EXPECT().URL().Return("https://example.com")

			endpoint, err := app.CreateEndpoint("456")
			So(err, ShouldBeNil)
			So(endpoint, ShouldEqual,
				"https://example.com/AAAAAAAAAAAAAAAAAAAAAGMKig==")
		})

		Convey("Should reject invalid keys", func() {
			app.SetTokenKey("lLyhlLk8qus1ky4ER8yjN5o=")
			app.SetEndpointHandler(mckEndHandler)

			_, err := app.CreateEndpoint("123")
			So(err, ShouldEqual, aes.KeySizeError(17))
		})

		Convey("Should return a relative URL without an update handler", func() {
			app.SetTokenKey("O03rpLsdafhIhJEjEJt-CgVHyqHI650oy0pZZvplKDc=")
			endpoint, err := app.CreateEndpoint("789")
			So(err, ShouldBeNil)
			So(endpoint, ShouldEqual, "/AAAAAAAAAAAAAAAAAAAAAPfdsA==")
		})
	})
}

func GeoIpDecoderSpec(c gs.Context) {
	t := &ts.SimpleT{}
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	pConfig := NewPipelineConfig(nil)
	pConfig.Globals.ShareDir = "/foo/bar/baz"

	c.Specify("A GeoIpDecoder", func() {
		decoder := new(GeoIpDecoder)
		decoder.SetPipelineConfig(pConfig)
		rec := new(geoip.GeoIPRecord)
		conf := decoder.ConfigStruct().(*GeoIpDecoderConfig)

		c.Expect(conf.DatabaseFile, gs.Equals, "/foo/bar/baz/GeoLiteCity.dat")

		supply := make(chan *PipelinePack, 1)
		pack := NewPipelinePack(supply)

		nf, _ := message.NewField("remote_host", "74.125.142.147", "")
		pack.Message.AddField(nf)

		decoder.SourceIpField = "remote_host"
		conf.SourceIpField = "remote_host"
		decoder.Init(conf)

		rec.CountryCode = "US"
		rec.CountryCode3 = "USA"
		rec.CountryName = "United States"
		rec.Region = "CA"
		rec.City = "Mountain View"
		rec.PostalCode = "94043"
		rec.Latitude = 37.4192
		rec.Longitude = -122.0574
		rec.AreaCode = 650
		rec.CharSet = 1
		rec.ContinentCode = "NA"

		c.Specify("Test GeoIpDecoder Output", func() {
			buf := decoder.GeoBuff(rec)
			nf, _ = message.NewField("geoip", buf.Bytes(), "")
			pack.Message.AddField(nf)

			b, ok := pack.Message.GetFieldValue("geoip")
			c.Expect(ok, gs.IsTrue)

			c.Expect(string(b.([]byte)), gs.Equals, `{"latitude":37.4192008972168,"longitude":-122.0574035644531,"location":[-122.0574035644531,37.4192008972168],"coordinates":["-122.0574035644531","37.4192008972168"],"countrycode":"US","countrycode3":"USA","countryname":"United States","region":"CA","city":"Mountain View","postalcode":"94043","areacode":650,"charset":1,"continentcode":"NA"}`)
		})

	})
}

func TestSocketListenConfig(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	mckLogger := NewMockLogger(mockCtrl)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).AnyTimes()
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(),
		gomock.Any(), gomock.Any()).AnyTimes()
	mckStat := NewMockStatistician(mockCtrl)
	mckListenerConfig := NewMockListenerConfig(mockCtrl)

	app := NewApplication()
	app.hostname = "example.org"
	app.SetLogger(mckLogger)
	app.SetMetrics(mckStat)

	sh := NewSocketHandler()
	sh.setApp(app)

	// Should forward Listen errors.
	listenErr := errors.New("splines not reticulated")
	mckListenerConfig.EXPECT().Listen().Return(nil, listenErr)
	if err := sh.listenWithConfig(mckListenerConfig); err != listenErr {
		t.Errorf("Wrong error: got %#v; want %#v", err, listenErr)
	}

	// Should use the wss:// scheme if UseTLS returns true.
	ml := newMockListener(netAddr{"test", "[::1]:8080"})
	gomock.InOrder(
		mckListenerConfig.EXPECT().Listen().Return(ml, nil),
		mckListenerConfig.EXPECT().UseTLS().Return(true),
		mckListenerConfig.EXPECT().GetMaxConns().Return(1),
	)
	if err := sh.listenWithConfig(mckListenerConfig); err != nil {
		t.Errorf("Error setting listener: %s", err)
	}
	if maxConns := sh.MaxConns(); maxConns != 1 {
		t.Errorf("Mismatched maximum connection count: got %d; want 1",
			maxConns)
	}
	expectedURL := "wss://example.org:8080"
	if url := sh.URL(); url != expectedURL {
		t.Errorf("Mismatched handler URL: got %q; want %q",
			url, expectedURL)
	}
}

// A more thorough test of notMatcher
func TestNotMatcher(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockMatcher := mock_matcher.NewMockMatcher(ctrl)
	notMatcher := gomock.Not(mockMatcher)

	mockMatcher.EXPECT().Matches(4).Return(true)
	if match := notMatcher.Matches(4); match {
		t.Errorf("notMatcher should not match 4")
	}

	mockMatcher.EXPECT().Matches(5).Return(false)
	if match := notMatcher.Matches(5); !match {
		t.Errorf("notMatcher should match 5")
	}
}

// decodeMessageAndVerifyOutput takes a decoder conf, message payload, and a fn -> the fn is a number of
// assertions to verify that the message after decoding is as expected.
func decodeMessageAndVerifyOutput(c gs.Context, conf *JsonDecoderConfig, payload string, fn packVerifier) {
	t := &pipeline_ts.SimpleT{}
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	// 1. Initialize test decoder
	decoder := new(JsonDecoder)
	err := decoder.Init(conf)
	c.Assume(err, gs.IsNil)
	dRunner := pipelinemock.NewMockDecoderRunner(ctrl)
	decoder.SetDecoderRunner(dRunner)

	// 2. Set payload to be tested, and decode it
	supply := make(chan *PipelinePack, 1)
	pack := NewPipelinePack(supply)
	pack.Message.SetPayload(payload)
	_, err = decoder.Decode(pack)

	// 3. Assert outcome of decoding
	fn(c, pack)
	pack.Zero()
}

func BenchmarkCreateEndpoint(b *testing.B) {
	mockCtrl := gomock.NewController(b)
	defer mockCtrl.Finish()

	app := NewApplication()
	app.endpointTemplate = testEndpointTemplate
	app.SetTokenKey("")
	mckEndHandler := NewMockHandler(mockCtrl)
	mckEndHandler.EXPECT().URL().Return("https://example.com").Times(b.N)
	app.SetEndpointHandler(mckEndHandler)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		endpoint, err := app.CreateEndpoint("123")
		if err != nil {
			b.Fatalf("Error generating update endpoint: %s", err)
		}
		expected := "https://example.com/123"
		if endpoint != expected {
			b.Fatalf("Wrong endpoint: got %q; want %q", endpoint, expected)
		}
	}
}

// TestRunWithTimestamp tests that if a TimestampColumn is provided in the config
// then the Heka message's timestamp gets added to the message with that column name.
func TestRunWithTimestamp(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	mockOR := pipelinemock.NewMockOutputRunner(mockCtrl)
	mockPH := pipelinemock.NewMockPluginHelper(mockCtrl)
	mockFirehose := NewMockRecordPutter(mockCtrl)

	firehoseOutput := FirehoseOutput{
		client:          mockFirehose,
		timestampColumn: "created",
	}

	testChan := make(chan *pipeline.PipelinePack)
	mockOR.EXPECT().InChan().Return(testChan)

	// Send test input through the channel
	input := `{}`
	timestamp := time.Date(2015, 07, 1, 13, 14, 15, 0, time.UTC).UnixNano()
	go func() {
		testPack := pipeline.PipelinePack{
			Message: &message.Message{
				Payload:   &input,
				Timestamp: &timestamp,
			},
		}

		testChan <- &testPack
		close(testChan)
	}()

	expected := `{"created":"2015-07-01 13:14:15.000"}`
	mockFirehose.EXPECT().PutRecord([]byte(expected)).Return(nil)

	err := firehoseOutput.Run(mockOR, mockPH)
	assert.NoError(t, err, "did not expect err for valid Run()")
}

func ScribbleDecoderSpec(c gs.Context) {
	t := &pipeline_ts.SimpleT{}
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	c.Specify("A ScribbleDecoder", func() {
		decoder := new(ScribbleDecoder)
		config := decoder.ConfigStruct().(*ScribbleDecoderConfig)
		myType := "myType"
		myPayload := "myPayload"
		config.MessageFields = MessageTemplate{"Type": myType, "Payload": myPayload}
		supply := make(chan *PipelinePack, 1)
		pack := NewPipelinePack(supply)

		c.Specify("sets basic values correctly", func() {
			decoder.Init(config)
			packs, err := decoder.Decode(pack)
			c.Expect(err, gs.IsNil)
			c.Expect(len(packs), gs.Equals, 1)
			c.Expect(pack.Message.GetType(), gs.Equals, myType)
			c.Expect(pack.Message.GetPayload(), gs.Equals, myPayload)
		})
	})
}

func TestLocatorReadyNotify(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	uaid := "fce61180716a40ed8e79bf5ff0ba34bc"

	mckLogger := NewMockLogger(mockCtrl)
	mckLogger.EXPECT().ShouldLog(gomock.Any()).Return(true).AnyTimes()
	mckLogger.EXPECT().Log(gomock.Any(), gomock.Any(), gomock.Any(),
		gomock.Any()).AnyTimes()

	var (
		// routerPipes maps fake peer addresses to their respective pipes. Used
		// by dialRouter to connect to peers.
		routerPipes = make(map[netAddr]*pipeListener)

		// contacts is a list of peer URLs for the locator.
		contacts []string
	)

	// Fake listener for the sender's router, used to test self-routing.
	sndRouterAddr := netAddr{"tcp", "snd-router.example.com:3000"}
	sndRouterPipe := newPipeListener()
	defer sndRouterPipe.Close()
	routerPipes[sndRouterAddr] = sndRouterPipe
	contacts = append(contacts, "http://snd-router.example.com:3000")

	// Fake listener for the receiver's router, used to test routing updates
	// to different hosts.
	recvRouterAddr := netAddr{"tcp", "recv-router.example.com:3000"}
	recvRouterPipe := newPipeListener()
	defer recvRouterPipe.Close()
	routerPipes[recvRouterAddr] = recvRouterPipe
	contacts = append(contacts, "http://recv-router.example.com:3000")

	// Fake listener for the receiver's WebSocket handler, used to accept a
	// WebSocket client connection.
	socketHandlerPipe := newPipeListener()
	defer socketHandlerPipe.Close()

	// Fake locator.
	mckLocator := NewMockLocator(mockCtrl)
	mckLocator.EXPECT().Contacts(uaid).Return(contacts, nil).Times(2)

	// Fake dialer to connect to each peer's routing listener.
	dialRouter := func(network, address string) (net.Conn, error) {
		if pipe, ok := routerPipes[netAddr{network, address}]; ok {
			return pipe.Dial(network, address)
		}
		return nil, &netErr{temporary: false, timeout: false}
	}
	// Configures a fake router for the app.
	setRouter := func(app *Application, listener net.Listener) {
		r := NewBroadcastRouter()
		r.setApp(app)
		r.setClientOptions(10, 3*time.Second, 3*time.Second) // Defaults.
		r.setClientTransport(&http.Transport{Dial: dialRouter})
		r.listenWithConfig(listenerConfig{listener: listener})
		r.maxDataLen = 4096
		r.server = newServeWaiter(&http.Server{Handler: r.ServeMux()})
		app.SetRouter(r)
	}

	// sndApp is the server broadcasting the update. The locator returns the
	// addresses of the sender and receiver to test self-routing.
	sndApp := NewApplication()
	sndApp.SetLogger(mckLogger)
	sndStat := NewMockStatistician(mockCtrl)
	sndApp.SetMetrics(sndStat)
	sndStore := NewMockStore(mockCtrl)
	sndApp.SetStore(sndStore)
	sndApp.SetLocator(mckLocator)
	// Set up a fake router for the sender.
	setRouter(sndApp, sndRouterPipe)

	// recvApp is the server receiving the update.
	recvApp := NewApplication()
	recvApp.SetLogger(mckLogger)
	recvStat := NewMockStatistician(mockCtrl)
	recvApp.SetMetrics(recvStat)
	recvStore := NewMockStore(mockCtrl)
	recvApp.SetStore(recvStore)
	// Wrap the fake locator in a type that implements ReadyNotifier.
	recvLocator := newMockReadyNotifier(mckLocator)
	recvApp.SetLocator(recvLocator)
	// Set up a fake WebSocket handler for the receiver.
	recvSocketHandler := NewSocketHandler()
	recvSocketHandler.setApp(recvApp)
	recvSocketHandler.listenWithConfig(listenerConfig{
		listener: socketHandlerPipe})
	recvSocketHandler.server = newServeWaiter(&http.Server{Handler: recvSocketHandler.ServeMux()})
	recvApp.SetSocketHandler(recvSocketHandler)
	// Set up a fake router for the receiver.
	setRouter(recvApp, recvRouterPipe)

	chid := "2b7c5c27d6224bfeaf1c158c3c57fca3"
	version := int64(2)
	data := "I'm a little teapot, short and stout."

	var wg sync.WaitGroup // Waits for the client to close.
//.........这里部分代码省略.........

func ReportSpec(c gs.Context) {
	t := new(ts.SimpleT)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	pConfig := NewPipelineConfig(nil)
	chanSize := pConfig.Globals.PluginChanSize

	checkForFields := func(c gs.Context, msg *message.Message) {
		f0Val, ok := msg.GetFieldValue(f0.GetName())
		c.Expect(ok, gs.IsTrue)
		c.Expect(f0Val.(int64), gs.Equals, f0.GetValue().(int64))
		f1Val, ok := msg.GetFieldValue(f1.GetName())
		c.Expect(ok, gs.IsTrue)
		c.Expect(f1Val.(string), gs.Equals, f1.GetValue().(string))
	}

	hasChannelData := func(msg *message.Message) (ok bool) {
		capVal, _ := msg.GetFieldValue("InChanCapacity")
		lenVal, _ := msg.GetFieldValue("InChanLength")
		var i int64
		if i, ok = capVal.(int64); !ok {
			return
		}
		if ok = (i == int64(chanSize)); !ok {
			return
		}
		if i, ok = lenVal.(int64); !ok {
			return
		}
		ok = (i == int64(0))
		return
	}

	fName := "counter"
	filter := new(CounterFilter)
	fRunner := NewFORunner(fName, filter, nil, chanSize)
	var err error
	fRunner.matcher, err = NewMatchRunner("Type == ''", "", fRunner, chanSize)
	c.Assume(err, gs.IsNil)
	fRunner.matcher.inChan = make(chan *PipelinePack, chanSize)
	leakCount := 10
	fRunner.SetLeakCount(leakCount)

	iName := "stat_accum"
	input := new(StatAccumInput)
	iRunner := NewInputRunner(iName, input, nil, false)

	c.Specify("`PopulateReportMsg`", func() {
		msg := ts.GetTestMessage()

		c.Specify("w/ a filter", func() {
			err := PopulateReportMsg(fRunner, msg)
			c.Assume(err, gs.IsNil)

			c.Specify("invokes `ReportMsg` on the filter", func() {
				checkForFields(c, msg)
			})

			c.Specify("adds the channel data", func() {
				c.Expect(hasChannelData(msg), gs.IsTrue)
			})

			c.Specify("has its leak count set properly", func() {
				leakVal, ok := msg.GetFieldValue("LeakCount")
				c.Assume(ok, gs.IsTrue)
				i, ok := leakVal.(int64)
				c.Assume(ok, gs.IsTrue)
				c.Expect(int(i), gs.Equals, leakCount)
			})
		})

		c.Specify("w/ an input", func() {
			err := PopulateReportMsg(iRunner, msg)
			c.Assume(err, gs.IsNil)

			c.Specify("invokes `ReportMsg` on the input", func() {
				checkForFields(c, msg)
			})

			c.Specify("doesn't add any channel data", func() {
				capVal, ok := msg.GetFieldValue("InChanCapacity")
				c.Expect(capVal, gs.IsNil)
				c.Expect(ok, gs.IsFalse)
				lenVal, ok := msg.GetFieldValue("InChanLength")
				c.Expect(lenVal, gs.IsNil)
				c.Expect(ok, gs.IsFalse)
			})
		})
	})

	c.Specify("PipelineConfig", func() {
		pc := NewPipelineConfig(nil)
		// Initialize all of the PipelinePacks that we'll need
		pc.reportRecycleChan <- NewPipelinePack(pc.reportRecycleChan)

		pc.FilterRunners = map[string]FilterRunner{fName: fRunner}
		pc.InputRunners = map[string]InputRunner{iName: iRunner}

		c.Specify("returns full set of accurate reports", func() {
//.........这里部分代码省略.........