Golang security.SetReadFileFn函數代碼示例

func TestMain(m *testing.M) {
	security.SetReadFileFn(securitytest.Asset)
	randutil.SeedForTests()
	serverutils.InitTestServerFactory(server.TestServerFactory)
	serverutils.InitTestClusterFactory(testcluster.TestClusterFactory)
	os.Exit(m.Run())
}

func TestQuit(t *testing.T) {
	defer leaktest.AfterTest(t)()
	c := newCLITest()
	c.Run("quit")
	// Wait until this async command stops the server.
	<-c.Stopper().IsStopped()

	// NB: if this test is ever flaky due to port reuse, we could run against
	// :0 (which however changes some of the errors we get).
	// One way of getting that is:
	//	c.TestServer.Cfg.AdvertiseAddr = "127.0.0.1:0"

	styled := func(s string) string {
		const preamble = `unable to connect or connection lost.

Please check the address and credentials such as certificates \(if attempting to
communicate with a secure cluster\).

`
		return preamble + s
	}

	for i, test := range []struct {
		cmd, expOutPattern string
	}{
		// Error returned from GRPC to internal/client (which has to pass it
		// up the stack as a roachpb.NewError(roachpb.NewSendError(.)).
		{`debug kv inc a`, styled(
			`increment failed: failed to send RPC: rpc error: code = 14 desc = grpc: the connection is unavailable`),
		},
		// Error returned directly from GRPC.
		{`quit`, styled(
			`rpc error: code = 14 desc = grpc: the connection is unavailable`),
		},
		// Going through the SQL client libraries gives a *net.OpError which
		// we also handle.
		{`zone ls`, styled(
			`dial tcp .*: getsockopt: connection refused`),
		},
		// A real error before we hit any networking ones.
		{`debug kv scan b a`, `start key must be smaller than end key`},
		{`debug kv inc`, `invalid arguments`},
	} {
		out, err := c.RunWithCapture(test.cmd)
		if err != nil {
			t.Fatal(errors.Wrap(err, strconv.Itoa(i)))
		}
		exp := test.cmd + "\n" + test.expOutPattern
		re := regexp.MustCompile(exp)
		if !re.MatchString(out) {
			t.Errorf("expected '%s' to match pattern\n%s\ngot:\n%s", test.cmd, exp, out)
		}
	}
	// Manually run the cleanup functions (intentionally only on success,
	// preserving the logs on failure).
	c.cleanupFunc()
	security.SetReadFileFn(securitytest.Asset)
}

func newCLITest(t *testing.T, insecure bool) (cliTest, error) {
	c := cliTest{}

	certsDir, err := ioutil.TempDir("", "cli-test")
	if err != nil {
		return cliTest{}, err
	}
	c.certsDir = certsDir

	// Reset the client context for each test. We don't reset the
	// pointer (because they are tied into the flags), but instead
	// overwrite the existing struct's values.
	baseCfg.InitDefaults()
	cliCtx.InitCLIDefaults()

	s, err := serverutils.StartServerRaw(base.TestServerArgs{Insecure: insecure})
	if err != nil {
		return cliTest{}, err
	}
	c.TestServer = s.(*server.TestServer)

	if insecure {
		c.cleanupFunc = func() error { return nil }
	} else {
		// Copy these assets to disk from embedded strings, so this test can
		// run from a standalone binary.
		// Disable embedded certs, or the security library will try to load
		// our real files as embedded assets.
		security.ResetReadFileFn()

		assets := []string{
			filepath.Join(security.EmbeddedCertsDir, security.EmbeddedCACert),
			filepath.Join(security.EmbeddedCertsDir, security.EmbeddedCAKey),
			filepath.Join(security.EmbeddedCertsDir, security.EmbeddedNodeCert),
			filepath.Join(security.EmbeddedCertsDir, security.EmbeddedNodeKey),
			filepath.Join(security.EmbeddedCertsDir, security.EmbeddedRootCert),
			filepath.Join(security.EmbeddedCertsDir, security.EmbeddedRootKey),
		}

		for _, a := range assets {
			securitytest.RestrictedCopy(nil, a, certsDir, filepath.Base(a))
		}

		c.cleanupFunc = func() error {
			security.SetReadFileFn(securitytest.Asset)
			return os.RemoveAll(certsDir)
		}
	}

	// Ensure that CLI error messages are logged to stdout, where they
	// can be captured.
	osStderr = os.Stdout

	return c, nil
}

func init() {
	security.SetReadFileFn(securitytest.Asset)
}

func (c cliTest) stop() {
	c.cleanupFunc()
	security.SetReadFileFn(securitytest.Asset)
	c.Stopper().Stop()
}

func TestMain(m *testing.M) {
	security.SetReadFileFn(securitytest.Asset)
	serverutils.InitTestServerFactory(server.TestServerFactory)
	os.Exit(m.Run())
}

// ResetTest sets up the test environment. In particular, it embeds the
// EmbeddedCertsDir folder and makes the tls package load from there.
func ResetTest() {
	security.SetReadFileFn(securitytest.Asset)
}