diff --git a/.gitignore b/.gitignore index 076c8c5..28dd888 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ .vscode *.exe clash-speedtest +.gocache +.gomodcache \ No newline at end of file diff --git a/main.go b/main.go index 7da3958..39c778f 100644 --- a/main.go +++ b/main.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "strings" + "sync" "time" tea "github.com/charmbracelet/bubbletea" @@ -72,7 +73,7 @@ func main() { if !*fastMode { requestedMode, err = speedtester.ParseSpeedMode(*speedMode) if err != nil { - log.Fatalln("parse speed mode failed: %s", err) + log.Fatalf("parse speed mode failed: %v", err) } } @@ -94,13 +95,13 @@ func main() { UserAgent: *userAgent, }) if err != nil { - log.Fatalln("create speed tester failed: %s", err) + log.Fatalf("create speed tester failed: %v", err) } effectiveMode := speedTester.Mode() allProxies, err := speedTester.LoadProxies() if err != nil { - log.Fatalln("load proxies failed: %s", err) + log.Fatalf("load proxies failed: %v", err) } outputMode := output.DetermineOutputMode(output.IsTerminalFile) @@ -110,7 +111,7 @@ func main() { var err error tsvWriter, err = output.NewTSVWriter(os.Stdout, effectiveMode) if err != nil { - log.Fatalln("create TSV writer failed: %s", err) + log.Fatalf("create TSV writer failed: %v", err) } } @@ -120,47 +121,110 @@ func main() { collectResults := *outputPath != "" // Run TUI for Interactive mode resultChannel := make(chan *speedtester.Result, len(allProxies)) - resultsDone := make(chan struct{}) - saveResult := make(chan error, 1) - - // Start testing in goroutine to send results to channel - go func() { - speedTester.TestProxies(allProxies, func(result *speedtester.Result) { - if collectResults { - results = append(results, result) + var resultsMu sync.Mutex + var activeRunDone <-chan struct{} + setActiveRun := func(done <-chan struct{}) { + resultsMu.Lock() + activeRunDone = done + resultsMu.Unlock() + } + storeFullRunResults := func(runResults []*speedtester.Result) { + if !collectResults { + return + } + resultsMu.Lock() + results = runResults + resultsMu.Unlock() + } + replaceStoredResult := func(result *speedtester.Result) { + if !collectResults { + return + } + resultsMu.Lock() + defer resultsMu.Unlock() + for i, existing := range results { + if existing.ProxyName == result.ProxyName { + results[i] = result + return } - resultChannel <- result - }) - close(resultChannel) - close(resultsDone) - }() - - if collectResults { - // Save results once all tests finish, without blocking the TUI loop. + } + results = append(results, result) + } + startTestRun := func(proxies map[string]*speedtester.CProxy, replaceAll bool) { + done := make(chan struct{}) + setActiveRun(done) go func() { - <-resultsDone - results = output.SortResults(results, effectiveMode) - saveResult <- saveConfig(results, effectiveMode) + defer close(done) + runResults := make([]*speedtester.Result, 0, len(proxies)) + if len(proxies) == 1 { + for name, proxy := range proxies { + result := speedTester.TestProxy(name, proxy) + if replaceAll { + runResults = append(runResults, result) + } else { + replaceStoredResult(result) + } + resultChannel <- result + } + } else { + speedTester.TestProxies(proxies, func(result *speedtester.Result) { + if replaceAll { + runResults = append(runResults, result) + } else { + replaceStoredResult(result) + } + resultChannel <- result + }) + } + if replaceAll { + storeFullRunResults(runResults) + } + resultChannel <- nil }() } + startTestRun(allProxies, true) + // Create and run TUI + model := tui.NewTUIModel(effectiveMode, len(allProxies), resultChannel) + model.SetRetestCallbacks( + func(_ chan<- *speedtester.Result) { + startTestRun(allProxies, true) + }, + func(name string, out chan<- *speedtester.Result) { + proxy, ok := allProxies[name] + if !ok { + out <- nil + return + } + startTestRun(map[string]*speedtester.CProxy{name: proxy}, false) + }, + ) p := tea.NewProgram( - tui.NewTUIModel(effectiveMode, len(allProxies), resultChannel), + model, tea.WithAltScreen(), tea.WithMouseAllMotion(), ) if _, err := p.Run(); err != nil { - log.Fatalln("TUI failed: %s", err) + log.Fatalf("TUI failed: %v", err) } if !collectResults { return } - err = <-saveResult + resultsMu.Lock() + done := activeRunDone + resultsMu.Unlock() + if done != nil { + <-done + } + resultsMu.Lock() + results = output.SortResults(results, effectiveMode) + resultsMu.Unlock() + err = saveConfig(results, effectiveMode) if err != nil { - log.Fatalln("save config file failed: %s", err) + log.Fatalf("save config file failed: %v", err) } fmt.Printf("\nsave config file to: %s\n", *outputPath) return @@ -182,7 +246,7 @@ func main() { if *outputPath != "" { err = saveConfig(results, effectiveMode) if err != nil { - log.Fatalln("save config file failed: %s", err) + log.Fatalf("save config file failed: %v", err) } fmt.Printf("\nsave config file to: %s\n", *outputPath) } diff --git a/speedtester/speedtester.go b/speedtester/speedtester.go index a117cb6..6f8eab0 100644 --- a/speedtester/speedtester.go +++ b/speedtester/speedtester.go @@ -342,10 +342,14 @@ func buildProxyServerPortKey(proxy *CProxy) (string, bool) { func (st *SpeedTester) TestProxies(proxies map[string]*CProxy, tester func(result *Result)) { for name, proxy := range proxies { - tester(st.testProxy(name, proxy)) + tester(st.TestProxy(name, proxy)) } } +func (st *SpeedTester) TestProxy(name string, proxy *CProxy) *Result { + return st.testProxy(name, proxy) +} + type Result struct { ProxyName string `json:"proxy_name"` ProxyType string `json:"proxy_type"` diff --git a/tui/detail.go b/tui/detail.go index b7616c4..69c6a1c 100644 --- a/tui/detail.go +++ b/tui/detail.go @@ -98,7 +98,7 @@ func buildDetailContent(result *speedtester.Result, width int, mode speedtester. lines = appendWrappedValue(lines, "Upload Error:", result.FormatUploadError(), width) } } - lines = append(lines, "", "Press ESC to close details.") + lines = append(lines, "", "Press R to retest this node. Press ESC to close details.") return strings.Join(lines, "\n") } diff --git a/tui/detail_test.go b/tui/detail_test.go index 21c1268..bbe094d 100644 --- a/tui/detail_test.go +++ b/tui/detail_test.go @@ -178,6 +178,9 @@ func TestBuildDetailContentDownloadOnly(t *testing.T) { if strings.Contains(content, result.UploadError) { t.Fatalf("expected download-only detail to omit upload error, got %q", content) } + if !strings.Contains(content, "Press R to retest this node.") { + t.Fatalf("expected detail to include retest hint, got %q", content) + } } func TestDetailPanelHeightUpdatesOnSelectionChange(t *testing.T) { diff --git a/tui/help.go b/tui/help.go index 42ddfec..4d839d4 100644 --- a/tui/help.go +++ b/tui/help.go @@ -14,6 +14,8 @@ type helpState struct { type helpKeyMap struct { Quit key.Binding + RetestAll key.Binding + RetestOne key.Binding CloseDetail key.Binding Table table.KeyMap } @@ -27,6 +29,14 @@ func newHelpState(tableKeys table.KeyMap) helpState { key.WithKeys("q", "ctrl+c"), key.WithHelp("q/ctrl+c", "quit"), ), + RetestAll: key.NewBinding( + key.WithKeys("r"), + key.WithHelp("r", "retest all"), + ), + RetestOne: key.NewBinding( + key.WithKeys("r"), + key.WithHelp("r", "retest node"), + ), CloseDetail: key.NewBinding( key.WithKeys("esc"), key.WithHelp("esc", "close details"), @@ -43,6 +53,8 @@ func (h *helpState) setWidth(width int) { func (h *helpState) setDetailVisible(visible bool) { h.keyMap.CloseDetail.SetEnabled(visible) + h.keyMap.RetestOne.SetEnabled(visible) + h.keyMap.RetestAll.SetEnabled(!visible) } func (h helpState) view() string { @@ -61,6 +73,8 @@ func (km helpKeyMap) ShortHelp() []key.Binding { return []key.Binding{ km.Table.LineUp, km.Table.LineDown, + km.RetestAll, + km.RetestOne, km.Quit, km.CloseDetail, } @@ -70,6 +84,6 @@ func (km helpKeyMap) FullHelp() [][]key.Binding { return [][]key.Binding{ {km.Table.LineUp, km.Table.LineDown, km.Table.GotoTop, km.Table.GotoBottom}, {km.Table.PageUp, km.Table.PageDown, km.Table.HalfPageUp, km.Table.HalfPageDown}, - {km.CloseDetail, km.Quit}, + {km.RetestAll, km.RetestOne, km.CloseDetail, km.Quit}, } } diff --git a/tui/help_test.go b/tui/help_test.go index 3d9571f..4cde133 100644 --- a/tui/help_test.go +++ b/tui/help_test.go @@ -27,13 +27,25 @@ func TestHelpViewShowsQuitAndDetailKeys(t *testing.T) { if !strings.Contains(helpView, "q/ctrl+c") { t.Fatalf("expected help to include quit shortcut, got %q", helpView) } + if !strings.Contains(helpView, "retest all") { + t.Fatalf("expected help to include retest-all shortcut, got %q", helpView) + } if strings.Contains(helpView, "esc") { t.Fatalf("expected help to hide detail shortcut when detail is closed, got %q", helpView) } + if strings.Contains(helpView, "retest node") { + t.Fatalf("expected help to hide retest-node shortcut when detail is closed, got %q", helpView) + } model.toggleDetail(result) helpView = model.help.view() if !strings.Contains(helpView, "esc") { t.Fatalf("expected help to include detail shortcut when detail is visible, got %q", helpView) } + if !strings.Contains(helpView, "retest node") { + t.Fatalf("expected help to include retest-node shortcut when detail is visible, got %q", helpView) + } + if strings.Contains(helpView, "retest all") { + t.Fatalf("expected help to hide retest-all shortcut when detail is visible, got %q", helpView) + } } diff --git a/tui/model.go b/tui/model.go index e3fed4b..13af07a 100644 --- a/tui/model.go +++ b/tui/model.go @@ -55,6 +55,9 @@ type tuiModel struct { flushScheduled bool detailHeight int perf *perfTracker + retestAll func(chan<- *speedtester.Result) + retestOne func(string, chan<- *speedtester.Result) + retestingName string } const ( @@ -126,9 +129,18 @@ func NewTUIModel(mode speedtester.SpeedMode, totalProxies int, resultChannel cha flushScheduled: false, detailHeight: 0, perf: newPerfTracker(), + retestingName: "", } } +func (m *tuiModel) SetRetestCallbacks( + retestAll func(chan<- *speedtester.Result), + retestOne func(string, chan<- *speedtester.Result), +) { + m.retestAll = retestAll + m.retestOne = retestOne +} + // Init initializes the TUI model func (m tuiModel) Init() tea.Cmd { return tea.Batch( @@ -141,7 +153,7 @@ func (m tuiModel) Init() tea.Cmd { func (m tuiModel) waitForResult() tea.Cmd { return func() tea.Msg { result, ok := <-m.resultChannel - if !ok { + if !ok || result == nil { return doneMsg{} } return resultMsg{result: result} @@ -181,6 +193,23 @@ func (m tuiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case "q", "ctrl+c": m.quitting = true return m, tea.Quit + case "r": + if m.testing { + return m, nil + } + if m.detailVisible && m.detailResult != nil { + if m.retestOne == nil { + return m, nil + } + m.testing = true + m.retestingName = m.detailResult.ProxyName + return m, tea.Batch(m.runRetestOneCmd(m.detailResult.ProxyName), m.waitForResult()) + } + if m.retestAll == nil { + return m, nil + } + m.beginFullRetest() + return m, tea.Batch(m.progress.SetPercent(0), m.runRetestAllCmd(), m.waitForResult()) } m.table, cmd = m.table.Update(msg) @@ -229,11 +258,19 @@ func (m tuiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, nil case resultMsg: - m.currentProxy++ - m.results = append(m.results, msg.result) - m.recordSequence(msg.result) + if m.retestingName != "" { + m.replaceRetestedResult(msg.result) + } else { + m.currentProxy++ + m.results = append(m.results, msg.result) + m.recordSequence(msg.result) + } m.resultsDirty = true - progressCmd := m.progress.SetPercent(float64(m.currentProxy) / float64(m.totalProxies)) + percent := 0.0 + if m.totalProxies > 0 { + percent = float64(m.currentProxy) / float64(m.totalProxies) + } + progressCmd := m.progress.SetPercent(percent) cmds := []tea.Cmd{progressCmd, m.waitForResult()} if !m.flushScheduled { m.flushScheduled = true @@ -244,6 +281,7 @@ func (m tuiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case doneMsg: m.testing = false m.flushScheduled = false + m.retestingName = "" m.flushResultsIfDirty() progressCmd := m.progress.SetPercent(1.0) return m, progressCmd @@ -277,6 +315,62 @@ func (m tuiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, cmd } +func (m *tuiModel) beginFullRetest() { + m.currentProxy = 0 + m.results = make([]*speedtester.Result, 0, m.totalProxies) + m.sequence = make(map[*speedtester.Result]int) + m.nextSequence = 0 + m.testing = true + m.startTime = time.Now() + m.resultsDirty = false + m.flushScheduled = false + m.retestingName = "" + m.detailVisible = false + m.detailResult = nil + m.detailHeight = 0 + m.selectedIndex = -1 + m.help.setDetailVisible(false) + m.table.SetRows([]table.Row{}) + m.table.SetCursor(0) + m.updateTableLayout() +} + +func (m tuiModel) runRetestAllCmd() tea.Cmd { + return func() tea.Msg { + m.retestAll(m.resultChannel) + return nil + } +} + +func (m tuiModel) runRetestOneCmd(name string) tea.Cmd { + return func() tea.Msg { + m.retestOne(name, m.resultChannel) + return nil + } +} + +func (m *tuiModel) replaceRetestedResult(result *speedtester.Result) { + for i, existing := range m.results { + if existing.ProxyName != m.retestingName { + continue + } + m.results[i] = result + if seq, ok := m.sequence[existing]; ok { + m.sequence[result] = seq + } else { + m.recordSequence(result) + } + delete(m.sequence, existing) + if m.detailResult == existing { + m.detailResult = result + } + return + } + m.currentProxy++ + m.results = append(m.results, result) + m.recordSequence(result) +} + // View renders the TUI func (m tuiModel) View() string { if m.quitting { diff --git a/tui/model_update_test.go b/tui/model_update_test.go index 953ab0e..2f6d085 100644 --- a/tui/model_update_test.go +++ b/tui/model_update_test.go @@ -4,6 +4,8 @@ import ( "testing" "time" + "github.com/charmbracelet/bubbles/progress" + tea "github.com/charmbracelet/bubbletea" "github.com/faceair/clash-speedtest/speedtester" ) @@ -211,3 +213,145 @@ func TestTUIModelUpdateFastMode(t *testing.T) { t.Error("Expected third result to be result1") } } + +func TestTUIModelRetestAllResetsState(t *testing.T) { + resultChannel := make(chan *speedtester.Result, 4) + model := NewTUIModel(speedtester.SpeedModeDownload, 2, resultChannel) + model.testing = false + model.currentProxy = 2 + model.progress.SetPercent(1.0) + model.results = []*speedtester.Result{ + {ProxyName: "Proxy 1", ProxyType: "SS", ProxyConfig: map[string]any{}}, + {ProxyName: "Proxy 2", ProxyType: "Trojan", ProxyConfig: map[string]any{}}, + } + model.updateTableRows() + called := false + model.SetRetestCallbacks( + func(out chan<- *speedtester.Result) { + called = true + out <- nil + }, + nil, + ) + + updated, cmd := model.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'r'}}) + updatedModel := updated.(tuiModel) + if cmd == nil { + t.Fatalf("expected retest-all to return a command") + } + progressCmd := updatedModel.progress.SetPercent(0) + progressUpdated, _ := updatedModel.progress.Update(progressCmd()) + updatedModel.progress = progressUpdated.(progress.Model) + updatedModel.runRetestAllCmd()() + if !called { + t.Fatalf("expected retest-all callback to run") + } + if !updatedModel.testing { + t.Fatalf("expected model to enter testing state") + } + if updatedModel.currentProxy != 0 { + t.Fatalf("expected currentProxy reset, got %d", updatedModel.currentProxy) + } + if len(updatedModel.results) != 0 { + t.Fatalf("expected results to clear, got %d", len(updatedModel.results)) + } + if updatedModel.selectedIndex != -1 { + t.Fatalf("expected selection to reset, got %d", updatedModel.selectedIndex) + } + if updatedModel.progress.Percent() != 0 { + t.Fatalf("expected progress percent reset to 0, got %f", updatedModel.progress.Percent()) + } +} + +func TestTUIModelRetestOneReplacesExistingResult(t *testing.T) { + resultChannel := make(chan *speedtester.Result, 4) + model := NewTUIModel(speedtester.SpeedModeDownload, 2, resultChannel) + oldSlow := &speedtester.Result{ + ProxyName: "Slow", + ProxyType: "SS", + Latency: 100 * time.Millisecond, + DownloadSpeed: 2 * 1024 * 1024, + ProxyConfig: map[string]any{}, + } + oldFast := &speedtester.Result{ + ProxyName: "Fast", + ProxyType: "Trojan", + Latency: 200 * time.Millisecond, + DownloadSpeed: 10 * 1024 * 1024, + ProxyConfig: map[string]any{}, + } + model.results = []*speedtester.Result{oldFast, oldSlow} + model.recordSequence(oldFast) + model.recordSequence(oldSlow) + model.testing = false + model.detailVisible = true + model.detailResult = oldSlow + model.selectedIndex = 1 + model.updateTableRows() + + newSlow := &speedtester.Result{ + ProxyName: "Slow", + ProxyType: "SS", + Latency: 90 * time.Millisecond, + DownloadSpeed: 20 * 1024 * 1024, + ProxyConfig: map[string]any{}, + } + calledName := "" + model.SetRetestCallbacks(nil, func(name string, out chan<- *speedtester.Result) { + calledName = name + out <- newSlow + out <- nil + }) + + updated, cmd := model.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'r'}}) + updatedModel := updated.(tuiModel) + if cmd == nil { + t.Fatalf("expected retest-one to return a command") + } + updatedModel.runRetestOneCmd(updatedModel.retestingName)() + if calledName != "Slow" { + t.Fatalf("expected retest-one callback to receive selected node, got %q", calledName) + } + + updated, _ = updatedModel.Update(resultMsg{result: newSlow}) + updatedModel = updated.(tuiModel) + updated, _ = updatedModel.Update(flushResultsMsg{}) + updatedModel = updated.(tuiModel) + + if len(updatedModel.results) != 2 { + t.Fatalf("expected result count to stay the same, got %d", len(updatedModel.results)) + } + if updatedModel.results[0] != newSlow { + t.Fatalf("expected retested node to be resorted to first position") + } + if updatedModel.detailResult != newSlow { + t.Fatalf("expected detail result to follow replacement") + } + if updatedModel.selectedIndex != 0 { + t.Fatalf("expected selection to follow resorted node, got %d", updatedModel.selectedIndex) + } + if updatedModel.currentProxy != 0 { + t.Fatalf("expected currentProxy to remain unchanged during single retest, got %d", updatedModel.currentProxy) + } +} + +func TestTUIModelIgnoresRetestWhileTesting(t *testing.T) { + resultChannel := make(chan *speedtester.Result, 1) + model := NewTUIModel(speedtester.SpeedModeDownload, 1, resultChannel) + called := false + model.SetRetestCallbacks( + func(chan<- *speedtester.Result) { called = true }, + func(string, chan<- *speedtester.Result) { called = true }, + ) + + updated, cmd := model.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'r'}}) + if called { + t.Fatalf("expected retest callbacks not to run while testing") + } + if cmd != nil { + t.Fatalf("expected no command when retest is ignored") + } + if !updated.(tuiModel).testing { + t.Fatalf("expected testing state to remain true") + } +}