cloudflare/cloudflared
Publicmirrored fromhttps://github.com/cloudflare/cloudflaredAvailable
awsuploader/directory_upload_manager_test.go
137lines · modecode
| 1 | package awsuploader |
| 2 | |
| 3 | import ( |
| 4 | "errors" |
| 5 | "io/ioutil" |
| 6 | "math/rand" |
| 7 | "os" |
| 8 | "path/filepath" |
| 9 | "testing" |
| 10 | "time" |
| 11 | |
| 12 | "github.com/cloudflare/cloudflared/logger" |
| 13 | ) |
| 14 | |
| 15 | type MockUploader struct { |
| 16 | shouldFail bool |
| 17 | } |
| 18 | |
| 19 | func (m *MockUploader) Upload(filepath string) error { |
| 20 | if m.shouldFail { |
| 21 | return errors.New("upload set to fail") |
| 22 | } |
| 23 | return nil |
| 24 | } |
| 25 | |
| 26 | func NewMockUploader(shouldFail bool) Uploader { |
| 27 | return &MockUploader{shouldFail: shouldFail} |
| 28 | } |
| 29 | |
| 30 | func getDirectoryPath(t *testing.T) string { |
| 31 | dir, err := os.Getwd() |
| 32 | if err != nil { |
| 33 | t.Fatal("couldn't create the test directory!", err) |
| 34 | } |
| 35 | return filepath.Join(dir, "uploads") |
| 36 | } |
| 37 | |
| 38 | func setupTestDirectory(t *testing.T) string { |
| 39 | path := getDirectoryPath(t) |
| 40 | os.RemoveAll(path) |
| 41 | time.Sleep(100 * time.Millisecond) //short way to wait for the OS to delete the folder |
| 42 | err := os.MkdirAll(path, os.ModePerm) |
| 43 | if err != nil { |
| 44 | t.Fatal("couldn't create the test directory!", err) |
| 45 | } |
| 46 | return path |
| 47 | } |
| 48 | |
| 49 | func createUploadManager(t *testing.T, shouldFailUpload bool) *DirectoryUploadManager { |
| 50 | rootDirectory := setupTestDirectory(t) |
| 51 | uploader := NewMockUploader(shouldFailUpload) |
| 52 | logger := logger.NewOutputWriter(logger.NewMockWriteManager()) |
| 53 | shutdownC := make(chan struct{}) |
| 54 | return NewDirectoryUploadManager(logger, uploader, rootDirectory, 1*time.Second, shutdownC) |
| 55 | } |
| 56 | |
| 57 | func createFile(t *testing.T, fileName string) (*os.File, string) { |
| 58 | path := filepath.Join(getDirectoryPath(t), fileName) |
| 59 | f, err := os.Create(path) |
| 60 | if err != nil { |
| 61 | t.Fatal("upload to create file for sweep test", err) |
| 62 | } |
| 63 | return f, path |
| 64 | } |
| 65 | |
| 66 | func TestUploadSuccess(t *testing.T) { |
| 67 | manager := createUploadManager(t, false) |
| 68 | path := filepath.Join(getDirectoryPath(t), "test_file") |
| 69 | if err := manager.Upload(path); err != nil { |
| 70 | t.Fatal("the upload request method failed", err) |
| 71 | } |
| 72 | } |
| 73 | |
| 74 | func TestUploadFailure(t *testing.T) { |
| 75 | manager := createUploadManager(t, true) |
| 76 | path := filepath.Join(getDirectoryPath(t), "test_file") |
| 77 | if err := manager.Upload(path); err == nil { |
| 78 | t.Fatal("the upload request method should have failed and didn't", err) |
| 79 | } |
| 80 | } |
| 81 | |
| 82 | func TestSweepSuccess(t *testing.T) { |
| 83 | manager := createUploadManager(t, false) |
| 84 | f, path := createFile(t, "test_file") |
| 85 | defer f.Close() |
| 86 | |
| 87 | manager.Start() |
| 88 | time.Sleep(2 * time.Second) |
| 89 | if _, err := os.Stat(path); os.IsExist(err) { |
| 90 | //the file should have been deleted |
| 91 | t.Fatal("the manager failed to delete the file", err) |
| 92 | } |
| 93 | } |
| 94 | |
| 95 | func TestSweepFailure(t *testing.T) { |
| 96 | manager := createUploadManager(t, true) |
| 97 | f, path := createFile(t, "test_file") |
| 98 | defer f.Close() |
| 99 | |
| 100 | manager.Start() |
| 101 | time.Sleep(2 * time.Second) |
| 102 | _, serr := f.Stat() |
| 103 | if serr != nil { |
| 104 | //the file should still exist |
| 105 | os.Remove(path) |
| 106 | t.Fatal("the manager failed to delete the file", serr) |
| 107 | } |
| 108 | } |
| 109 | |
| 110 | func TestHighLoad(t *testing.T) { |
| 111 | manager := createUploadManager(t, false) |
| 112 | for i := 0; i < 30; i++ { |
| 113 | f, _ := createFile(t, randomString(6)) |
| 114 | defer f.Close() |
| 115 | } |
| 116 | manager.Start() |
| 117 | time.Sleep(4 * time.Second) |
| 118 | |
| 119 | directory := getDirectoryPath(t) |
| 120 | files, err := ioutil.ReadDir(directory) |
| 121 | if err != nil || len(files) > 0 { |
| 122 | t.Fatalf("the manager failed to upload all the files: %s files left: %d", err, len(files)) |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | // LowerCase [a-z] |
| 127 | const randSet = "abcdefghijklmnopqrstuvwxyz" |
| 128 | |
| 129 | // String returns a string of length 'n' from a set of letters 'lset' |
| 130 | func randomString(n int) string { |
| 131 | b := make([]byte, n) |
| 132 | lsetLen := len(randSet) |
| 133 | for i := range b { |
| 134 | b[i] = randSet[rand.Intn(lsetLen)] |
| 135 | } |
| 136 | return string(b) |
| 137 | } |
| 138 | |