github.com/qri-io/qri@v0.10.1-0.20220104210721-c771715036cb/transform/startf/ds/dataset_test.go (about)

     1  package ds
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  
     7  	"github.com/google/go-cmp/cmp"
     8  	"github.com/qri-io/dataset"
     9  	"github.com/qri-io/qfs"
    10  	"github.com/qri-io/starlib/dataframe"
    11  	"github.com/qri-io/starlib/testdata"
    12  	"go.starlark.net/resolve"
    13  	"go.starlark.net/starlark"
    14  	"go.starlark.net/starlarktest"
    15  )
    16  
    17  func callMethod(thread *starlark.Thread, v starlark.HasAttrs, name string, tuple starlark.Tuple) (starlark.Value, error) {
    18  	method, err := v.Attr(name)
    19  	if err != nil {
    20  		return nil, err
    21  	}
    22  	if method == nil {
    23  		return nil, fmt.Errorf("method %s does not exist", name)
    24  	}
    25  	return starlark.Call(thread, method, tuple, nil)
    26  }
    27  
    28  func TestCannotSetIfReadOnly(t *testing.T) {
    29  	ds := NewDataset(&dataset.Dataset{}, nil)
    30  	ds.Freeze()
    31  	expect := "cannot set, Dataset is frozen"
    32  	err := ds.SetField("body", starlark.NewList([]starlark.Value{starlark.NewList([]starlark.Value{starlark.String("a")})}))
    33  	if err == nil {
    34  		t.Fatal("expected error, did not get one")
    35  	}
    36  	if err.Error() != expect {
    37  		t.Errorf("expected error: %s, got: %s", expect, err)
    38  	}
    39  }
    40  
    41  func TestSetAndGetBody(t *testing.T) {
    42  	outconf := &dataframe.OutputConfig{}
    43  	ds := NewDataset(&dataset.Dataset{}, outconf)
    44  	err := ds.SetField("body", starlark.NewList([]starlark.Value{starlark.NewList([]starlark.Value{starlark.String("a")})}))
    45  	if err != nil {
    46  		t.Fatal(err)
    47  	}
    48  	expect := `     0
    49  0    a`
    50  	bd, _ := ds.Attr("body")
    51  	actual := bd.String()
    52  	if diff := cmp.Diff(expect, actual); diff != "" {
    53  		t.Errorf("result mismatch (-want +got):\n%s", diff)
    54  	}
    55  }
    56  
    57  func TestFile(t *testing.T) {
    58  	resolve.AllowFloat = true
    59  	thread := &starlark.Thread{Load: newLoader()}
    60  	starlarktest.SetReporter(thread, t)
    61  	dataframe.SetOutputSize(thread, 0, 0)
    62  
    63  	// Execute test file
    64  	_, err := starlark.ExecFile(thread, "testdata/test.star", nil, starlark.StringDict{
    65  		"csv_ds": csvDataset(),
    66  	})
    67  	if err != nil {
    68  		if ee, ok := err.(*starlark.EvalError); ok {
    69  			t.Error(ee.Backtrace())
    70  		} else {
    71  			t.Error(err)
    72  		}
    73  	}
    74  }
    75  
    76  // load implements the 'load' operation as used in the evaluator tests.
    77  func newLoader() func(thread *starlark.Thread, module string) (starlark.StringDict, error) {
    78  	return testdata.NewLoader(LoadModule, ModuleName)
    79  }
    80  
    81  func csvDataset() *Dataset {
    82  	text := `title,count,is great
    83  foo,1,true
    84  bar,2,false
    85  bat,3,meh
    86  `
    87  	ds := &dataset.Dataset{
    88  		Structure: &dataset.Structure{
    89  			Format: "csv",
    90  			FormatConfig: map[string]interface{}{
    91  				"headerRow": true,
    92  			},
    93  			Schema: map[string]interface{}{
    94  				"type": "array",
    95  				"items": map[string]interface{}{
    96  					"type": "array",
    97  					"items": []interface{}{
    98  						map[string]interface{}{"title": "title", "type": "string"},
    99  						map[string]interface{}{"title": "count", "type": "integer"},
   100  						map[string]interface{}{"title": "is great", "type": "string"},
   101  					},
   102  				},
   103  			},
   104  		},
   105  	}
   106  	ds.SetBodyFile(qfs.NewMemfileBytes("body.csv", []byte(text)))
   107  
   108  	outconf := &dataframe.OutputConfig{}
   109  	d := NewDataset(ds, outconf)
   110  	return d
   111  }
   112  
   113  func TestCreateColumnsFromStructure(t *testing.T) {
   114  	cases := []struct {
   115  		desc   string
   116  		schema map[string]interface{}
   117  		expect []string
   118  	}{
   119  		{
   120  			// description
   121  			"typical schema and the columns it turns into",
   122  			//schema
   123  			map[string]interface{}{
   124  				"items": map[string]interface{}{
   125  					"items": []interface{}{
   126  						map[string]interface{}{
   127  							"title": "name",
   128  							"type":  "string",
   129  						},
   130  						map[string]interface{}{
   131  							"title": "sound",
   132  							"type":  "string",
   133  						},
   134  						map[string]interface{}{
   135  							"title": "weight",
   136  							"type":  "number",
   137  						},
   138  					},
   139  				},
   140  			},
   141  			//expect
   142  			[]string{"name", "sound", "weight"},
   143  		},
   144  
   145  		{
   146  			// description
   147  			"base array schema will not return any columns",
   148  			//schema
   149  			map[string]interface{}{
   150  				"type": "array",
   151  			},
   152  			//expect
   153  			nil,
   154  		},
   155  	}
   156  	for i, c := range cases {
   157  		ds := &dataset.Dataset{
   158  			Structure: &dataset.Structure{
   159  				Schema: c.schema,
   160  			},
   161  		}
   162  		d := NewDataset(ds, nil)
   163  		actual := d.createColumnsFromStructure()
   164  		if diff := cmp.Diff(c.expect, actual); diff != "" {
   165  			t.Errorf("case %d %s: mismatch (-want +got):\n%s", i, c.desc, diff)
   166  		}
   167  	}
   168  }