zenmux-usage-cli/cmd/zenmux-usage/main_test.go

package main

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

// newTestServer returns an httptest server that maps API keys to status codes.
// Keys ending with "-401" return 401; "-422" return 422; everything else returns 200 with the sample body.
func newTestServer(t *testing.T) *httptest.Server {
	t.Helper()
	sample, err := os.ReadFile(filepath.Join("..", "..", "internal", "api", "testdata", "sample.json"))
	if err != nil {
		t.Fatalf("read sample: %v", err)
	}
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization")
		switch {
		case strings.HasSuffix(auth, "-401"):
			w.WriteHeader(http.StatusUnauthorized)
		case strings.HasSuffix(auth, "-422"):
			w.WriteHeader(http.StatusUnprocessableEntity)
		default:
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write(sample)
		}
	}))
}

func writeConfig(t *testing.T, body string) string {
	t.Helper()
	dir := t.TempDir()
	p := filepath.Join(dir, "config.yaml")
	if err := os.WriteFile(p, []byte(body), 0o600); err != nil {
		t.Fatal(err)
	}
	return p
}

func TestRun_Version(t *testing.T) {
	var out, errb bytes.Buffer
	code := run([]string{"--version"}, &out, &errb)
	if code != exitOK {
		t.Errorf("exit = %d want %d", code, exitOK)
	}
	if !strings.Contains(out.String(), "zenmux-usage") {
		t.Errorf("version output missing program name: %q", out.String())
	}
}

func TestRun_UnknownAccountExits2(t *testing.T) {
	cfg := writeConfig(t, `
accounts:
  - name: personal
    api_key: sk-p
`)
	var out, errb bytes.Buffer
	code := run([]string{"--config", cfg, "--account", "missing"}, &out, &errb)
	if code != exitUsage {
		t.Errorf("exit = %d want %d, stderr=%q", code, exitUsage, errb.String())
	}
	if !strings.Contains(errb.String(), "account not found") {
		t.Errorf("stderr missing 'account not found': %q", errb.String())
	}
}

func TestRun_NoConfigNoEnvExits3(t *testing.T) {
	t.Setenv("ZENMUX_MANAGEMENT_API_KEY", "")
	t.Setenv("XDG_CONFIG_HOME", t.TempDir()) // point default path at an empty dir
	var out, errb bytes.Buffer
	code := run(nil, &out, &errb)
	if code != exitNoAccount {
		t.Errorf("exit = %d want %d, stderr=%q", code, exitNoAccount, errb.String())
	}
}

func TestRun_ExplicitConfigMissingExits3(t *testing.T) {
	var out, errb bytes.Buffer
	code := run([]string{"--config", "/nonexistent/does/not/exist.yaml"}, &out, &errb)
	if code != exitNoAccount {
		t.Errorf("exit = %d want %d, stderr=%q", code, exitNoAccount, errb.String())
	}
}

func TestRun_MalformedConfigExits7(t *testing.T) {
	cfg := writeConfig(t, "accounts:\n  - name: a\n    api_key: [")
	var out, errb bytes.Buffer
	code := run([]string{"--config", cfg}, &out, &errb)
	if code != exitConfigParse {
		t.Errorf("exit = %d want %d, stderr=%q", code, exitConfigParse, errb.String())
	}
}

func TestRun_InvalidTimeoutExits2(t *testing.T) {
	var out, errb bytes.Buffer
	code := run([]string{"--timeout", "0s"}, &out, &errb)
	if code != exitUsage {
		t.Errorf("exit = %d want %d", code, exitUsage)
	}
}

func TestRun_UnexpectedPositionalExits2(t *testing.T) {
	var out, errb bytes.Buffer
	code := run([]string{"extra"}, &out, &errb)
	if code != exitUsage {
		t.Errorf("exit = %d want %d", code, exitUsage)
	}
}

func TestRun_HelpExits0(t *testing.T) {
	var out, errb bytes.Buffer
	code := run([]string{"--help"}, &out, &errb)
	if code != exitOK {
		t.Errorf("exit = %d want %d", code, exitOK)
	}
}

