github.com/attic-labs/noms@v0.0.0-20210827224422-e5fa29d95e8b/samples/go/csv/write_test.go (about)

     1  // Copyright 2016 Attic Labs, Inc. All rights reserved.
     2  // Licensed under the Apache License, version 2.0:
     3  // http://www.apache.org/licenses/LICENSE-2.0
     4  
     5  package csv
     6  
     7  import (
     8  	"bytes"
     9  	"encoding/csv"
    10  	"fmt"
    11  	"io"
    12  	"io/ioutil"
    13  	"os"
    14  	"strings"
    15  	"testing"
    16  
    17  	"github.com/attic-labs/noms/go/chunks"
    18  	"github.com/attic-labs/noms/go/d"
    19  	"github.com/attic-labs/noms/go/datas"
    20  	"github.com/attic-labs/noms/go/types"
    21  	"github.com/attic-labs/noms/go/util/clienttest"
    22  	"github.com/stretchr/testify/suite"
    23  )
    24  
    25  const (
    26  	TEST_ROW_STRUCT_NAME = "row"
    27  	TEST_ROW_FIELDS      = "anid,month,rainfall,year"
    28  	TEST_DATA_SIZE       = 200
    29  	TEST_YEAR            = 2012
    30  )
    31  
    32  func TestCSVWrite(t *testing.T) {
    33  	suite.Run(t, &csvWriteTestSuite{})
    34  }
    35  
    36  type csvWriteTestSuite struct {
    37  	clienttest.ClientTestSuite
    38  	fieldTypes    []*types.Type
    39  	rowStructDesc types.StructDesc
    40  	comma         rune
    41  	tmpFileName   string
    42  }
    43  
    44  func typesToKinds(ts []*types.Type) KindSlice {
    45  	kinds := make(KindSlice, len(ts))
    46  	for i, t := range ts {
    47  		kinds[i] = t.TargetKind()
    48  	}
    49  	return kinds
    50  }
    51  
    52  func (s *csvWriteTestSuite) SetupTest() {
    53  	input, err := ioutil.TempFile(s.TempDir, "")
    54  	d.Chk.NoError(err)
    55  	s.tmpFileName = input.Name()
    56  	defer input.Close()
    57  
    58  	fieldNames := strings.Split(TEST_ROW_FIELDS, ",")
    59  	s.fieldTypes = []*types.Type{types.StringType, types.NumberType, types.NumberType, types.NumberType}
    60  	fields := make([]types.StructField, len(fieldNames))
    61  	for i, name := range fieldNames {
    62  		fields[i] = types.StructField{
    63  			Name: name,
    64  			Type: s.fieldTypes[i],
    65  		}
    66  	}
    67  	rowStructType := types.MakeStructType(TEST_ROW_STRUCT_NAME, fields...)
    68  	s.rowStructDesc = rowStructType.Desc.(types.StructDesc)
    69  	s.comma, _ = StringToRune(",")
    70  	createCsvTestExpectationFile(input)
    71  }
    72  
    73  func (s *csvWriteTestSuite) TearDownTest() {
    74  	os.Remove(s.tmpFileName)
    75  }
    76  
    77  func createCsvTestExpectationFile(w io.Writer) {
    78  	_, err := io.WriteString(w, TEST_ROW_FIELDS)
    79  	d.Chk.NoError(err)
    80  	_, err = io.WriteString(w, "\n")
    81  	d.Chk.NoError(err)
    82  	for i := 0; i < TEST_DATA_SIZE; i++ {
    83  		_, err = io.WriteString(w, fmt.Sprintf("a - %3d,%d,%d,%d\n", i, i%12, i%32, TEST_YEAR+i%4))
    84  		d.Chk.NoError(err)
    85  	}
    86  }
    87  
    88  func startReadingCsvTestExpectationFile(s *csvWriteTestSuite) (cr *csv.Reader, headers []string) {
    89  	res, err := os.Open(s.tmpFileName)
    90  	d.PanicIfError(err)
    91  	cr = NewCSVReader(res, s.comma)
    92  	headers, err = cr.Read()
    93  	d.PanicIfError(err)
    94  	return
    95  }
    96  
    97  func createTestList(s *csvWriteTestSuite) types.List {
    98  	storage := &chunks.MemoryStorage{}
    99  	db := datas.NewDatabase(storage.NewView())
   100  	cr, headers := startReadingCsvTestExpectationFile(s)
   101  	l := ReadToList(cr, TEST_ROW_STRUCT_NAME, headers, typesToKinds(s.fieldTypes), db, LIMIT)
   102  	return l
   103  }
   104  
   105  func createTestMap(s *csvWriteTestSuite) types.Map {
   106  	storage := &chunks.MemoryStorage{}
   107  	db := datas.NewDatabase(storage.NewView())
   108  	cr, headers := startReadingCsvTestExpectationFile(s)
   109  	return ReadToMap(cr, TEST_ROW_STRUCT_NAME, headers, []string{"anid"}, typesToKinds(s.fieldTypes), db, LIMIT)
   110  }
   111  
   112  func createTestNestedMap(s *csvWriteTestSuite) types.Map {
   113  	storage := &chunks.MemoryStorage{}
   114  	db := datas.NewDatabase(storage.NewView())
   115  	cr, headers := startReadingCsvTestExpectationFile(s)
   116  	return ReadToMap(cr, TEST_ROW_STRUCT_NAME, headers, []string{"anid", "year"}, typesToKinds(s.fieldTypes), db, LIMIT)
   117  }
   118  
   119  func verifyOutput(s *csvWriteTestSuite, r io.Reader) {
   120  	res, err := os.Open(s.tmpFileName)
   121  	d.PanicIfError(err)
   122  	actual, err := ioutil.ReadAll(r)
   123  	d.Chk.NoError(err)
   124  	expected, err := ioutil.ReadAll(res)
   125  	d.Chk.NoError(err)
   126  	s.True(string(expected) == string(actual), "csv files are different")
   127  }
   128  
   129  func (s *csvWriteTestSuite) TestCSVWriteList() {
   130  	l := createTestList(s)
   131  	w := new(bytes.Buffer)
   132  	s.True(TEST_DATA_SIZE == l.Len(), "list length")
   133  	WriteList(l, s.rowStructDesc, s.comma, w)
   134  	verifyOutput(s, w)
   135  }
   136  
   137  func (s *csvWriteTestSuite) TestCSVWriteMap() {
   138  	m := createTestMap(s)
   139  	w := new(bytes.Buffer)
   140  	s.True(TEST_DATA_SIZE == m.Len(), "map length")
   141  	WriteMap(m, s.rowStructDesc, s.comma, w)
   142  	verifyOutput(s, w)
   143  }
   144  
   145  func (s *csvWriteTestSuite) TestCSVWriteNestedMap() {
   146  	m := createTestNestedMap(s)
   147  	w := new(bytes.Buffer)
   148  	s.True(TEST_DATA_SIZE == m.Len(), "nested map length")
   149  	WriteMap(m, s.rowStructDesc, s.comma, w)
   150  	verifyOutput(s, w)
   151  }