Fix error in memorystore

files/filestore_fs_test.go

@@ -7,20 +7,23 @@ import (
 )
 
 func TestFSFileStore(t *testing.T) {
+	newFunc := func() files.FileStore {
 		dir := t.TempDir()
 		s, err := files.NewFSFileStore(dir)
 		if err != nil {
 			t.Fatalf("Error creating store: %s", err)
 		}
+		return s
+	}
 
-	RunFilestoreTest(s, t)
+	RunFilestoreTest(newFunc, t)
 	persistentDir := t.TempDir()
-	newFunc := func() files.FileStore {
+	persistentFunc := func() files.FileStore {
 		s, err := files.NewFSFileStore(persistentDir)
 		if err != nil {
 			t.Fatalf("Error creating store: %s", err)
 		}
 		return s
 	}
-	RunPersistentFilestoreTest(newFunc, t)
+	RunPersistentFilestoreTest(persistentFunc, t)
 }

files/filestore_memory.go

@@ -57,11 +57,16 @@ func (s *MemoryFileStore) Get(id string) (*File, error) {
 		return nil, fmt.Errorf("no such item")
 	}
 
+	body := new(bytes.Buffer)
+	if _, err := body.Write(fd.Body.Bytes()); err != nil {
+		return nil, err
+	}
+
 	f := &File{
 		ID:        fd.ID,
 		MaxViews:  fd.MaxViews,
 		ExpiresOn: fd.ExpiresOn,
-		Body:      io.NopCloser(&fd.Body),
+		Body:      io.NopCloser(body),
 		FileSize:  fd.FileSize,
 	}
 

files/filestore_memory_test.go

@@ -7,7 +7,9 @@ import (
 )
 
 func TestMemoryFileStore(t *testing.T) {
-	s := files.NewMemoryFileStore()
+	newFunc := func() files.FileStore {
+		return files.NewMemoryFileStore()
+	}
 
-	RunFilestoreTest(s, t)
+	RunFilestoreTest(newFunc, t)
 }

files/filestore_test.go

@@ -12,8 +12,11 @@ import (
 	"github.com/google/uuid"
 )
 
-func RunFilestoreTest(s files.FileStore, t *testing.T) {
+var ignoreBody = cmp.FilterPath(func(p cmp.Path) bool { return p.String() == "Body" }, cmp.Ignore())
+
+func RunFilestoreTest(newStoreFunc func() files.FileStore, t *testing.T) {
 	t.Run("Basic", func(t *testing.T) {
+		s := newStoreFunc()
 		// Create
 		dataString := "TEST_LOL_OMG"
 		id := uuid.Must(uuid.NewRandom()).String()
@@ -58,7 +61,6 @@ func RunFilestoreTest(s files.FileStore, t *testing.T) {
 			FileSize:  int64(len(dataString)),
 		}
 
-		ignoreBody := cmp.FilterPath(func(p cmp.Path) bool { return p.String() == "Body" }, cmp.Ignore())
 		if diff := cmp.Diff(retrieved, expected, ignoreBody); diff != "" {
 			t.Errorf("File comparison failed: %s", diff)
 		}
@@ -88,9 +90,54 @@ func RunFilestoreTest(s files.FileStore, t *testing.T) {
 			t.Fatalf("List after delete has wrong length: %d", len(ids))
 		}
 	})
+	t.Run("MultipleGet", func(t *testing.T) {
+		s := newStoreFunc()
+
+		fileContents := "multiple get test !"
+		body := io.NopCloser(strings.NewReader(fileContents))
+		file := &files.File{
+			ID:               uuid.NewString(),
+			OriginalFilename: "multiple.txt",
+			MaxViews:         999,
+			ExpiresOn:        time.Now().Add(1 * time.Hour),
+			Body:             body,
+			FileSize:         int64(len(fileContents)),
+		}
+
+		if err := s.Store(file); err != nil {
+			t.Fatalf("Error storing file: %s", err)
+		}
+
+		first, err := s.Get(file.ID)
+		if err != nil {
+			t.Errorf("Error retrieving first file: %s", err)
+		}
+
+		firstBody := new(bytes.Buffer)
+		io.Copy(firstBody, first.Body)
+		first.Body.Close()
+
+		if diff := cmp.Diff(firstBody.String(), fileContents); diff != "" {
+			t.Fatalf("File contents mismatch: %s", diff)
+		}
+
+		second, err := s.Get(file.ID)
+		if err != nil {
+			t.Errorf("Error retrieving first file: %s", err)
+		}
+
+		secondBody := new(bytes.Buffer)
+		io.Copy(secondBody, second.Body)
+		first.Body.Close()
+
+		if diff := cmp.Diff(secondBody.String(), fileContents); diff != "" {
+			t.Fatalf("File contents mismatch: %s", diff)
+		}
+	})
 }
 
 func RunPersistentFilestoreTest(newStoreFunc func() files.FileStore, t *testing.T) {
+	t.Run("Basics", func(t *testing.T) {
 		s := newStoreFunc()
 
 		files := []struct {
@@ -133,7 +180,6 @@ func RunPersistentFilestoreTest(newStoreFunc func() files.FileStore, t *testing.
 				t.Fatalf("Unable to retrieve file: %s", err)
 			}
 
-		ignoreBody := cmp.FilterPath(func(p cmp.Path) bool { return p.String() == "Body" }, cmp.Ignore())
 			if !cmp.Equal(retrieved, f.File, ignoreBody) {
 				t.Errorf("Mismatch: %s", cmp.Diff(retrieved, f.File))
 			}
@@ -155,7 +201,6 @@ func RunPersistentFilestoreTest(newStoreFunc func() files.FileStore, t *testing.
 				t.Fatalf("Unable to retrieve file: %s", err)
 			}
 
-		ignoreBody := cmp.FilterPath(func(p cmp.Path) bool { return p.String() == "Body" }, cmp.Ignore())
 			if !cmp.Equal(retrieved, f.File, ignoreBody) {
 				t.Errorf("Mismatch: %s", cmp.Diff(retrieved, f.File))
 			}
@@ -168,4 +213,5 @@ func RunPersistentFilestoreTest(newStoreFunc func() files.FileStore, t *testing.
 				t.Fatalf("Data does not match. %s", cmp.Diff(buf.String(), f.ExpectedData))
 			}
 		}
+	})
 }