mirror of
https://github.com/crawlab-team/crawlab.git
synced 2026-01-21 17:21:09 +01:00
140 lines
3.5 KiB
Go
140 lines
3.5 KiB
Go
package export
|
|
|
|
import (
|
|
"encoding/csv"
|
|
"fmt"
|
|
"github.com/crawlab-team/crawlab/core/constants"
|
|
"github.com/crawlab-team/crawlab/db/mongo"
|
|
"github.com/stretchr/testify/require"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/bson/primitive"
|
|
"os"
|
|
"strconv"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestCsvService_Export(t *testing.T) {
|
|
// test data rows
|
|
var rows []interface{}
|
|
for i := 0; i < 10; i++ {
|
|
data := bson.M{
|
|
"no": i,
|
|
"string_field": "test",
|
|
"int_field": 1,
|
|
"float_field": 1.1,
|
|
"bool_field": true,
|
|
"time_field": time.Now(),
|
|
"object_id_field": primitive.NewObjectID(),
|
|
}
|
|
rows = append(rows, data)
|
|
}
|
|
|
|
// test mongo collection name
|
|
collectionName := "test_collection"
|
|
|
|
// test mongo collection
|
|
collection := mongo.GetMongoCol(collectionName)
|
|
|
|
// delete records of test mongo collection after test
|
|
t.Cleanup(func() {
|
|
_ = collection.Delete(bson.M{})
|
|
})
|
|
|
|
// save test data rows to mongo collection
|
|
_, err := collection.InsertMany(rows)
|
|
require.Nil(t, err)
|
|
|
|
// export service
|
|
csvSvc := NewCsvService()
|
|
|
|
// export
|
|
exportId, err := csvSvc.Export(collectionName, collectionName, nil)
|
|
require.Nil(t, err)
|
|
|
|
// get export
|
|
export, err := csvSvc.GetExport(exportId)
|
|
require.Nil(t, err)
|
|
require.NotNil(t, export)
|
|
require.Equal(t, exportId, export.GetId())
|
|
require.NotNil(t, export.GetDownloadPath())
|
|
|
|
// wait for export to finish with timeout of 5 seconds
|
|
timeout := time.After(5 * time.Second)
|
|
finished := false
|
|
for {
|
|
if finished {
|
|
break
|
|
}
|
|
select {
|
|
case <-timeout:
|
|
t.Fatal("export timeout")
|
|
default:
|
|
if export.GetStatus() == constants.TaskStatusFinished {
|
|
finished = true
|
|
continue
|
|
}
|
|
time.Sleep(100 * time.Millisecond)
|
|
}
|
|
}
|
|
|
|
// export file path
|
|
exportFilePath := export.GetDownloadPath()
|
|
require.FileExists(t, exportFilePath)
|
|
|
|
// csv file
|
|
csvFile, err := os.Open(exportFilePath)
|
|
require.Nil(t, err)
|
|
defer csvFile.Close()
|
|
|
|
// csv file reader
|
|
csvFileReader := csv.NewReader(csvFile)
|
|
|
|
// csv file rows
|
|
csvFileRows, err := csvFileReader.ReadAll()
|
|
require.Nil(t, err)
|
|
require.Equal(t, len(rows), len(csvFileRows)-1)
|
|
|
|
// csv file columns
|
|
csvFileColumns := csvFileRows[0]
|
|
|
|
// iterate csv file records and compare with test data rows
|
|
for i, row := range rows {
|
|
// csv file record
|
|
csvFileRecord := csvFileRows[i+1]
|
|
|
|
// iterate csv file columns and compare with test data row
|
|
for j, column := range csvFileColumns {
|
|
// csv file column value
|
|
csvFileColumnValue := csvFileRecord[j]
|
|
|
|
// compare csv file column value with test data row
|
|
switch column {
|
|
case "no":
|
|
// convert int to string
|
|
stringValue := fmt.Sprintf("%d", row.(bson.M)["no"].(int))
|
|
require.Equal(t, stringValue, csvFileColumnValue)
|
|
case "string_field":
|
|
require.Equal(t, row.(bson.M)["string_field"], csvFileColumnValue)
|
|
case "int_field":
|
|
// convert int to string
|
|
stringValue := fmt.Sprintf("%d", row.(bson.M)["int_field"])
|
|
require.Equal(t, stringValue, csvFileColumnValue)
|
|
case "float_field":
|
|
// convert string to float
|
|
floatValue, err := strconv.ParseFloat(csvFileColumnValue, 64)
|
|
require.Nil(t, err)
|
|
require.Equal(t, row.(bson.M)["float_field"].(float64), floatValue)
|
|
case "bool_field":
|
|
// convert bool to string
|
|
stringValue := fmt.Sprintf("%t", row.(bson.M)["bool_field"])
|
|
require.Equal(t, stringValue, csvFileColumnValue)
|
|
case "time_field":
|
|
// convert time to string
|
|
stringValue := row.(bson.M)["time_field"].(time.Time).Format("2006-01-02 15:04:05")
|
|
require.Equal(t, stringValue, csvFileColumnValue)
|
|
}
|
|
}
|
|
}
|
|
}
|