github.com/vipernet-xyz/tm@v0.34.24/libs/tempfile/tempfile_test.go (about)

     1  package tempfile
     2  
     3  // Need access to internal variables, so can't use _test package
     4  
     5  import (
     6  	"bytes"
     7  	"fmt"
     8  	"os"
     9  	testing "testing"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	tmrand "github.com/vipernet-xyz/tm/libs/rand"
    14  )
    15  
    16  func TestWriteFileAtomic(t *testing.T) {
    17  	var (
    18  		data             = []byte(tmrand.Str(tmrand.Intn(2048)))
    19  		old              = tmrand.Bytes(tmrand.Intn(2048))
    20  		perm os.FileMode = 0o600
    21  	)
    22  
    23  	f, err := os.CreateTemp("/tmp", "write-atomic-test-")
    24  	if err != nil {
    25  		t.Fatal(err)
    26  	}
    27  	defer os.Remove(f.Name())
    28  
    29  	if err = os.WriteFile(f.Name(), old, 0o600); err != nil {
    30  		t.Fatal(err)
    31  	}
    32  
    33  	if err = WriteFileAtomic(f.Name(), data, perm); err != nil {
    34  		t.Fatal(err)
    35  	}
    36  
    37  	rData, err := os.ReadFile(f.Name())
    38  	if err != nil {
    39  		t.Fatal(err)
    40  	}
    41  
    42  	if !bytes.Equal(data, rData) {
    43  		t.Fatalf("data mismatch: %v != %v", data, rData)
    44  	}
    45  
    46  	stat, err := os.Stat(f.Name())
    47  	if err != nil {
    48  		t.Fatal(err)
    49  	}
    50  
    51  	if have, want := stat.Mode().Perm(), perm; have != want {
    52  		t.Errorf("have %v, want %v", have, want)
    53  	}
    54  }
    55  
    56  // This tests atomic write file when there is a single duplicate file.
    57  // Expected behavior is for a new file to be created, and the original write file to be unaltered.
    58  func TestWriteFileAtomicDuplicateFile(t *testing.T) {
    59  	var (
    60  		defaultSeed    uint64 = 1
    61  		testString            = "This is a glorious test string"
    62  		expectedString        = "Did the test file's string appear here?"
    63  
    64  		fileToWrite = "/tmp/TestWriteFileAtomicDuplicateFile-test.txt"
    65  	)
    66  	// Create a file at the seed, and reset the seed.
    67  	atomicWriteFileRand = defaultSeed
    68  	firstFileRand := randWriteFileSuffix()
    69  	atomicWriteFileRand = defaultSeed
    70  	fname := "/tmp/" + atomicWriteFilePrefix + firstFileRand
    71  	f, err := os.OpenFile(fname, atomicWriteFileFlag, 0o777)
    72  	defer os.Remove(fname)
    73  	// Defer here, in case there is a panic in WriteFileAtomic.
    74  	defer os.Remove(fileToWrite)
    75  
    76  	require.NoError(t, err)
    77  	_, err = f.WriteString(testString)
    78  	require.NoError(t, err)
    79  	err = WriteFileAtomic(fileToWrite, []byte(expectedString), 0o777)
    80  	require.NoError(t, err)
    81  	// Check that the first atomic file was untouched
    82  	firstAtomicFileBytes, err := os.ReadFile(fname)
    83  	require.NoError(t, err, "Error reading first atomic file")
    84  	require.Equal(t, []byte(testString), firstAtomicFileBytes, "First atomic file was overwritten")
    85  	// Check that the resultant file is correct
    86  	resultantFileBytes, err := os.ReadFile(fileToWrite)
    87  	require.NoError(t, err, "Error reading resultant file")
    88  	require.Equal(t, []byte(expectedString), resultantFileBytes, "Written file had incorrect bytes")
    89  
    90  	// Check that the intermediate write file was deleted
    91  	// Get the second write files' randomness
    92  	atomicWriteFileRand = defaultSeed
    93  	_ = randWriteFileSuffix()
    94  	secondFileRand := randWriteFileSuffix()
    95  	_, err = os.Stat("/tmp/" + atomicWriteFilePrefix + secondFileRand)
    96  	require.True(t, os.IsNotExist(err), "Intermittent atomic write file not deleted")
    97  }
    98  
    99  // This tests atomic write file when there are many duplicate files.
   100  // Expected behavior is for a new file to be created under a completely new seed,
   101  // and the original write files to be unaltered.
   102  func TestWriteFileAtomicManyDuplicates(t *testing.T) {
   103  	var (
   104  		defaultSeed    uint64 = 2
   105  		testString            = "This is a glorious test string, from file %d"
   106  		expectedString        = "Did any of the test file's string appear here?"
   107  
   108  		fileToWrite = "/tmp/TestWriteFileAtomicDuplicateFile-test.txt"
   109  	)
   110  	// Initialize all of the atomic write files
   111  	atomicWriteFileRand = defaultSeed
   112  	for i := 0; i < atomicWriteFileMaxNumConflicts+2; i++ {
   113  		fileRand := randWriteFileSuffix()
   114  		fname := "/tmp/" + atomicWriteFilePrefix + fileRand
   115  		f, err := os.OpenFile(fname, atomicWriteFileFlag, 0o777)
   116  		require.Nil(t, err)
   117  		_, err = f.WriteString(fmt.Sprintf(testString, i))
   118  		require.NoError(t, err)
   119  		defer os.Remove(fname)
   120  	}
   121  
   122  	atomicWriteFileRand = defaultSeed
   123  	// Defer here, in case there is a panic in WriteFileAtomic.
   124  	defer os.Remove(fileToWrite)
   125  
   126  	err := WriteFileAtomic(fileToWrite, []byte(expectedString), 0o777)
   127  	require.NoError(t, err)
   128  	// Check that all intermittent atomic file were untouched
   129  	atomicWriteFileRand = defaultSeed
   130  	for i := 0; i < atomicWriteFileMaxNumConflicts+2; i++ {
   131  		fileRand := randWriteFileSuffix()
   132  		fname := "/tmp/" + atomicWriteFilePrefix + fileRand
   133  		firstAtomicFileBytes, err := os.ReadFile(fname)
   134  		require.Nil(t, err, "Error reading first atomic file")
   135  		require.Equal(t, []byte(fmt.Sprintf(testString, i)), firstAtomicFileBytes,
   136  			"atomic write file %d was overwritten", i)
   137  	}
   138  
   139  	// Check that the resultant file is correct
   140  	resultantFileBytes, err := os.ReadFile(fileToWrite)
   141  	require.Nil(t, err, "Error reading resultant file")
   142  	require.Equal(t, []byte(expectedString), resultantFileBytes, "Written file had incorrect bytes")
   143  }