feat(auth): fail-fast for unauthenticated requests before MCP call
- Add token validation in executeInvocation to reject requests early - Return clear Chinese error message: '未登录,请先执行 dws auth login' - Include actionable hint and suggested command in error response - Avoid wasting network resources on HTTP 400 from MCP gateway Also includes: - Add DWS_PERF_TIMING env for CLI execution timing analysis - Add TimingCollector to track cmd_init, auth_token, mcp_call durations Test changes: - Add --token flag to tests that require authentication - Add TestRuntimeRunnerRejectsUnauthenticatedRequest test case
This commit is contained in:
1 parent
0257d1f084
commit
14f558facf
6 files changed
+424
-13
No files matched your search
@@ -261,7 +261,7 @@ func TestExecuteWritesRecoveryEventIDToStderrOnCapturedFailure(t *testing.T) {
|
||||
|
||||
oldArgs := os.Args
|
||||
defer func() { os.Args = oldArgs }()
|
||||
os.Args = []string{"dws", "mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`}
|
||||
os.Args = []string{"dws", "mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"}
|
||||
|
||||
stdoutR, stdoutW, err := os.Pipe()
|
||||
if err != nil {
|
||||
|
||||
+11
-2
@@ -52,7 +52,9 @@ const recoveryEventStderrPrefix = "RECOVERY_EVENT_ID="
|
||||
// Execute runs the root command and returns the process exit code.
|
||||
func Execute() int {
|
||||
totalStart := time.Now()
|
||||
timing := NewTimingCollector()
|
||||
defer func() {
|
||||
timing.PrintIfEnabled()
|
||||
if os.Getenv("DWS_PERF_DEBUG") != "" {
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] Execute total: %v\n", time.Since(totalStart))
|
||||
}
|
||||
@@ -61,12 +63,17 @@ func Execute() int {
|
||||
ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||||
defer cancel()
|
||||
|
||||
// Attach timing collector to context for use by child components
|
||||
ctx = WithTimingCollector(ctx, timing)
|
||||
|
||||
initStart := time.Now()
|
||||
recovery.ResetRuntimeState()
|
||||
engine := newPipelineEngine()
|
||||
root := NewRootCommandWithEngine(ctx, engine)
|
||||
initDuration := time.Since(initStart)
|
||||
timing.Record("cmd_init", initDuration)
|
||||
if os.Getenv("DWS_PERF_DEBUG") != "" {
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] command init: %v\n", time.Since(initStart))
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] command init: %v\n", initDuration)
|
||||
}
|
||||
|
||||
// Run PreParse handlers on raw argv before Cobra parses flags.
|
||||
@@ -76,8 +83,10 @@ func Execute() int {
|
||||
|
||||
execStart := time.Now()
|
||||
executed, err := root.ExecuteC()
|
||||
execDuration := time.Since(execStart)
|
||||
timing.Record("cobra_exec", execDuration)
|
||||
if os.Getenv("DWS_PERF_DEBUG") != "" {
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] cobra ExecuteC: %v\n", time.Since(execStart))
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] cobra ExecuteC: %v\n", execDuration)
|
||||
}
|
||||
if err != nil {
|
||||
if executed == nil {
|
||||
|
||||
+25
-4
@@ -101,7 +101,9 @@ func (r *runtimeRunner) Run(ctx context.Context, invocation executor.Invocation)
|
||||
}
|
||||
}
|
||||
|
||||
catalogStart := time.Now()
|
||||
catalog, err := r.loader.Load(ctx)
|
||||
RecordTiming(ctx, "catalog_load", time.Since(catalogStart))
|
||||
if err != nil {
|
||||
return executor.Result{}, err
|
||||
}
|
||||
@@ -132,9 +134,11 @@ func (r *runtimeRunner) executeInvocation(ctx context.Context, endpoint string,
|
||||
}
|
||||
|
||||
authStart := time.Now()
|
||||
tc := r.transport.WithAuth(r.resolveAuthToken(ctx), resolveIdentityHeaders())
|
||||
authToken := r.resolveAuthToken(ctx)
|
||||
authDuration := time.Since(authStart)
|
||||
RecordTiming(ctx, "auth_token", authDuration)
|
||||
if os.Getenv("DWS_PERF_DEBUG") != "" {
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] resolveAuthToken: %v\n", time.Since(authStart))
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] resolveAuthToken: %v\n", authDuration)
|
||||
}
|
||||
|
||||
if invocation.DryRun {
|
||||
@@ -166,10 +170,25 @@ func (r *runtimeRunner) executeInvocation(ctx context.Context, endpoint string,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Fail-fast: reject unauthenticated requests before making network calls.
|
||||
// This provides a clear error message instead of cryptic HTTP 400 from MCP.
|
||||
if strings.TrimSpace(authToken) == "" {
|
||||
return executor.Result{}, apperrors.NewAuth(
|
||||
"未登录,请先执行 dws auth login",
|
||||
apperrors.WithReason("not_authenticated"),
|
||||
apperrors.WithHint("运行 'dws auth login' 完成登录后重试"),
|
||||
apperrors.WithActions("dws auth login"),
|
||||
)
|
||||
}
|
||||
|
||||
tc := r.transport.WithAuth(authToken, resolveIdentityHeaders())
|
||||
|
||||
callStart := time.Now()
|
||||
callResult, err := tc.CallTool(ctx, endpoint, invocation.Tool, invocation.Params)
|
||||
callDuration := time.Since(callStart)
|
||||
RecordTiming(ctx, "mcp_call", callDuration)
|
||||
if os.Getenv("DWS_PERF_DEBUG") != "" {
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] MCP CallTool: %v\n", time.Since(callStart))
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] MCP CallTool: %v\n", callDuration)
|
||||
}
|
||||
if err != nil {
|
||||
captureRuntimeFailure(invocation, err, err)
|
||||
@@ -247,8 +266,10 @@ func getCachedRuntimeToken(ctx context.Context) string {
|
||||
cachedRuntimeTokenOnce.Do(func() {
|
||||
loadStart := time.Now()
|
||||
defer func() {
|
||||
loadDuration := time.Since(loadStart)
|
||||
RecordTiming(ctx, "keychain_load", loadDuration)
|
||||
if os.Getenv("DWS_PERF_DEBUG") != "" {
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] getCachedRuntimeToken (first load): %v\n", time.Since(loadStart))
|
||||
_, _ = fmt.Fprintf(os.Stderr, "[PERF] getCachedRuntimeToken (first load): %v\n", loadDuration)
|
||||
}
|
||||
}()
|
||||
|
||||
|
||||
@@ -45,7 +45,7 @@ func TestRuntimeRunnerIncludesContentScanReportWhenEnabled(t *testing.T) {
|
||||
var out bytes.Buffer
|
||||
cmd.SetOut(&out)
|
||||
cmd.SetErr(&out)
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`})
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"})
|
||||
|
||||
if err := cmd.Execute(); err != nil {
|
||||
t.Fatalf("Execute() error = %v", err)
|
||||
@@ -90,7 +90,7 @@ func TestRuntimeRunnerBlocksUnsafeContentWhenEnforced(t *testing.T) {
|
||||
cmd := NewRootCommand()
|
||||
cmd.SetOut(&bytes.Buffer{})
|
||||
cmd.SetErr(&bytes.Buffer{})
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`})
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"})
|
||||
|
||||
err := cmd.Execute()
|
||||
if err == nil {
|
||||
@@ -121,7 +121,7 @@ func TestCanonicalCommandUsesRuntimeRunnerWhenEnabled(t *testing.T) {
|
||||
var out bytes.Buffer
|
||||
cmd.SetOut(&out)
|
||||
cmd.SetErr(&out)
|
||||
cmd.SetArgs([]string{"mcp", "doc", "create_document", "--json", `{"title":"Quarterly"}`, "--yes"})
|
||||
cmd.SetArgs([]string{"mcp", "doc", "create_document", "--json", `{"title":"Quarterly"}`, "--yes", "--token", "test-token"})
|
||||
|
||||
if err := cmd.Execute(); err != nil {
|
||||
t.Fatalf("Execute() error = %v", err)
|
||||
@@ -255,6 +255,37 @@ func TestRuntimeRunnerInjectsAuthTokenFromFlag(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestRuntimeRunnerRejectsUnauthenticatedRequest verifies that requests without
|
||||
// a valid token are rejected with a clear error before making any network call.
|
||||
func TestRuntimeRunnerRejectsUnauthenticatedRequest(t *testing.T) {
|
||||
setupRuntimeCommandTest(t)
|
||||
server := mockmcp.DefaultServer()
|
||||
defer server.Close()
|
||||
|
||||
t.Setenv(cli.CatalogFixtureEnv, writeDocCatalogFixture(t, server.RemoteURL("/server/doc"), false))
|
||||
|
||||
cmd := NewRootCommand()
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd.SetOut(&stdout)
|
||||
cmd.SetErr(&stderr)
|
||||
// No --token flag, should be rejected
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`})
|
||||
|
||||
err := cmd.Execute()
|
||||
if err == nil {
|
||||
t.Fatal("Execute() error = nil, want authentication error")
|
||||
}
|
||||
|
||||
// Verify we get a clear auth error, not a cryptic HTTP 400
|
||||
errMsg := err.Error()
|
||||
if !strings.Contains(errMsg, "未登录") {
|
||||
t.Fatalf("Execute() error = %v, want error containing '未登录'", err)
|
||||
}
|
||||
if !strings.Contains(errMsg, "auth login") {
|
||||
t.Fatalf("Execute() error = %v, want error containing 'auth login'", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeRunnerFallsBackForUnavailableProduct(t *testing.T) {
|
||||
setupRuntimeCommandTest(t)
|
||||
server := mockmcp.DefaultServer()
|
||||
@@ -415,7 +446,7 @@ func TestCanonicalSensitiveToolAcceptsInteractiveConfirmation(t *testing.T) {
|
||||
cmd.SetOut(&out)
|
||||
cmd.SetErr(&errOut)
|
||||
cmd.SetIn(strings.NewReader("yes\n"))
|
||||
cmd.SetArgs([]string{"mcp", "doc", "create_document", "--json", `{"title":"Quarterly"}`})
|
||||
cmd.SetArgs([]string{"mcp", "doc", "create_document", "--json", `{"title":"Quarterly"}`, "--token", "test-token"})
|
||||
|
||||
if err := cmd.Execute(); err != nil {
|
||||
t.Fatalf("Execute() error = %v", err)
|
||||
@@ -465,7 +496,7 @@ func TestRuntimeRunnerUsesProductEndpointOverride(t *testing.T) {
|
||||
var out bytes.Buffer
|
||||
cmd.SetOut(&out)
|
||||
cmd.SetErr(&out)
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`})
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"})
|
||||
|
||||
if err := cmd.Execute(); err != nil {
|
||||
t.Fatalf("Execute() error = %v", err)
|
||||
@@ -628,7 +659,7 @@ func TestRuntimeRunnerReturnsErrorWhenMCPIsErrorTrue(t *testing.T) {
|
||||
cmd := NewRootCommand()
|
||||
cmd.SetOut(&bytes.Buffer{})
|
||||
cmd.SetErr(&bytes.Buffer{})
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`})
|
||||
cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"})
|
||||
|
||||
err := cmd.Execute()
|
||||
if err == nil {
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
// Copyright 2026 Alibaba Group
|
||||
// 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
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Environment variable to enable performance timing output.
|
||||
const PerfTimingEnv = "DWS_PERF_TIMING"
|
||||
|
||||
// timingContextKey is the context key for TimingCollector.
|
||||
type timingContextKey struct{}
|
||||
|
||||
// TimingEntry represents a single timing measurement.
|
||||
type TimingEntry struct {
|
||||
Name string
|
||||
Duration time.Duration
|
||||
Timestamp time.Time
|
||||
Seq int // insertion order
|
||||
}
|
||||
|
||||
// TimingCollector collects timing measurements for a single command execution.
|
||||
// It is safe for concurrent use.
|
||||
type TimingCollector struct {
|
||||
mu sync.Mutex
|
||||
start time.Time
|
||||
entries []TimingEntry
|
||||
seq int
|
||||
}
|
||||
|
||||
// NewTimingCollector creates a new collector with the start time set to now.
|
||||
func NewTimingCollector() *TimingCollector {
|
||||
return &TimingCollector{
|
||||
start: time.Now(),
|
||||
entries: make([]TimingEntry, 0, 16),
|
||||
}
|
||||
}
|
||||
|
||||
// Record adds a timing entry with the given name and duration.
|
||||
func (tc *TimingCollector) Record(name string, d time.Duration) {
|
||||
if tc == nil {
|
||||
return
|
||||
}
|
||||
tc.mu.Lock()
|
||||
defer tc.mu.Unlock()
|
||||
tc.entries = append(tc.entries, TimingEntry{
|
||||
Name: name,
|
||||
Duration: d,
|
||||
Timestamp: time.Now(),
|
||||
Seq: tc.seq,
|
||||
})
|
||||
tc.seq++
|
||||
}
|
||||
|
||||
// StartTimer returns a function that, when called, records the elapsed time
|
||||
// since StartTimer was called. This is convenient for defer usage:
|
||||
//
|
||||
// defer tc.StartTimer("operation")()
|
||||
func (tc *TimingCollector) StartTimer(name string) func() {
|
||||
if tc == nil {
|
||||
return func() {}
|
||||
}
|
||||
start := time.Now()
|
||||
return func() {
|
||||
tc.Record(name, time.Since(start))
|
||||
}
|
||||
}
|
||||
|
||||
// Total returns the total elapsed time since the collector was created.
|
||||
func (tc *TimingCollector) Total() time.Duration {
|
||||
if tc == nil {
|
||||
return 0
|
||||
}
|
||||
return time.Since(tc.start)
|
||||
}
|
||||
|
||||
// Entries returns a copy of all recorded entries in insertion order.
|
||||
func (tc *TimingCollector) Entries() []TimingEntry {
|
||||
if tc == nil {
|
||||
return nil
|
||||
}
|
||||
tc.mu.Lock()
|
||||
defer tc.mu.Unlock()
|
||||
result := make([]TimingEntry, len(tc.entries))
|
||||
copy(result, tc.entries)
|
||||
sort.Slice(result, func(i, j int) bool {
|
||||
return result[i].Seq < result[j].Seq
|
||||
})
|
||||
return result
|
||||
}
|
||||
|
||||
// Print writes a summary of all timing entries to the given writer.
|
||||
func (tc *TimingCollector) Print(w io.Writer) {
|
||||
if tc == nil || w == nil {
|
||||
return
|
||||
}
|
||||
entries := tc.Entries()
|
||||
if len(entries) == 0 {
|
||||
fmt.Fprintf(w, "\n[Timing] Total: %v (no detailed entries)\n", tc.Total().Truncate(time.Millisecond))
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Fprintln(w)
|
||||
fmt.Fprintln(w, "[Timing] Execution breakdown:")
|
||||
for _, e := range entries {
|
||||
fmt.Fprintf(w, " %-30s %v\n", e.Name, e.Duration.Truncate(time.Millisecond))
|
||||
}
|
||||
fmt.Fprintf(w, " %-30s %v\n", "──────────────────────────────", "──────────")
|
||||
fmt.Fprintf(w, " %-30s %v\n", "Total", tc.Total().Truncate(time.Millisecond))
|
||||
}
|
||||
|
||||
// PrintIfEnabled prints timing info to stderr if DWS_PERF_TIMING is set.
|
||||
func (tc *TimingCollector) PrintIfEnabled() {
|
||||
if tc == nil {
|
||||
return
|
||||
}
|
||||
if os.Getenv(PerfTimingEnv) == "" {
|
||||
return
|
||||
}
|
||||
tc.Print(os.Stderr)
|
||||
}
|
||||
|
||||
// WithTimingCollector returns a new context with the TimingCollector attached.
|
||||
func WithTimingCollector(ctx context.Context, tc *TimingCollector) context.Context {
|
||||
return context.WithValue(ctx, timingContextKey{}, tc)
|
||||
}
|
||||
|
||||
// TimingCollectorFromContext extracts the TimingCollector from context, or nil.
|
||||
func TimingCollectorFromContext(ctx context.Context) *TimingCollector {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
tc, _ := ctx.Value(timingContextKey{}).(*TimingCollector)
|
||||
return tc
|
||||
}
|
||||
|
||||
// RecordTiming is a convenience function to record timing to the collector in context.
|
||||
func RecordTiming(ctx context.Context, name string, d time.Duration) {
|
||||
if tc := TimingCollectorFromContext(ctx); tc != nil {
|
||||
tc.Record(name, d)
|
||||
}
|
||||
}
|
||||
|
||||
// StartTiming is a convenience function that returns a stop function for defer usage.
|
||||
// Example:
|
||||
//
|
||||
// defer StartTiming(ctx, "operation")()
|
||||
func StartTiming(ctx context.Context, name string) func() {
|
||||
tc := TimingCollectorFromContext(ctx)
|
||||
if tc == nil {
|
||||
return func() {}
|
||||
}
|
||||
return tc.StartTimer(name)
|
||||
}
|
||||
|
||||
// IsPerfTimingEnabled returns true if performance timing output is enabled.
|
||||
func IsPerfTimingEnabled() bool {
|
||||
return os.Getenv(PerfTimingEnv) != ""
|
||||
}
|
||||
@@ -0,0 +1,173 @@
|
||||
// Copyright 2026 Alibaba Group
|
||||
// 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
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestTimingCollector_Basic(t *testing.T) {
|
||||
tc := NewTimingCollector()
|
||||
if tc == nil {
|
||||
t.Fatal("NewTimingCollector returned nil")
|
||||
}
|
||||
|
||||
// Record some timings
|
||||
tc.Record("op1", 10*time.Millisecond)
|
||||
tc.Record("op2", 20*time.Millisecond)
|
||||
|
||||
entries := tc.Entries()
|
||||
if len(entries) != 2 {
|
||||
t.Errorf("expected 2 entries, got %d", len(entries))
|
||||
}
|
||||
|
||||
// Check ordering
|
||||
if entries[0].Name != "op1" {
|
||||
t.Errorf("expected first entry to be 'op1', got %q", entries[0].Name)
|
||||
}
|
||||
if entries[1].Name != "op2" {
|
||||
t.Errorf("expected second entry to be 'op2', got %q", entries[1].Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimingCollector_StartTimer(t *testing.T) {
|
||||
tc := NewTimingCollector()
|
||||
|
||||
stop := tc.StartTimer("timed_op")
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
stop()
|
||||
|
||||
entries := tc.Entries()
|
||||
if len(entries) != 1 {
|
||||
t.Fatalf("expected 1 entry, got %d", len(entries))
|
||||
}
|
||||
if entries[0].Name != "timed_op" {
|
||||
t.Errorf("expected entry name 'timed_op', got %q", entries[0].Name)
|
||||
}
|
||||
if entries[0].Duration < 5*time.Millisecond {
|
||||
t.Errorf("expected duration >= 5ms, got %v", entries[0].Duration)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimingCollector_NilSafe(t *testing.T) {
|
||||
var tc *TimingCollector
|
||||
|
||||
// Should not panic on nil collector
|
||||
tc.Record("op", 10*time.Millisecond)
|
||||
stop := tc.StartTimer("op")
|
||||
stop()
|
||||
_ = tc.Total()
|
||||
_ = tc.Entries()
|
||||
tc.Print(nil)
|
||||
tc.PrintIfEnabled()
|
||||
}
|
||||
|
||||
func TestTimingCollector_Print(t *testing.T) {
|
||||
tc := NewTimingCollector()
|
||||
tc.Record("auth_token", 44*time.Millisecond)
|
||||
tc.Record("mcp_call", 150*time.Millisecond)
|
||||
|
||||
var buf bytes.Buffer
|
||||
tc.Print(&buf)
|
||||
|
||||
output := buf.String()
|
||||
if !strings.Contains(output, "[Timing]") {
|
||||
t.Error("output should contain [Timing] header")
|
||||
}
|
||||
if !strings.Contains(output, "auth_token") {
|
||||
t.Error("output should contain 'auth_token'")
|
||||
}
|
||||
if !strings.Contains(output, "mcp_call") {
|
||||
t.Error("output should contain 'mcp_call'")
|
||||
}
|
||||
if !strings.Contains(output, "Total") {
|
||||
t.Error("output should contain 'Total'")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimingCollector_PrintIfEnabled(t *testing.T) {
|
||||
// Set environment variable
|
||||
os.Setenv(PerfTimingEnv, "1")
|
||||
defer os.Unsetenv(PerfTimingEnv)
|
||||
|
||||
tc := NewTimingCollector()
|
||||
tc.Record("test_op", 10*time.Millisecond)
|
||||
|
||||
// This should not panic and should print to stderr
|
||||
tc.PrintIfEnabled()
|
||||
}
|
||||
|
||||
func TestTimingCollector_ContextIntegration(t *testing.T) {
|
||||
tc := NewTimingCollector()
|
||||
ctx := WithTimingCollector(context.Background(), tc)
|
||||
|
||||
// Retrieve from context
|
||||
retrieved := TimingCollectorFromContext(ctx)
|
||||
if retrieved != tc {
|
||||
t.Error("TimingCollectorFromContext should return the same collector")
|
||||
}
|
||||
|
||||
// Use convenience functions
|
||||
RecordTiming(ctx, "ctx_op", 30*time.Millisecond)
|
||||
stop := StartTiming(ctx, "ctx_timed")
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
stop()
|
||||
|
||||
entries := tc.Entries()
|
||||
if len(entries) != 2 {
|
||||
t.Errorf("expected 2 entries, got %d", len(entries))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimingCollectorFromContext_NilContext(t *testing.T) {
|
||||
tc := TimingCollectorFromContext(nil)
|
||||
if tc != nil {
|
||||
t.Error("TimingCollectorFromContext(nil) should return nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimingCollectorFromContext_NoCollector(t *testing.T) {
|
||||
tc := TimingCollectorFromContext(context.Background())
|
||||
if tc != nil {
|
||||
t.Error("TimingCollectorFromContext with no collector should return nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartTiming_NoCollector(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
stop := StartTiming(ctx, "no_collector")
|
||||
// Should not panic
|
||||
stop()
|
||||
}
|
||||
|
||||
func TestIsPerfTimingEnabled(t *testing.T) {
|
||||
// Clear the env var first
|
||||
os.Unsetenv(PerfTimingEnv)
|
||||
|
||||
if IsPerfTimingEnabled() {
|
||||
t.Error("IsPerfTimingEnabled should return false when env var is not set")
|
||||
}
|
||||
|
||||
os.Setenv(PerfTimingEnv, "1")
|
||||
defer os.Unsetenv(PerfTimingEnv)
|
||||
|
||||
if !IsPerfTimingEnabled() {
|
||||
t.Error("IsPerfTimingEnabled should return true when env var is set")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user