func TestRun_SingleAccountSuccess(t *testing.T) {
	srv := newTestServer(t)
	defer srv.Close()
	t.Setenv("ZENMUX_BASE_URL", srv.URL)

	var out, errb bytes.Buffer
	code := run([]string{"--api-key", "sk-ok", "--no-color"}, &out, &errb)
	if code != exitOK {
		t.Fatalf("exit = %d want %d; stderr=%q", code, exitOK, errb.String())
	}
	if !strings.Contains(out.String(), "Ultra plan") {
		t.Errorf("expected Ultra plan in output, got %q", out.String())
	}
}

func TestRun_SingleAccountAuthFailExits4(t *testing.T) {
	srv := newTestServer(t)
	defer srv.Close()
	t.Setenv("ZENMUX_BASE_URL", srv.URL)

	var out, errb bytes.Buffer
	code := run([]string{"--api-key", "sk-401", "--no-color"}, &out, &errb)
	if code != exitAuth {
		t.Fatalf("exit = %d want %d", code, exitAuth)
	}
}

func TestRun_MultiAccountMixedExits1(t *testing.T) {
	srv := newTestServer(t)
	defer srv.Close()
	t.Setenv("ZENMUX_BASE_URL", srv.URL)
	cfg := writeConfig(t, `
accounts:
  - name: ok
    api_key: sk-ok
  - name: bad
    api_key: sk-401
`)

	var out, errb bytes.Buffer
	code := run([]string{"--config", cfg, "--no-color"}, &out, &errb)
	if code != exitGeneric {
		t.Fatalf("exit = %d want %d (mixed)", code, exitGeneric)
	}
	if !strings.Contains(out.String(), "━━━ ok ━━━") || !strings.Contains(out.String(), "━━━ bad ━━━") {
		t.Errorf("both account headers should appear; got %q", out.String())
	}
	if !strings.Contains(out.String(), "error:") {
		t.Errorf("failed account should have error line; got %q", out.String())
	}
}

func TestRun_AllRateLimitedExits5(t *testing.T) {
	srv := newTestServer(t)
	defer srv.Close()
	t.Setenv("ZENMUX_BASE_URL", srv.URL)
	cfg := writeConfig(t, `
accounts:
  - name: a
    api_key: sk-422
  - name: b
    api_key: sk-422
`)
	var out, errb bytes.Buffer
	code := run([]string{"--config", cfg, "--no-color"}, &out, &errb)
	if code != exitRateLimited {
		t.Fatalf("exit = %d want %d", code, exitRateLimited)
	}
}

func TestRun_JSONSingleAccountPassthrough(t *testing.T) {
	srv := newTestServer(t)
	defer srv.Close()
	t.Setenv("ZENMUX_BASE_URL", srv.URL)

	var out, errb bytes.Buffer
	code := run([]string{"--api-key", "sk-ok", "--json"}, &out, &errb)
	if code != exitOK {
		t.Fatalf("exit = %d want %d", code, exitOK)
	}
	if !strings.Contains(out.String(), `"tier": "ultra"`) {
		t.Errorf("expected raw passthrough including plan tier; got %q", out.String())
	}
}

func TestRun_AllNetworkTimeoutExits6(t *testing.T) {
	// Slow server forces a timeout.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(200 * time.Millisecond)
	}))
	defer srv.Close()
	t.Setenv("ZENMUX_BASE_URL", srv.URL)

	var out, errb bytes.Buffer
	code := run([]string{"--api-key", "sk-slow", "--timeout", "10ms", "--no-color"}, &out, &errb)
	if code != exitNetwork {
		t.Fatalf("exit = %d want %d", code, exitNetwork)
	}
}

func TestRun_JSONMultiAccount(t *testing.T) {
	srv := newTestServer(t)
	defer srv.Close()
	t.Setenv("ZENMUX_BASE_URL", srv.URL)
	cfg := writeConfig(t, `
accounts:
  - name: ok
    api_key: sk-ok
  - name: bad
    api_key: sk-401
`)
	var out, errb bytes.Buffer
	code := run([]string{"--config", cfg, "--json"}, &out, &errb)
	// Mixed outcomes → 1
	if code != exitGeneric {
		t.Fatalf("exit = %d want %d", code, exitGeneric)
	}
	// Array shape.
	o := strings.TrimSpace(out.String())
	if !strings.HasPrefix(o, "[") || !strings.HasSuffix(o, "]") {
		t.Errorf("expected JSON array, got %q", o)
	}
	if !strings.Contains(o, `"account": "ok"`) || !strings.Contains(o, `"account": "bad"`) {
		t.Errorf("missing account keys in multi JSON: %q", o)
	}
}