From 1dd971e16b6dffb4de6e7dd9cdd8c4d4330fe3bc Mon Sep 17 00:00:00 2001 From: NotAShelf Date: Wed, 18 Jun 2025 14:01:26 +0300 Subject: [PATCH] treewide: refactor; defer cli handling to cobra --- cmd/root.go | 91 +++++++++++++++++++++++++++++++++++++ go.mod | 10 ++-- go.sum | 18 ++++---- internal/tct/client.go | 32 +++++++++++++ internal/tct/client_test.go | 46 +++++++++++++++++++ main.go | 77 ++----------------------------- main_test.go | 62 ------------------------- 7 files changed, 185 insertions(+), 151 deletions(-) create mode 100644 cmd/root.go create mode 100644 internal/tct/client.go create mode 100644 internal/tct/client_test.go delete mode 100644 main_test.go diff --git a/cmd/root.go b/cmd/root.go new file mode 100644 index 0000000..4b4a7f1 --- /dev/null +++ b/cmd/root.go @@ -0,0 +1,91 @@ +package cmd + +import ( + "context" + "fmt" + "os" + "os/signal" + "sync" + "syscall" + "time" + + "github.com/spf13/cobra" + "notashelf.dev/tct/internal/tct" +) + +var ( + url string + maxRequests int + delay time.Duration + Version string // will be set by main package +) + +var rootCmd = &cobra.Command{ + Use: "tct", + Short: "TCP Connection Timer - find optimal parallel request count", + Long: `A tool to measure and find the optimal number of parallel TCP requests for a given URL.`, + Run: run, +} + +func Execute() { + if Version != "" { + rootCmd.Version = Version + } + if err := rootCmd.Execute(); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} + +func init() { + rootCmd.Flags().StringVarP(&url, "url", "u", "http://example.com", "URL to fetch") + rootCmd.Flags().IntVarP(&maxRequests, "max", "m", 100, "Maximum number of parallel requests") + rootCmd.Flags().DurationVarP(&delay, "delay", "d", 0, "Delay between requests") +} + +func run(cmd *cobra.Command, args []string) { + client := tct.NewClient() + + // Context that listens for the interrupt signal + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + optimal := 0 + minTime := time.Duration(0) + + for i := 1; i <= maxRequests; i++ { + // Check for cancellation before starting the loop + if ctx.Err() != nil { + break + } + + start := time.Now() + wg := &sync.WaitGroup{} + wg.Add(i) + + for range i { + go func() { + defer wg.Done() + select { + case <-ctx.Done(): + fmt.Println("Operation cancelled.") + return + default: + time.Sleep(delay) + client.MakeRequest(url) + } + }() + } + + wg.Wait() + duration := time.Since(start) + fmt.Printf("Parallel Requests: %d, Time Taken: %s\n", i, duration) + + if minTime == 0 || duration < minTime { + minTime = duration + optimal = i + } + } + + fmt.Printf("\nOptimal Number of Parallel TCP Requests: %d\n", optimal) +} diff --git a/go.mod b/go.mod index 8bff24f..ee35e6b 100644 --- a/go.mod +++ b/go.mod @@ -1,12 +1,10 @@ module notashelf.dev/tct -go 1.22.2 +go 1.24.3 -require github.com/stretchr/testify v1.9.0 +require github.com/spf13/cobra v1.9.1 require ( - github.com/davecgh/go-spew v1.1.1 // indirect - github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/stretchr/objx v0.5.2 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect + github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/spf13/pflag v1.0.6 // indirect ) diff --git a/go.sum b/go.sum index 8565daa..ffae55e 100644 --- a/go.sum +++ b/go.sum @@ -1,12 +1,10 @@ -github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= -github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= -github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= -github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= -github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= -github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= -gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= +github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= +github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/spf13/cobra v1.9.1 h1:CXSaggrXdbHK9CF+8ywj8Amf7PBRmPCOJugH954Nnlo= +github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0= +github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o= +github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= -gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/tct/client.go b/internal/tct/client.go new file mode 100644 index 0000000..3c198e8 --- /dev/null +++ b/internal/tct/client.go @@ -0,0 +1,32 @@ +package tct + +import ( + "fmt" + "io" + "net/http" +) + +type HttpClient interface { + Get(url string) (*http.Response, error) +} + +type Client struct { + httpClient HttpClient +} + +func NewClient() *Client { + return &Client{ + httpClient: &http.Client{}, + } +} + +func (c *Client) MakeRequest(url string) { + resp, err := c.httpClient.Get(url) + if err != nil { + fmt.Println(err) + return + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + _ = body +} diff --git a/internal/tct/client_test.go b/internal/tct/client_test.go new file mode 100644 index 0000000..09efacb --- /dev/null +++ b/internal/tct/client_test.go @@ -0,0 +1,46 @@ +package tct + +import ( + "net/http" + "testing" +) + +type MockHttpClient struct { + GetFunc func(url string) (*http.Response, error) +} + +func (m *MockHttpClient) Get(url string) (*http.Response, error) { + return m.GetFunc(url) +} + +func TestClientMakeRequest(t *testing.T) { + tests := []struct { + name string + url string + }{ + { + name: "valid URL", + url: "http://example.com", + }, + { + name: "another valid URL", + url: "https://google.com", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mockClient := &MockHttpClient{ + GetFunc: func(url string) (*http.Response, error) { + return &http.Response{ + StatusCode: 200, + Body: http.NoBody, + }, nil + }, + } + + client := &Client{httpClient: mockClient} + client.MakeRequest(tt.url) + }) + } +} diff --git a/main.go b/main.go index 25c69aa..d854c8e 100644 --- a/main.go +++ b/main.go @@ -1,79 +1,10 @@ package main -import ( - "context" - "flag" - "fmt" - "io" - "net/http" - "os" - "os/signal" - "sync" - "syscall" - "time" -) +import "notashelf.dev/tct/cmd" -type HttpClient interface { - Get(url string) (*http.Response, error) -} - -var httpClient HttpClient = &http.Client{} +var version = "dev" // will be set by build process func main() { - urlPtr := flag.String("url", "http://example.com", "URL to fetch") - maxRequestsPtr := flag.Int("max", 100, "Maximum number of parallel requests") - delayPtr := flag.Duration("delay", 0, "Delay between requests") - - flag.Parse() - url := *urlPtr - maxRequests := *maxRequestsPtr - delay := *delayPtr - - // context that listens for the interrupt signal. - ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) - defer stop() - - optimal := 0 - minTime := time.Duration(0) - for i := 1; i <= maxRequests; i++ { - // check for cancellation before starting the loop - if ctx.Err() != nil { - break - } - start := time.Now() - wg := &sync.WaitGroup{} - wg.Add(i) - for j := 0; j < i; j++ { - go func() { - defer wg.Done() - select { - case <-ctx.Done(): - fmt.Println("Operation cancelled.") - return - default: - time.Sleep(delay) - makeRequest(httpClient, url) - } - }() - } - wg.Wait() - duration := time.Since(start) - fmt.Printf("Parallel Requests: %d, Time Taken: %s\n", i, duration) - if minTime == 0 || duration < minTime { - minTime = duration - optimal = i - } - } - fmt.Printf("\nOptimal Number of Parallel TCP Requests: %d\n", optimal) -} - -func makeRequest(client HttpClient, url string) { - resp, err := client.Get(url) - if err != nil { - fmt.Println(err) - return - } - defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) - _ = body + cmd.Version = version + cmd.Execute() } diff --git a/main_test.go b/main_test.go deleted file mode 100644 index b6186c5..0000000 --- a/main_test.go +++ /dev/null @@ -1,62 +0,0 @@ -package main - -import ( - "bytes" - "errors" - "io" - "net/http" - "net/http/httptest" - "testing" - - "github.com/stretchr/testify/mock" -) - -type MockHttpClient struct { - mock.Mock -} - -func (m *MockHttpClient) Get(url string) (*http.Response, error) { - args := m.Called(url) - return args.Get(0).(*http.Response), args.Error(1) -} - -func TestMakeRequest(t *testing.T) { - tests := []struct { - name string - statusCode int - responseBody string - expectedError error - }{ - { - name: "Successful request", - statusCode: http.StatusOK, - responseBody: "OK", - }, - { - name: "Failed request", - statusCode: http.StatusInternalServerError, - responseBody: "", - expectedError: errors.New("request failed"), - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - mockClient := new(MockHttpClient) - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(tt.statusCode) - w.Write([]byte(tt.responseBody)) - })) - defer server.Close() - - mockClient.On("Get", server.URL).Return(&http.Response{ - StatusCode: tt.statusCode, - Body: io.NopCloser(bytes.NewBufferString(tt.responseBody)), - }, tt.expectedError) - - makeRequest(mockClient, server.URL) - - mockClient.AssertExpectations(t) - }) - } -}