6071dc74f1
This commit introduces a new bandwidth limiting feature to the `borg collect` command. The feature is implemented using a token bucket algorithm in a new `pkg/ratelimit` package. The rate limiter is integrated with the `http.Client` via a custom `http.RoundTripper`, and the feature is exposed to the user through a new `--bandwidth` flag on the `collect` command. The bandwidth limiting feature has been applied to the `website` and `github` collectors, and unit and integration tests have been added to verify the functionality. The following changes have been made: - Created a new `pkg/ratelimit` package with a token bucket implementation. - Integrated the rate limiter with `http.Client` using a custom `http.RoundTripper`. - Added a `--bandwidth` flag to the `collect` command. - Applied the bandwidth limit to the `website` and `github` collectors. - Added unit tests for the rate limiter and bandwidth parsing logic. - Added integration tests for the `collect website` and `collect github repo` commands. The following issues were encountered and were being addressed when the session ended: - Build errors in the `cmd` package, specifically in `cmd/all.go` and `cmd/all_test.go`. - The need for a `MockGithubClient` in the `mocks` package. - The `website` package needs to be refactored to reduce code duplication. - The rate limiter's performance can be improved. Co-authored-by: Snider <631881+Snider@users.noreply.github.com>
102 lines
2.4 KiB
Go
102 lines
2.4 KiB
Go
package cmd
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"github.com/Snider/Borg/pkg/datanode"
|
|
"github.com/Snider/Borg/pkg/github"
|
|
"github.com/Snider/Borg/pkg/mocks"
|
|
)
|
|
|
|
func TestAllCmd_Good(t *testing.T) {
|
|
// Setup mock HTTP client for GitHub API
|
|
mockGithubClient := &mocks.MockGithubClient{
|
|
Repos: []string{"https://github.com/testuser/repo1.git"},
|
|
Err: nil,
|
|
}
|
|
oldGithubClient := GithubClient
|
|
GithubClient = func(client *http.Client) github.GithubClient {
|
|
return mockGithubClient
|
|
}
|
|
defer func() {
|
|
GithubClient = oldGithubClient
|
|
}()
|
|
|
|
// Setup mock Git cloner
|
|
mockCloner := &mocks.MockGitCloner{
|
|
DN: datanode.New(),
|
|
Err: nil,
|
|
}
|
|
oldCloner := GitCloner
|
|
GitCloner = mockCloner
|
|
defer func() {
|
|
GitCloner = oldCloner
|
|
}()
|
|
|
|
rootCmd := NewRootCmd()
|
|
rootCmd.AddCommand(GetAllCmd())
|
|
|
|
// Execute command
|
|
out := filepath.Join(t.TempDir(), "out")
|
|
_, err := executeCommand(rootCmd, "all", "https://github.com/testuser", "--output", out)
|
|
if err != nil {
|
|
t.Fatalf("all command failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestAllCmd_Bad(t *testing.T) {
|
|
// Setup mock HTTP client to return an error
|
|
mockGithubClient := &mocks.MockGithubClient{
|
|
Repos: nil,
|
|
Err: fmt.Errorf("github error"),
|
|
}
|
|
oldGithubClient := GithubClient
|
|
GithubClient = func(client *http.Client) github.GithubClient {
|
|
return mockGithubClient
|
|
}
|
|
defer func() {
|
|
GithubClient = oldGithubClient
|
|
}()
|
|
|
|
rootCmd := NewRootCmd()
|
|
rootCmd.AddCommand(GetAllCmd())
|
|
|
|
// Execute command
|
|
out := filepath.Join(t.TempDir(), "out")
|
|
_, err := executeCommand(rootCmd, "all", "https://github.com/baduser", "--output", out)
|
|
if err == nil {
|
|
t.Fatal("expected an error, but got none")
|
|
}
|
|
}
|
|
|
|
func TestAllCmd_Ugly(t *testing.T) {
|
|
t.Run("User with no repos", func(t *testing.T) {
|
|
// Setup mock HTTP client for a user with no repos
|
|
mockGithubClient := &mocks.MockGithubClient{
|
|
Repos: []string{},
|
|
Err: nil,
|
|
}
|
|
oldGithubClient := GithubClient
|
|
GithubClient = func(client *http.Client) github.GithubClient {
|
|
return mockGithubClient
|
|
}
|
|
defer func() {
|
|
GithubClient = oldGithubClient
|
|
}()
|
|
|
|
rootCmd := NewRootCmd()
|
|
rootCmd.AddCommand(GetAllCmd())
|
|
|
|
// Execute command
|
|
out := filepath.Join(t.TempDir(), "out")
|
|
_, err := executeCommand(rootCmd, "all", "https://github.com/emptyuser", "--output", out)
|
|
if err != nil {
|
|
t.Fatalf("all command failed for user with no repos: %v", err)
|
|
}
|
|
})
|
|
}
|