dingtalk-workspace-cli/internal/app/audit_runtime_test.go

package app

import (
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/audit"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/auth"
)

// TestAuditIdentityReresolvesOnProfileSwitch guards the reviewer's finding that a
// long-running process (e.g. serve mode) must attribute events to the ACTIVE
// runtime profile rather than reusing a process-global first Actor. It also
// asserts the per-profile cache avoids redundant token loads within one profile.
func TestAuditIdentityReresolvesOnProfileSwitch(t *testing.T) {
	prevLoader := loadTokenForProfile
	prevProfile := auth.RuntimeProfile()
	t.Cleanup(func() {
		loadTokenForProfile = prevLoader
		auth.SetRuntimeProfile(prevProfile)
		resetAuditIdentityCache()
	})
	resetAuditIdentityCache()

	var mu sync.Mutex
	calls := map[string]int{}
	loadTokenForProfile = func(_ /*configDir*/, profile string) (*auth.TokenData, error) {
		mu.Lock()
		calls[profile]++
		mu.Unlock()
		switch profile {
		case "orgA":
			return &auth.TokenData{UserID: "ua", UserName: "Alice", CorpID: "ca", CorpName: "CorpA"}, nil
		case "orgB":
			return &auth.TokenData{UserID: "ub", UserName: "Bob", CorpID: "cb", CorpName: "CorpB"}, nil
		default:
			return nil, nil
		}
	}

	auth.SetRuntimeProfile("orgA")
	if actor, _ := auditIdentity(); actor.UserID != "ua" || actor.CorpName != "CorpA" {
		t.Fatalf("orgA: got %+v, want Alice/CorpA", actor)
	}
	// Second call under the same profile must hit the cache (no extra load).
	if actor, _ := auditIdentity(); actor.UserID != "ua" {
		t.Fatalf("orgA cached: got %+v", actor)
	}

	auth.SetRuntimeProfile("orgB")
	if actor, _ := auditIdentity(); actor.UserID != "ub" || actor.CorpName != "CorpB" {
		t.Fatalf("orgB: got %+v, want Bob/CorpB (stale Actor reused?)", actor)
	}

	mu.Lock()
	defer mu.Unlock()
	if calls["orgA"] != 1 {
		t.Fatalf("orgA loaded %d times, want 1 (cache miss?)", calls["orgA"])
	}
	if calls["orgB"] != 1 {
		t.Fatalf("orgB loaded %d times, want 1", calls["orgB"])
	}
}

func resetAuditIdentityCache() {
	auditIDMu.Lock()
	defer auditIDMu.Unlock()
	cachedActor = audit.Actor{}
	cachedAgentID = ""
	cachedProfile = ""
	identityLoaded = false
}

// TestCloseAuditSinkDrainsOnErrorPath guards the reviewer's V5 finding: when a
// command's RunE returns an error, Cobra skips PersistentPostRunE, so the audit
// drain must instead happen through the unconditional defer in Execute that calls
// CloseAuditSink. This test wires a real forwarder-backed sink into the shared
// slot and asserts CloseAuditSink flushes the queued forward exactly as the
// error-path defer would, and that a second call is a harmless no-op.
func TestCloseAuditSinkDrainsOnErrorPath(t *testing.T) {
	var delivered int64
	var releaseOnce sync.Once
	release := make(chan struct{})
	releaseFn := func() { releaseOnce.Do(func() { close(release) }) }
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		<-release // hold the request until the drain awaits it
		atomic.AddInt64(&delivered, 1)
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()
	defer releaseFn() // LIFO: unblock any in-flight handler before srv.Close()

	writer, err := audit.NewDateRotatingWriter(t.TempDir(), 0)
	if err != nil {
		t.Fatal(err)
	}
	fwd := audit.NewHTTPForwarder(srv.URL, "", audit.RedactNone, nil)
	sink := audit.NewFileSink(writer, audit.NewChain(""), fwd)

	prevSink := sharedAuditSink
	t.Cleanup(func() {
		sharedAuditSink = prevSink
		auditCloseOnce = sync.Once{}
	})
	sharedAuditSink = sink
	auditCloseOnce = sync.Once{}

	if err := sink.Emit(&audit.Event{Timestamp: time.Unix(0, 0), Product: "calendar", Command: "event_list", Result: "error"}); err != nil {
		t.Fatalf("emit: %v", err)
	}
	if got := atomic.LoadInt64(&delivered); got != 0 {
		t.Fatalf("forward delivered before drain: %d", got)
	}

	// Let the held request complete, then drain via the same entry point the
	// error-path defer uses. CloseAuditSink blocks until the forward goroutine
	// observes the HTTP response, so the counter is settled when it returns.
	releaseFn()
	CloseAuditSink()

	if got := atomic.LoadInt64(&delivered); got != 1 {
		t.Fatalf("forward not drained on error path: delivered=%d, want 1", got)
	}

	// Idempotent: the success-path PersistentPostRunE and the defer both call it.
	CloseAuditSink()
}