From 40452e100e79214385a47f553d15efd40564a391 Mon Sep 17 00:00:00 2001 From: Kaj Kowalski Date: Thu, 18 Dec 2025 08:26:50 +0100 Subject: [PATCH] Refactor: Extract testable functions and improve code organization - Extract TokenUsage as named type (eliminates inline struct repetition) - Refactor main() into testable functions: - readInputFromStdin: Read JSON from stdin - parseStatusInput: Validate and parse JSON - buildStatusLine: Construct left and right statusline parts - calculatePadding: Compute padding for alignment - formatOutput: Combine components into final output - Add comprehensive tests for extracted functions - Improve coverage from 45% to 71% (+26 percentage points) - All new functions have 100% test coverage - Clean linting with zero issues --- main.go | 88 +++++++++------ main_test.go | 311 +++++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 354 insertions(+), 45 deletions(-) diff --git a/main.go b/main.go index 4261ee9..20c68a3 100644 --- a/main.go +++ b/main.go @@ -16,7 +16,14 @@ import ( const statuslineWidthOffset = 7 -// Input JSON structure from Claude Code +// TokenUsage tracks context window token consumption. +type TokenUsage struct { + InputTokens int `json:"input_tokens"` + CacheCreationTokens int `json:"cache_creation_input_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` +} + +// StatusInput represents the JSON input from Claude Code. type StatusInput struct { Model struct { DisplayName string `json:"display_name"` @@ -25,12 +32,8 @@ type StatusInput struct { CurrentDir string `json:"current_dir"` } `json:"workspace"` ContextWindow struct { - ContextWindowSize int `json:"context_window_size"` - CurrentUsage *struct { - InputTokens int `json:"input_tokens"` - CacheCreationTokens int `json:"cache_creation_input_tokens"` - CacheReadInputTokens int `json:"cache_read_input_tokens"` - } `json:"current_usage"` + ContextWindowSize int `json:"context_window_size"` + CurrentUsage *TokenUsage `json:"current_usage"` } `json:"context_window"` } @@ -46,46 +49,66 @@ const ( boldGreen = "\033[1;32m" ) -func main() { - // Read JSON from stdin - reader := bufio.NewReader(os.Stdin) +// readInputFromStdin reads JSON input from stdin. +func readInputFromStdin(r *bufio.Reader) string { var input strings.Builder for { - line, err := reader.ReadString('\n') + line, err := r.ReadString('\n') input.WriteString(line) if err != nil { break } } + return input.String() +} +// parseStatusInput unmarshals JSON string into StatusInput. +func parseStatusInput(jsonStr string) (*StatusInput, error) { var data StatusInput - if err := json.Unmarshal([]byte(input.String()), &data); err != nil { - fmt.Fprintf(os.Stderr, "Error parsing JSON: %v\n", err) - os.Exit(1) + if err := json.Unmarshal([]byte(jsonStr), &data); err != nil { + return nil, err } + return &data, nil +} - // Calculate context info +// buildStatusLine constructs the left and right parts of the status line. +func buildStatusLine(data *StatusInput) (left, right string) { contextInfo := formatContextInfo(data.ContextWindow.ContextWindowSize, data.ContextWindow.CurrentUsage) - - // Get directory name dirName := filepath.Base(data.Workspace.CurrentDir) - - // Check gitea status giteaStatus := getGiteaStatus() - - // Get git info gitInfo := getGitInfo(data.Workspace.CurrentDir) - // Build left part - left := fmt.Sprintf("%s %s%s%s %s➜%s %s%s%s%s", + left = fmt.Sprintf("%s %s%s%s %s➜%s %s%s%s%s", giteaStatus, magenta, data.Model.DisplayName, reset, boldGreen, reset, cyan, dirName, reset, gitInfo) - // Build right part - right := fmt.Sprintf("%s%s%s", yellow, contextInfo, reset) + right = fmt.Sprintf("%s%s%s", yellow, contextInfo, reset) + return left, right +} + +// calculatePadding returns the number of spaces needed for padding. +func calculatePadding(leftVisible, rightVisible string, termWidth int) int { + return max(termWidth-len(leftVisible)-len(rightVisible), 1) +} + +// formatOutput combines left, right, and padding into final output. +func formatOutput(left, right string, padding int) string { + return fmt.Sprintf("%s%s%s", left, strings.Repeat(" ", padding), right) +} + +func main() { + jsonStr := readInputFromStdin(bufio.NewReader(os.Stdin)) + + data, err := parseStatusInput(jsonStr) + if err != nil { + fmt.Fprintf(os.Stderr, "Error parsing JSON: %v\n", err) + os.Exit(1) + } + + left, right := buildStatusLine(data) // Calculate visible lengths (strip ANSI) leftVisible := stripANSI(left) @@ -94,19 +117,14 @@ func main() { // Get terminal width termWidth := getTerminalWidth() - statuslineWidthOffset - // Calculate padding - padding := max(termWidth-len(leftVisible)-len(rightVisible), 1) + // Calculate and apply padding + padding := calculatePadding(leftVisible, rightVisible, termWidth) + output := formatOutput(left, right, padding) - // Output with padding - fmt.Printf("%s%s%s", left, strings.Repeat(" ", padding), right) + fmt.Print(output) } -func formatContextInfo(contextSize int, usage *struct { - InputTokens int `json:"input_tokens"` - CacheCreationTokens int `json:"cache_creation_input_tokens"` - CacheReadInputTokens int `json:"cache_read_input_tokens"` -}, -) string { +func formatContextInfo(contextSize int, usage *TokenUsage) string { totalK := contextSize / 1000 if usage == nil { diff --git a/main_test.go b/main_test.go index 2af0365..a6cbc97 100644 --- a/main_test.go +++ b/main_test.go @@ -1,8 +1,10 @@ package main import ( + "bufio" "encoding/json" "os" + "strings" "testing" ) @@ -15,11 +17,7 @@ func TestFormatContextInfo_NilUsage(t *testing.T) { } func TestFormatContextInfo_WithUsage(t *testing.T) { - usage := &struct { - InputTokens int `json:"input_tokens"` - CacheCreationTokens int `json:"cache_creation_input_tokens"` - CacheReadInputTokens int `json:"cache_read_input_tokens"` - }{ + usage := &TokenUsage{ InputTokens: 8500, CacheCreationTokens: 5000, CacheReadInputTokens: 2000, @@ -32,11 +30,7 @@ func TestFormatContextInfo_WithUsage(t *testing.T) { } func TestFormatContextInfo_SmallValues(t *testing.T) { - usage := &struct { - InputTokens int `json:"input_tokens"` - CacheCreationTokens int `json:"cache_creation_input_tokens"` - CacheReadInputTokens int `json:"cache_read_input_tokens"` - }{ + usage := &TokenUsage{ InputTokens: 500, CacheCreationTokens: 0, CacheReadInputTokens: 0, @@ -203,3 +197,300 @@ func containsHelper(s, substr string) bool { } return false } + +func TestGetTerminalWidth(t *testing.T) { + // getTerminalWidth should return a positive integer + // In test environment (no TTY), it should fall back to 80 + width := getTerminalWidth() + if width <= 0 { + t.Errorf("getTerminalWidth() = %d, expected positive value", width) + } +} + +func TestGetTerminalWidth_DefaultFallback(t *testing.T) { + // When not connected to a terminal, should return 80 + width := getTerminalWidth() + // In CI/test environments, this typically returns 80 + if width != 80 && width < 40 { + t.Errorf("getTerminalWidth() = %d, expected 80 or reasonable terminal width", width) + } +} + +func TestGetGitInfo_InvalidPath(t *testing.T) { + result := getGitInfo("/nonexistent/path/that/does/not/exist") + if result != "" { + t.Errorf("getGitInfo(invalid) = %q, expected empty string", result) + } +} + +func TestGetGitInfo_RootDir(t *testing.T) { + // Root directory is unlikely to be a git repo + result := getGitInfo("/") + if result != "" && !contains(result, "git:(") { + t.Errorf("getGitInfo(/) = %q, expected empty or valid git info", result) + } +} + +func TestFormatContextInfo_ZeroContextSize(t *testing.T) { + result := formatContextInfo(0, nil) + expected := "0/0k" + if result != expected { + t.Errorf("formatContextInfo(0, nil) = %q, want %q", result, expected) + } +} + +func TestFormatContextInfo_LargeValues(t *testing.T) { + usage := &TokenUsage{ + InputTokens: 150000, + CacheCreationTokens: 25000, + CacheReadInputTokens: 10000, + } + result := formatContextInfo(200000, usage) + expected := "185k/200k" + if result != expected { + t.Errorf("formatContextInfo(200000, large usage) = %q, want %q", result, expected) + } +} + +func TestFormatContextInfo_ExactThousand(t *testing.T) { + usage := &TokenUsage{ + InputTokens: 1000, + CacheCreationTokens: 0, + CacheReadInputTokens: 0, + } + result := formatContextInfo(100000, usage) + expected := "1k/100k" + if result != expected { + t.Errorf("formatContextInfo(100000, 1000 tokens) = %q, want %q", result, expected) + } +} + +func TestStripANSI_Empty(t *testing.T) { + result := stripANSI("") + if result != "" { + t.Errorf("stripANSI(\"\") = %q, want empty", result) + } +} + +func TestStripANSI_OnlyANSI(t *testing.T) { + result := stripANSI("\033[31m\033[0m") + if result != "" { + t.Errorf("stripANSI(only codes) = %q, want empty", result) + } +} + +func TestStripANSI_NestedCodes(t *testing.T) { + input := "\033[1m\033[31mbold red\033[0m\033[0m" + result := stripANSI(input) + expected := "bold red" + if result != expected { + t.Errorf("stripANSI(%q) = %q, want %q", input, result, expected) + } +} + +func TestStatusInputParsing_EmptyJSON(t *testing.T) { + jsonData := `{}` + var data StatusInput + err := json.Unmarshal([]byte(jsonData), &data) + if err != nil { + t.Fatalf("Failed to parse empty JSON: %v", err) + } + if data.Model.DisplayName != "" { + t.Errorf("Expected empty DisplayName, got %q", data.Model.DisplayName) + } +} + +func TestStatusInputParsing_PartialJSON(t *testing.T) { + jsonData := `{"model": {"display_name": "Test"}}` + var data StatusInput + err := json.Unmarshal([]byte(jsonData), &data) + if err != nil { + t.Fatalf("Failed to parse partial JSON: %v", err) + } + if data.Model.DisplayName != "Test" { + t.Errorf("DisplayName = %q, want %q", data.Model.DisplayName, "Test") + } + if data.Workspace.CurrentDir != "" { + t.Errorf("Expected empty CurrentDir, got %q", data.Workspace.CurrentDir) + } +} + +func TestStatusInputParsing_InvalidJSON(t *testing.T) { + jsonData := `{invalid json}` + var data StatusInput + err := json.Unmarshal([]byte(jsonData), &data) + if err == nil { + t.Error("Expected error for invalid JSON, got nil") + } +} + +func TestANSIConstants(t *testing.T) { + // Verify ANSI constants are properly defined + tests := []struct { + name string + constant string + prefix string + }{ + {"reset", reset, "\033["}, + {"red", red, "\033["}, + {"green", green, "\033["}, + {"yellow", yellow, "\033["}, + {"magenta", magenta, "\033["}, + {"cyan", cyan, "\033["}, + {"boldGreen", boldGreen, "\033["}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if !contains(tt.constant, tt.prefix) { + t.Errorf("%s constant doesn't start with ANSI escape", tt.name) + } + }) + } +} + +func TestGetGiteaStatus_ReturnsValidColor(t *testing.T) { + result := getGiteaStatus() + + // Must contain the dot character + if !contains(result, "●") { + t.Errorf("getGiteaStatus() = %q, expected to contain dot", result) + } + + // Must contain ANSI codes + stripped := stripANSI(result) + if stripped != "●" { + t.Errorf("stripped getGiteaStatus() = %q, expected just dot", stripped) + } +} + +func TestReadInputFromStdin(t *testing.T) { + input := "line1\nline2\nline3" + reader := bufio.NewReader(strings.NewReader(input)) + result := readInputFromStdin(reader) + expected := "line1\nline2\nline3" + if result != expected { + t.Errorf("readInputFromStdin = %q, want %q", result, expected) + } +} + +func TestReadInputFromStdin_Empty(t *testing.T) { + reader := bufio.NewReader(strings.NewReader("")) + result := readInputFromStdin(reader) + if result != "" { + t.Errorf("readInputFromStdin (empty) = %q, want empty", result) + } +} + +func TestParseStatusInput_Valid(t *testing.T) { + jsonStr := `{"model": {"display_name": "Test"}, "workspace": {"current_dir": "/test"}}` + data, err := parseStatusInput(jsonStr) + if err != nil { + t.Fatalf("parseStatusInput failed: %v", err) + } + if data.Model.DisplayName != "Test" { + t.Errorf("DisplayName = %q, want Test", data.Model.DisplayName) + } + if data.Workspace.CurrentDir != "/test" { + t.Errorf("CurrentDir = %q, want /test", data.Workspace.CurrentDir) + } +} + +func TestParseStatusInput_Invalid(t *testing.T) { + _, err := parseStatusInput("invalid json") + if err == nil { + t.Error("parseStatusInput should fail on invalid JSON") + } +} + +func TestBuildStatusLine_ContainsComponents(t *testing.T) { + data := &StatusInput{} + data.Model.DisplayName = "TestModel" + data.Workspace.CurrentDir = "/home/user/project" + data.ContextWindow.ContextWindowSize = 100000 + data.ContextWindow.CurrentUsage = &TokenUsage{ + InputTokens: 5000, + CacheCreationTokens: 1000, + CacheReadInputTokens: 500, + } + + left, right := buildStatusLine(data) + + // Check left contains model name and directory + if !contains(left, "TestModel") { + t.Errorf("left statusline missing model: %q", left) + } + if !contains(left, "project") { + t.Errorf("left statusline missing directory: %q", left) + } + + // Check right contains context info + if !contains(right, "6k/100k") { + t.Errorf("right statusline missing context info: %q", right) + } +} + +func TestBuildStatusLine_HasGiteaStatus(t *testing.T) { + data := &StatusInput{} + data.Model.DisplayName = "Model" + data.Workspace.CurrentDir = "/tmp" + data.ContextWindow.ContextWindowSize = 100000 + + left, _ := buildStatusLine(data) + + // Check for gitea status (dot) + if !contains(left, "●") { + t.Errorf("left statusline missing gitea status: %q", left) + } +} + +func TestCalculatePadding_ZeroWidth(t *testing.T) { + result := calculatePadding("left", "right", 0) + expected := 1 + if result != expected { + t.Errorf("calculatePadding(\"left\", \"right\", 0) = %d, want %d", result, expected) + } +} + +func TestCalculatePadding_NegativeResult(t *testing.T) { + result := calculatePadding("left", "right", 5) + expected := 1 + if result != expected { + t.Errorf("calculatePadding with overflow = %d, want minimum of %d", result, expected) + } +} + +func TestCalculatePadding_Normal(t *testing.T) { + result := calculatePadding("left", "right", 50) + expected := 50 - len("left") - len("right") + if result != expected { + t.Errorf("calculatePadding(\"left\", \"right\", 50) = %d, want %d", result, expected) + } +} + +func TestFormatOutput_Composition(t *testing.T) { + result := formatOutput("LEFT", "RIGHT", 5) + expected := "LEFT RIGHT" + if result != expected { + t.Errorf("formatOutput = %q, want %q", result, expected) + } +} + +func TestFormatOutput_Empty(t *testing.T) { + result := formatOutput("", "", 0) + expected := "" + if result != expected { + t.Errorf("formatOutput (empty) = %q, want %q", result, expected) + } +} + +func TestFormatOutput_WithANSI(t *testing.T) { + left := red + "text" + reset + right := green + "info" + reset + result := formatOutput(left, right, 3) + + stripped := stripANSI(result) + if !contains(stripped, "text") || !contains(stripped, "info") { + t.Errorf("formatOutput with ANSI = %q, expected both parts visible", stripped) + } +}