Skip to content

Commit fcdd18a

Browse files
JuneezeeJuneezee
authored
test: use T.TempDir to create temporary test directory (#647)
1 parent dfb1c18 commit fcdd18a

5 files changed

+27
-75
lines changed

‎encrypt_util.go

+3
Original file line numberDiff line numberDiff line change
@@ -176,10 +176,13 @@ func encryptFile(
176176
if err != nil {
177177
return nil, "", err
178178
}
179+
defer tmpOutputFile.Close()
179180
infile, err := os.OpenFile(filename, os.O_CREATE|os.O_RDONLY, readWriteFileMode)
180181
if err != nil {
181182
return nil, "", err
182183
}
184+
defer infile.Close()
185+
183186
meta, err := encryptStream(sfe, infile, tmpOutputFile, chunkSize)
184187
if err != nil {
185188
return nil, "", err

‎encrypt_util_test.go

+6-23
Original file line numberDiff line numberDiff line change
@@ -91,11 +91,7 @@ func TestEncryptDecryptFilePadding(t *testing.T) {
9191

9292
for _, test := range testcases {
9393
t.Run(fmt.Sprintf("%v_%v", test.numberOfBytesInEachRow, test.numberOfLines), func(t *testing.T) {
94-
tmpDir, err := os.MkdirTemp("", "data")
95-
if err != nil {
96-
t.Error(err)
97-
}
98-
tmpDir, err = generateKLinesOfNByteRows(test.numberOfLines, test.numberOfBytesInEachRow, tmpDir)
94+
tmpDir, err := generateKLinesOfNByteRows(test.numberOfLines, test.numberOfBytesInEachRow, t.TempDir())
9995
if err != nil {
10096
t.Error(err)
10197
}
@@ -114,11 +110,7 @@ func TestEncryptDecryptLargeFile(t *testing.T) {
114110

115111
numberOfFiles := 1
116112
numberOfLines := 10000
117-
tmpDir, err := os.MkdirTemp("", "data")
118-
if err != nil {
119-
t.Error(err)
120-
}
121-
tmpDir, err = generateKLinesOfNFiles(numberOfLines, numberOfFiles, false, tmpDir)
113+
tmpDir, err := generateKLinesOfNFiles(numberOfLines, numberOfFiles, false, t.TempDir())
122114
if err != nil {
123115
t.Error(err)
124116
}
@@ -127,7 +119,6 @@ func TestEncryptDecryptLargeFile(t *testing.T) {
127119
}
128120

129121
func encryptDecryptFile(t *testing.T, encMat snowflakeFileEncryption, expected int, tmpDir string) {
130-
defer os.RemoveAll(tmpDir)
131122
files, err := filepath.Glob(filepath.Join(tmpDir, "file*"))
132123
if err != nil {
133124
t.Error(err)
@@ -150,6 +141,8 @@ func encryptDecryptFile(t *testing.T, encMat snowflakeFileEncryption, expected i
150141
if err != nil {
151142
t.Error(err)
152143
}
144+
defer fd.Close()
145+
153146
scanner := bufio.NewScanner(fd)
154147
for scanner.Scan() {
155148
cnt++
@@ -163,12 +156,6 @@ func encryptDecryptFile(t *testing.T, encMat snowflakeFileEncryption, expected i
163156
}
164157

165158
func generateKLinesOfNByteRows(numLines int, numBytes int, tmpDir string) (string, error) {
166-
if tmpDir == "" {
167-
_, err := os.MkdirTemp(tmpDir, "data")
168-
if err != nil {
169-
return "", err
170-
}
171-
}
172159
fname := path.Join(tmpDir, "file"+strconv.FormatInt(int64(numLines*numBytes), 10))
173160
f, err := os.Create(fname)
174161
if err != nil {
@@ -185,12 +172,6 @@ func generateKLinesOfNByteRows(numLines int, numBytes int, tmpDir string) (strin
185172
}
186173

187174
func generateKLinesOfNFiles(k int, n int, compress bool, tmpDir string) (string, error) {
188-
if tmpDir == "" {
189-
_, err := os.MkdirTemp(tmpDir, "data")
190-
if err != nil {
191-
return "", err
192-
}
193-
}
194175
for i := 0; i < n; i++ {
195176
fname := path.Join(tmpDir, "file"+strconv.FormatInt(int64(i), 10))
196177
f, err := os.Create(fname)
@@ -248,6 +229,8 @@ func generateKLinesOfNFiles(k int, n int, compress bool, tmpDir string) (string,
248229
return "", err
249230
}
250231
w.Close()
232+
fOut.Close()
233+
fIn.Close()
251234
}
252235
}
253236
}

‎put_get_test.go

+10-27
Original file line numberDiff line numberDiff line change
@@ -34,19 +34,17 @@ func TestPutError(t *testing.T) {
3434
if isWindows {
3535
t.Skip("permission model is different")
3636
}
37-
tmpDir, err := os.MkdirTemp("", "putfiledir")
38-
if err != nil {
39-
t.Error(err)
40-
}
41-
defer os.RemoveAll(tmpDir)
37+
tmpDir := t.TempDir()
4238
file1 := filepath.Join(tmpDir, "file1")
4339
remoteLocation := filepath.Join(tmpDir, "remote_loc")
4440
f, err := os.Create(file1)
4541
if err != nil {
4642
t.Error(err)
4743
}
44+
defer f.Close()
4845
f.WriteString("test1")
4946
os.Chmod(file1, 0000)
47+
defer os.Chmod(file1, 0644)
5048

5149
data := &execResponseData{
5250
Command: string(uploadCommand),
@@ -253,11 +251,7 @@ func TestPutWithAutoCompressFalse(t *testing.T) {
253251
if runningOnGithubAction() && !runningOnAWS() {
254252
t.Skip("skipping non aws environment")
255253
}
256-
tmpDir, err := os.MkdirTemp("", "put")
257-
if err != nil {
258-
t.Error(err)
259-
}
260-
defer os.RemoveAll(tmpDir)
254+
tmpDir := t.TempDir()
261255
testData := filepath.Join(tmpDir, "data.txt")
262256
f, err := os.Create(testData)
263257
if err != nil {
@@ -294,11 +288,7 @@ func TestPutWithAutoCompressFalse(t *testing.T) {
294288
}
295289

296290
func TestPutOverwrite(t *testing.T) {
297-
tmpDir, err := os.MkdirTemp("", "data")
298-
if err != nil {
299-
t.Error(err)
300-
}
301-
defer os.RemoveAll(tmpDir)
291+
tmpDir := t.TempDir()
302292
testData := filepath.Join(tmpDir, "data.txt")
303293
f, err := os.Create(testData)
304294
if err != nil {
@@ -388,11 +378,7 @@ func TestPutGetStream(t *testing.T) {
388378
}
389379

390380
func testPutGet(t *testing.T, isStream bool) {
391-
tmpDir, err := os.MkdirTemp("", "put_get")
392-
if err != nil {
393-
t.Error(err)
394-
}
395-
defer os.RemoveAll(tmpDir)
381+
tmpDir := t.TempDir()
396382
fname := filepath.Join(tmpDir, "test_put_get.txt.gz")
397383
originalContents := "123,test1\n456,test2\n"
398384
tableName := randomString(5)
@@ -401,23 +387,19 @@ func testPutGet(t *testing.T, isStream bool) {
401387
gzw := gzip.NewWriter(&b)
402388
gzw.Write([]byte(originalContents))
403389
gzw.Close()
404-
if err = os.WriteFile(fname, b.Bytes(), readWriteFileMode); err != nil {
390+
if err := os.WriteFile(fname, b.Bytes(), readWriteFileMode); err != nil {
405391
t.Fatal("could not write to gzip file")
406392
}
407393

408394
runDBTest(t, func(dbt *DBTest) {
409395
dbt.mustExec("create or replace table " + tableName +
410396
" (a int, b string)")
397+
defer dbt.mustExec("drop table " + tableName)
411398
fileStream, err := os.Open(fname)
412399
if err != nil {
413400
t.Error(err)
414401
}
415-
defer func() {
416-
defer dbt.mustExec("drop table " + tableName)
417-
if fileStream != nil {
418-
fileStream.Close()
419-
}
420-
}()
402+
defer fileStream.Close()
421403

422404
var sqlText string
423405
var rows *RowsExtended
@@ -489,6 +471,7 @@ func testPutGet(t *testing.T, isStream bool) {
489471
if err != nil {
490472
t.Error(err)
491473
}
474+
defer gz.Close()
492475
var contents string
493476
for {
494477
c := make([]byte, defaultChunkBufferSize)

‎put_get_user_stage_test.go

+4-9
Original file line numberDiff line numberDiff line change
@@ -14,29 +14,24 @@ func TestPutGetFileSmallDataViaUserStage(t *testing.T) {
1414
if os.Getenv("AWS_ACCESS_KEY_ID") == "" {
1515
t.Skip("this test requires to change the internal parameter")
1616
}
17-
putGetUserStage(t, "", 5, 1, false)
17+
putGetUserStage(t, 5, 1, false)
1818
}
1919

2020
func TestPutGetStreamSmallDataViaUserStage(t *testing.T) {
2121
if os.Getenv("AWS_ACCESS_KEY_ID") == "" {
2222
t.Skip("this test requires to change the internal parameter")
2323
}
24-
putGetUserStage(t, "", 1, 1, true)
24+
putGetUserStage(t, 1, 1, true)
2525
}
2626

27-
func putGetUserStage(t *testing.T, tmpDir string, numberOfFiles int, numberOfLines int, isStream bool) {
27+
func putGetUserStage(t *testing.T, numberOfFiles int, numberOfLines int, isStream bool) {
2828
if os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
2929
t.Fatal("no aws secret access key found")
3030
}
31-
tmpDir, err := os.MkdirTemp(tmpDir, "data")
31+
tmpDir, err := generateKLinesOfNFiles(numberOfLines, numberOfFiles, false, t.TempDir())
3232
if err != nil {
3333
t.Error(err)
3434
}
35-
tmpDir, err = generateKLinesOfNFiles(numberOfLines, numberOfFiles, false, tmpDir)
36-
if err != nil {
37-
t.Error(err)
38-
}
39-
defer os.RemoveAll(tmpDir)
4035
var files string
4136
if isStream {
4237
list, err := os.ReadDir(tmpDir)

‎put_get_with_aws_test.go

+4-16
Original file line numberDiff line numberDiff line change
@@ -91,11 +91,7 @@ func TestPutWithInvalidToken(t *testing.T) {
9191
if !runningOnAWS() {
9292
t.Skip("skipping non aws environment")
9393
}
94-
tmpDir, err := os.MkdirTemp("", "aws_put")
95-
if err != nil {
96-
t.Error(err)
97-
}
98-
defer os.RemoveAll(tmpDir)
94+
tmpDir := t.TempDir()
9995
fname := filepath.Join(tmpDir, "test_put_get_with_aws.txt.gz")
10096
originalContents := "123,test1\n456,test2\n"
10197

@@ -189,11 +185,7 @@ func TestPretendToPutButList(t *testing.T) {
189185
if runningOnGithubAction() && !runningOnAWS() {
190186
t.Skip("skipping non aws environment")
191187
}
192-
tmpDir, err := os.MkdirTemp("", "aws_put")
193-
if err != nil {
194-
t.Error(err)
195-
}
196-
defer os.RemoveAll(tmpDir)
188+
tmpDir := t.TempDir()
197189
fname := filepath.Join(tmpDir, "test_put_get_with_aws.txt.gz")
198190
originalContents := "123,test1\n456,test2\n"
199191

@@ -244,11 +236,7 @@ func TestPutGetAWSStage(t *testing.T) {
244236
t.Skip("skipping non aws environment")
245237
}
246238

247-
tmpDir, err := os.MkdirTemp("", "put_get")
248-
if err != nil {
249-
t.Error(err)
250-
}
251-
defer os.RemoveAll(tmpDir)
239+
tmpDir := t.TempDir()
252240
name := "test_put_get.txt.gz"
253241
fname := filepath.Join(tmpDir, name)
254242
originalContents := "123,test1\n456,test2\n"
@@ -258,7 +246,7 @@ func TestPutGetAWSStage(t *testing.T) {
258246
gzw := gzip.NewWriter(&b)
259247
gzw.Write([]byte(originalContents))
260248
gzw.Close()
261-
if err = os.WriteFile(fname, b.Bytes(), readWriteFileMode); err != nil {
249+
if err := os.WriteFile(fname, b.Bytes(), readWriteFileMode); err != nil {
262250
t.Fatal("could not write to gzip file")
263251
}
264252

0 commit comments

Comments
 (0)