diff --git a/cmd/add.go b/cmd/add.go index 5448f76..0448c6e 100644 --- a/cmd/add.go +++ b/cmd/add.go @@ -21,7 +21,7 @@ func addAdd(cmd *cobra.Command) { addCmd.Flags().String("algo", internal.RecommendedAlgo, "Integrity algorithm") addCmd.Flags().String("filename", "", "Target file name to use when downloading the resource") addCmd.Flags().StringArray("tag", []string{}, "Resource tags") - addCmd.Flags().String("cache", "", "Artifactory cache URL") + addCmd.Flags().String("artifactory-cache-url", "", "Artifactory cache URL") cmd.AddCommand(addCmd) } @@ -31,15 +31,14 @@ func runAdd(cmd *cobra.Command, args []string) error { return err } // Get cache URL - cacheURL, err := cmd.Flags().GetString("cache") + ArtifactoryCacheURL, err := cmd.Flags().GetString("artifactory-cache-url") if err != nil { return err } - // Check token if cache is requested - if cacheURL != "" { - token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") + if ArtifactoryCacheURL != "" { + token := os.Getenv(internal.GRABIT_ARTIFACTORY_TOKEN_ENV_VAR) if token == "" { - return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set") + return fmt.Errorf("%s environment variable is not set", internal.GRABIT_ARTIFACTORY_TOKEN_ENV_VAR) } } lock, err := internal.NewLock(lockFile, true) @@ -58,7 +57,7 @@ func runAdd(cmd *cobra.Command, args []string) error { if err != nil { return err } - err = lock.AddResource(args, algo, tags, filename, cacheURL) + err = lock.AddResource(args, algo, tags, filename, ArtifactoryCacheURL) if err != nil { return err } diff --git a/cmd/add_test.go b/cmd/add_test.go index 8ec01b2..cdb97a6 100644 --- a/cmd/add_test.go +++ b/cmd/add_test.go @@ -2,49 +2,27 @@ package cmd import ( "fmt" - "net/http" "testing" + "github.com/cisco-open/grabit/internal" "github.com/cisco-open/grabit/test" "github.com/stretchr/testify/assert" ) func TestRunAdd(t *testing.T) { - // Set the GRABIT_ARTIFACTORY_TOKEN environment variable. - t.Setenv("GRABIT_ARTIFACTORY_TOKEN", "test-token") - - // Setup HTTP handler for the resource - handler := func(w http.ResponseWriter, r *http.Request) { - _, err := w.Write([]byte(`abcdef`)) - if err != nil { - t.Fatal(err) - } - } - port, server := test.HttpHandler(handler) - defer server.Close() - - // Setup dummy cache server - cacheHandler := func(w http.ResponseWriter, r *http.Request) { - if r.Method == "PUT" { - w.WriteHeader(http.StatusCreated) - } - } - cachePort, cacheServer := test.HttpHandler(cacheHandler) - defer cacheServer.Close() - - // Create empty lockfile - lockFile := test.TmpFile(t, "") - + port := test.TestHttpHandler("abcdef", t) cmd := NewRootCmd() - // Add cache URL to the command - cacheURL := fmt.Sprintf("http://localhost:%d", cachePort) - cmd.SetArgs([]string{ - "-f", lockFile, - "add", - fmt.Sprintf("http://localhost:%d/test.html", port), - "--cache", cacheURL, - }) + cmd.SetArgs([]string{"-f", test.TmpFile(t, ""), "add", fmt.Sprintf("http://localhost:%d/test.html", port)}) + err := cmd.Execute() + assert.Nil(t, err) +} +func TestRunAddWithArtifactoryCache(t *testing.T) { + t.Setenv(internal.GRABIT_ARTIFACTORY_TOKEN_ENV_VAR, "artifactory-token") + port := test.TestHttpHandler("abcdef", t) + artPort := test.TestHttpHandler("abcdef", t) + cmd := NewRootCmd() + cmd.SetArgs([]string{"-f", test.TmpFile(t, ""), "add", fmt.Sprintf("http://localhost:%d/test.html", port), "--artifactory-cache-url", fmt.Sprintf("http://localhost:%d/artifactory", artPort)}) err := cmd.Execute() assert.Nil(t, err) } diff --git a/cmd/delete.go b/cmd/delete.go index 27c5722..769bd45 100644 --- a/cmd/delete.go +++ b/cmd/delete.go @@ -28,7 +28,10 @@ func runDel(cmd *cobra.Command, args []string) error { return err } for _, r := range args { - lock.DeleteResource(r) + err = lock.DeleteResource(r) + if err != nil { + return err + } } err = lock.Save() if err != nil { diff --git a/internal/artifactory_cache_test.go b/internal/artifactory_cache_test.go deleted file mode 100644 index ef19ded..0000000 --- a/internal/artifactory_cache_test.go +++ /dev/null @@ -1,91 +0,0 @@ -package internal - -import ( - "fmt" - "net/http" - "testing" - - "github.com/cisco-open/grabit/test" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// TestAddWithArtifactoryCache verifies adding a resource with caching enabled. -func TestAddWithArtifactoryCacheTokenNotSet(t *testing.T) { - // Setup a simple HTTP handler that always returns "test content". - handler := func(w http.ResponseWriter, r *http.Request) { - w.Write([]byte(`test content`)) - } - // Start the HTTP server and get the port it runs on. - port, server := test.HttpHandler(handler) - defer server.Close() // Ensure the server is stopped after the test. - - // Create a temporary lock file for testing. - path := test.TmpFile(t, "") - lock, err := NewLock(path, true) - require.NoError(t, err) // Fail the test if the lock cannot be created. - - // Set up URLs for the source file and cache. - sourceURL := fmt.Sprintf("http://localhost:%d/test.txt", port) - cacheURL := fmt.Sprintf("http://localhost:%d", port) - - // Attempt to add a resource to the lock file. - err = lock.AddResource([]string{sourceURL}, "sha256", []string{}, "", cacheURL) - // Verify that the error message indicates the token is not set. - assert.Contains(t, err.Error(), "GRABIT_ARTIFACTORY_TOKEN environment variable is not set") -} - -// TestDeleteWithArtifactoryCache verifies deleting a resource with caching enabled. -func TestDeleteWithArtifactoryCache(t *testing.T) { - // Set the GRABIT_ARTIFACTORY_TOKEN environment variable. - t.Setenv("GRABIT_ARTIFACTORY_TOKEN", "test-token") - - // Setup an HTTP handler to handle DELETE requests. - handler := func(w http.ResponseWriter, r *http.Request) { - if r.Method == "DELETE" { // Respond with OK for DELETE requests. - w.WriteHeader(http.StatusOK) - } - } - // Start the HTTP server and get the port it runs on. - port, server := test.HttpHandler(handler) - defer server.Close() // Ensure the server is stopped after the test. - - // Set up URLs for the source file and cache. - sourceURL := fmt.Sprintf("http://localhost:%d/test.txt", port) - cacheURL := fmt.Sprintf("http://localhost:%d", port) - - // Create a lock file with the resource and cache information. - lockContent := fmt.Sprintf(`[[Resource]] - Urls = ['%s'] - Integrity = 'sha256-test' - CacheUri = '%s'`, sourceURL, cacheURL) - - lockPath := test.TmpFile(t, lockContent) - lock, err := NewLock(lockPath, false) - require.NoError(t, err) // Fail the test if the lock cannot be created. - - // Save the lock file before modifying it. - err = lock.Save() - require.NoError(t, err) // Fail the test if saving fails. - - // Delete the resource from the lock file. - lock.DeleteResource(sourceURL) - - // Checks that http.StatusOK was returned meaning the delete was successful - resp, err := http.Get(server.URL) - if err != nil { - require.NoError(t, err) - } - if resp.StatusCode != http.StatusOK { - t.Errorf("expected 200 but got %d", resp.StatusCode) - } - - // Save the lock file again after deletion. - err = lock.Save() - require.NoError(t, err) // Fail the test if saving fails. - - // Reload the lock file and verify the resource is gone. - newLock, err := NewLock(lockPath, false) - require.NoError(t, err) - assert.Equal(t, 0, len(newLock.conf.Resource)) // Ensure no resources remain. -} diff --git a/internal/lock.go b/internal/lock.go index 9f47acb..66a5120 100644 --- a/internal/lock.go +++ b/internal/lock.go @@ -8,26 +8,14 @@ import ( "context" "errors" "fmt" - "net/http" "os" - "path" - "path/filepath" "strconv" - "strings" - "github.com/carlmjohnson/requests" toml "github.com/pelletier/go-toml/v2" - "github.com/rs/zerolog/log" ) var COMMENT_PREFIX = "//" -// getArtifactoryURL constructs the URL for an artifact in Artifactory -func getArtifactoryURL(baseURL, integrity string) string { - return fmt.Sprintf("%s/%s", baseURL, integrity) - -} - // Lock represents a grabit lockfile. type Lock struct { path string @@ -68,90 +56,29 @@ func (l *Lock) AddResource(paths []string, algo string, tags []string, filename } } r, err := NewResourceFromUrl(paths, algo, tags, filename, cacheURL) + if err != nil { return err } - // If cache URL is provided, handles Artifactory upload - if cacheURL != "" { - token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") - if token == "" { - return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set") - } - - // Add context here - ctx := context.Background() - path, err := GetUrltoTempFile(paths[0], token, ctx) - if err != nil { - return fmt.Errorf("failed to get file for cache: %s", err) - } - defer os.Remove(path) - - // Upload to Artifactory using hash as filename - err = uploadToArtifactory(path, cacheURL, r.Integrity) - if err != nil { - return fmt.Errorf("failed to upload to cache: %v", err) - } - } - l.conf.Resource = append(l.conf.Resource, *r) return nil } -func uploadToArtifactory(filePath, cacheURL, integrity string) error { - token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") - if token == "" { - return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set") - } - - // Use the integrity value directly for the URL - artifactoryURL := getArtifactoryURL(cacheURL, integrity) - - // Upload the file using the requests package - err := requests. - URL(artifactoryURL). - Method(http.MethodPut). - Header("Authorization", fmt.Sprintf("Bearer %s", token)). - BodyFile(filePath). // Using BodyFile instead of ReadFile - Fetch(context.Background()) - - if err != nil { - return fmt.Errorf("upload failed: %w", err) - } - - return nil -} -func (l *Lock) DeleteResource(path string) { +func (l *Lock) DeleteResource(path string) error { newStatements := []Resource{} for _, r := range l.conf.Resource { if !r.Contains(path) { newStatements = append(newStatements, r) - } else if r.Contains(path) && r.CacheUri != "" { - token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") - if token == "" { - log.Warn().Msg("Warning: Unable to delete from Artifactory: GRABIT_ARTIFACTORY_TOKEN not set.") - - continue - } - - artifactoryURL := getArtifactoryURL(r.CacheUri, r.Integrity) - - err := deleteCache(artifactoryURL, token) + } else { + err := r.Delete() if err != nil { - log.Warn().Msg("Warning: Unable to delete from Artifactory") + return fmt.Errorf("Failed to delete resource '%s': %w", path, err) } } } l.conf.Resource = newStatements -} - -func deleteCache(url, token string) error { - // Create and send a DELETE request with an Authorization header. - return requests. - URL(url). - Method(http.MethodDelete). - Header("Authorization", fmt.Sprintf("Bearer %s", token)). - Fetch(context.Background()) + return nil } const NoFileMode = os.FileMode(0) @@ -246,47 +173,6 @@ func (l *Lock) Download(dir string, tags []string, notags []string, perm string, for i, r := range filteredResources { resource := r go func() { - // Try Artifactory first if available - if resource.CacheUri != "" { - token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") - if token != "" { - artifactoryURL := getArtifactoryURL(resource.CacheUri, resource.Integrity) - filename := resource.Filename - if filename == "" { - filename = path.Base(resource.Urls[0]) - } - fullPath := filepath.Join(dir, filename) - - // Use getUrl with bearer token - tmpPath, err := getUrl(artifactoryURL, fullPath, token, ctx) - if err == nil { - // integrity check - algo, err := getAlgoFromIntegrity(resource.Integrity) - if err != nil { - errorCh <- err - return - } - err = checkIntegrityFromFile(tmpPath, algo, resource.Integrity, artifactoryURL) - if err != nil { - errorCh <- err - return - } - if mode != NoFileMode { - err = os.Chmod(tmpPath, mode.Perm()) - } - if err == nil { - errorCh <- nil - if statusLine != nil { - statusLine.Increment(i) - } - return - } - } - if strings.Contains(err.Error(), "lookup invalid") || strings.Contains(err.Error(), "dial tcp") { - fmt.Printf("Failed to download from Artifactory, falling back to original URL: %v\n", err) - } - } - } err := resource.Download(dir, mode, ctx) errorCh <- err diff --git a/internal/lock_test.go b/internal/lock_test.go index f7ed791..11a92f7 100644 --- a/internal/lock_test.go +++ b/internal/lock_test.go @@ -56,7 +56,8 @@ func TestLockManipulations(t *testing.T) { assert.Equal(t, 2, len(lock.conf.Resource)) err = lock.Save() assert.Nil(t, err) - lock.DeleteResource(resource) + err = lock.DeleteResource(resource) + assert.Nil(t, err) assert.Equal(t, 1, len(lock.conf.Resource)) } diff --git a/internal/resource.go b/internal/resource.go index d93d01f..dbe5c3d 100644 --- a/internal/resource.go +++ b/internal/resource.go @@ -8,6 +8,7 @@ import ( "crypto/sha256" "encoding/hex" "fmt" + "net/http" "net/url" "os" "path" @@ -20,14 +21,20 @@ import ( // Resource represents an external resource to be downloaded. type Resource struct { - Urls []string - Integrity string - Tags []string `toml:",omitempty"` - Filename string `toml:",omitempty"` - CacheUri string `toml:",omitempty"` + Urls []string + Integrity string + Tags []string `toml:",omitempty"` + Filename string `toml:",omitempty"` + ArtifactoryCacheURL string `toml:",omitempty"` } -func NewResourceFromUrl(urls []string, algo string, tags []string, filename string, cacheURL string) (*Resource, error) { +const GRABIT_ARTIFACTORY_TOKEN_ENV_VAR = "GRABIT_ARTIFACTORY_TOKEN" + +func getArtifactoryToken() string { + return os.Getenv(GRABIT_ARTIFACTORY_TOKEN_ENV_VAR) +} + +func NewResourceFromUrl(urls []string, algo string, tags []string, filename string, ArtifactoryCacheURL string) (*Resource, error) { if len(urls) < 1 { return nil, fmt.Errorf("empty url list") } @@ -42,8 +49,57 @@ func NewResourceFromUrl(urls []string, algo string, tags []string, filename stri if err != nil { return nil, fmt.Errorf("failed to compute ressource integrity: %s", err) } + resource := &Resource{Urls: urls, Integrity: integrity, Tags: tags, Filename: filename, ArtifactoryCacheURL: ArtifactoryCacheURL} + // If cache URL is provided, upload file to Artifactory. + if resource.ArtifactoryCacheURL != "" { + err := resource.AddToCache(path) + if err != nil { + return nil, fmt.Errorf("failed to upload to cache: %s", err) + } + } + return resource, nil +} + +func (l *Resource) getCacheFullURL() string { + url, err := url.JoinPath(l.ArtifactoryCacheURL, l.Integrity) + if err != nil { + log.Fatal().Err(err) + } + return url +} + +func (l *Resource) AddToCache(filePath string) error { + token := getArtifactoryToken() + if token == "" { + return fmt.Errorf("%s environment variable is not set and is needed to upload to cache", GRABIT_ARTIFACTORY_TOKEN_ENV_VAR) + } + + err := requests. + URL(l.getCacheFullURL()). + Method(http.MethodPut). + Header("Authorization", fmt.Sprintf("Bearer %s", token)). + BodyFile(filePath). + Fetch(context.Background()) + if err != nil { + return fmt.Errorf("failed to upload to cache: %v", err) + } + return nil +} - return &Resource{Urls: urls, Integrity: integrity, Tags: tags, Filename: filename, CacheUri: cacheURL}, nil +func (l *Resource) Delete() error { + token := getArtifactoryToken() + if token == "" { + log.Warn().Msgf("%s environment variable is not set and is needed to delete the file from the cache", GRABIT_ARTIFACTORY_TOKEN_ENV_VAR) + return nil + } + url := l.getCacheFullURL() + err := requests. + URL(url). + Method(http.MethodDelete). + Header("Authorization", fmt.Sprintf("Bearer %s", token)). + Fetch(context.Background()) + log.Warn().Msgf("Error deleting file from cache (%s): %v", url, err) + return nil } // getUrl downloads the given resource and returns the path to it. @@ -91,35 +147,38 @@ func GetUrltoTempFile(u string, bearer string, ctx context.Context) (string, err } func (l *Resource) Download(dir string, mode os.FileMode, ctx context.Context) error { - // Check if a cache URL exists to use Artifactory first - if l.CacheUri != "" { - token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") + algo, err := getAlgoFromIntegrity(l.Integrity) + if err != nil { + return err + } + // Check if a cache URL exists to use Artifactory first. + if l.ArtifactoryCacheURL != "" { + token := getArtifactoryToken() if token != "" { - artifactoryURL := fmt.Sprintf("%s/%s", l.CacheUri, l.Integrity) + artifactoryURL := fmt.Sprintf("%s/%s", l.ArtifactoryCacheURL, l.Integrity) localName := l.Filename if localName == "" { localName = path.Base(l.Urls[0]) } resPath := filepath.Join(dir, localName) - // Use getUrl directly with bearer token for Artifactory tmpPath, err := getUrl(artifactoryURL, resPath, token, ctx) if err == nil { if mode != NoFileMode { err = os.Chmod(tmpPath, mode.Perm()) + if err != nil { + return fmt.Errorf("error changing target file permission: '%v'", err) + } } - if err == nil { - return nil // Success + err = checkIntegrityFromFile(resPath, algo, l.Integrity, artifactoryURL) + if err != nil { + return fmt.Errorf("cache file at '%s' with incorrect integrity: '%v'", artifactoryURL, err) } } - fmt.Printf("Failed to download from Artifactory, falling back to original URL: %v\n", err) + log.Warn().Msgf("Failed to download from Artifactory cache, falling back to original URL: %v\n", err) } } ok := false - algo, err := getAlgoFromIntegrity(l.Integrity) - if err != nil { - return err - } var downloadError error = nil for _, u := range l.Urls { @@ -169,15 +228,10 @@ func (l *Resource) Download(dir string, mode os.FileMode, ctx context.Context) e if mode != NoFileMode { err = os.Chmod(resPath, mode.Perm()) if err != nil { - return err + return fmt.Errorf("error changing target file permission: '%v'", err) } } ok = true - if l.CacheUri != "" && os.Getenv("NO_CACHE_UPLOAD") != "1" { - if uploadErr := uploadToArtifactory(resPath, l.CacheUri, l.Integrity); uploadErr != nil { - fmt.Printf("Warning: Failed to upload to cache: %v\n", uploadErr) - } - } break } if !ok { diff --git a/internal/resource_test.go b/internal/resource_test.go index 86b2b06..f05c73d 100644 --- a/internal/resource_test.go +++ b/internal/resource_test.go @@ -13,6 +13,7 @@ import ( "github.com/cisco-open/grabit/test" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestNewResourceFromUrl(t *testing.T) { @@ -34,7 +35,7 @@ func TestNewResourceFromUrl(t *testing.T) { { urls: []string{fmt.Sprintf("http://localhost:%d/test.html", port)}, valid: true, - res: Resource{Urls: []string{fmt.Sprintf("http://localhost:%d/test.html", port)}, Integrity: fmt.Sprintf("%s-vvV+x/U6bUC+tkCngKY5yDvCmsipgW8fxsXG3Nk8RyE=", algo), Tags: []string{}, Filename: "", CacheUri: ""}, + res: Resource{Urls: []string{fmt.Sprintf("http://localhost:%d/test.html", port)}, Integrity: fmt.Sprintf("%s-vvV+x/U6bUC+tkCngKY5yDvCmsipgW8fxsXG3Nk8RyE=", algo), Tags: []string{}, Filename: "", ArtifactoryCacheURL: ""}, }, { urls: []string{"invalid url"}, @@ -83,3 +84,87 @@ func TestResourceDownloadWithInValidFileAlreadyPresent(t *testing.T) { assert.Contains(t, err.Error(), "integrity mismatch") assert.Contains(t, err.Error(), "existing file") } + +func TestUseResourceWithCache(t *testing.T) { + content := `abcdef` + token := "test-token" + port, server := test.TestHttpHandlerWithServer(content, t) + fileName := "test.txt" + sourceURL := fmt.Sprintf("http://localhost:%d", port) + + artServer, artPort := test.NewRecorderHttpServer(func(w http.ResponseWriter, r *http.Request) { + if r.Method == "GET" { + _, err := w.Write([]byte(content)) + if err != nil { + t.Fatal(err) + } + } + }, t) + baseCacheURL := fmt.Sprintf("http://localhost:%d/", artPort) + + // Create resource, download it, upload it to cache. + t.Setenv(GRABIT_ARTIFACTORY_TOKEN_ENV_VAR, token) + resource, err := NewResourceFromUrl([]string{sourceURL}, "sha256", []string{}, fileName, baseCacheURL) + require.Nil(t, err) + server.Close() // Close origin server: file will be served from cache. + outputDir := test.TmpDir(t) + + // Download resource from cache. + err = resource.Download(outputDir, 0644, context.Background()) + require.Nil(t, err) + for _, file := range []string{"test.txt"} { + test.AssertFileContains(t, fmt.Sprintf("%s/%s", outputDir, file), content) + } + assert.Equal(t, 2, len(*artServer.Requests)) + assert.Equal(t, "PUT", (*artServer.Requests)[0].Method) + assert.Equal(t, []byte(content), (*artServer.Requests)[0].Body) + assert.Equal(t, []string([]string{fmt.Sprintf("Bearer %s", token)}), (*artServer.Requests)[0].Headers["Authorization"]) + assert.Equal(t, "GET", (*artServer.Requests)[1].Method) + + // Delete resource, deleting it from cache. + err = resource.Delete() + require.Nil(t, err) + assert.Equal(t, 3, len(*artServer.Requests)) + assert.Equal(t, "DELETE", (*artServer.Requests)[2].Method) + assert.Equal(t, []string([]string{fmt.Sprintf("Bearer %s", token)}), (*artServer.Requests)[2].Headers["Authorization"]) +} + +func TestResourceWithCacheCorruptedCache(t *testing.T) { + content := `abcdef` + token := "test-token" + port, server := test.TestHttpHandlerWithServer(content, t) + fileName := "test.txt" + sourceURL := fmt.Sprintf("http://localhost:%d", port) + + _, artPort := test.NewRecorderHttpServer(func(w http.ResponseWriter, r *http.Request) { + if r.Method == "GET" { + _, err := w.Write([]byte("invalid-content")) + if err != nil { + t.Fatal(err) + } + } + }, t) + baseCacheURL := fmt.Sprintf("http://localhost:%d/", artPort) + + // Create resource, download it, upload it to cache. + t.Setenv(GRABIT_ARTIFACTORY_TOKEN_ENV_VAR, token) + resource, err := NewResourceFromUrl([]string{sourceURL}, "sha256", []string{}, fileName, baseCacheURL) + require.Nil(t, err) + server.Close() // Close origin server: file will be served from cache. + outputDir := test.TmpDir(t) + + // Download resource from cache. + err = resource.Download(outputDir, 0644, context.Background()) + assert.NotNil(t, err) + assert.Contains(t, err.Error(), "cache file at") + assert.Contains(t, err.Error(), "with incorrect integrity") +} + +func TestResourceWithCacheNoToken(t *testing.T) { + t.Setenv(GRABIT_ARTIFACTORY_TOKEN_ENV_VAR, "") + fileName := "test.txt" + port := 33 + sourceURL := fmt.Sprintf("http://localhost:%d", port) + _, err := NewResourceFromUrl([]string{sourceURL}, "sha256", []string{}, fileName, "http://localhost:8080/") + assert.NotNil(t, err) +} diff --git a/test/utils.go b/test/utils.go index 3653641..5fcb1de 100644 --- a/test/utils.go +++ b/test/utils.go @@ -4,9 +4,11 @@ package test import ( + "bytes" "crypto/sha256" "encoding/base64" "fmt" + "io" "log" "net" "net/http" @@ -58,9 +60,51 @@ func HttpHandler(handler http.HandlerFunc) (int, *httptest.Server) { return l.Addr().(*net.TCPAddr).Port, s } +type RecordedRequest struct { + Method string + Url string + Body []byte + Headers map[string][]string +} + +func NewRecordedRequest(r *http.Request) *RecordedRequest { + var bodyBytes []byte + if r.Body != nil { + bodyBytes, _ = io.ReadAll(r.Body) + } + r.Body = io.NopCloser(bytes.NewReader(bodyBytes)) + return &RecordedRequest{ + Method: r.Method, + Url: r.URL.String(), + Body: bodyBytes, + Headers: r.Header} +} + +type RecorderHttpServer struct { + *httptest.Server + Requests *[]RecordedRequest +} + +func NewRecorderHttpServer(handler http.HandlerFunc, t *testing.T) (*RecorderHttpServer, int) { + requests := make([]RecordedRequest, 0) + + outerHandler := func(w http.ResponseWriter, r *http.Request) { + requests = append(requests, *NewRecordedRequest(r)) + handler(w, r) + } + port, server := HttpHandler(outerHandler) + t.Cleanup(func() { server.Close() }) + return &RecorderHttpServer{Server: server, Requests: &requests}, port +} + // TestHttpHandler creates a new HTTP server and returns the port and serves // the given content. Its lifetime is tied to the given testing.T object. func TestHttpHandler(content string, t *testing.T) int { + port, _ := TestHttpHandlerWithServer(content, t) + return port +} + +func TestHttpHandlerWithServer(content string, t *testing.T) (int, *httptest.Server) { handler := func(w http.ResponseWriter, r *http.Request) { _, err := w.Write([]byte(content)) if err != nil { @@ -69,7 +113,7 @@ func TestHttpHandler(content string, t *testing.T) int { } port, server := HttpHandler(handler) t.Cleanup(func() { server.Close() }) - return port + return port, server } // AssertFileContains asserts that the file at the given path exists and