use a second client in the test

Signed-off-by: Grant Linville <[email protected]>

Author: g-linville

datasets.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"os"
 )
 
 type DatasetElementMeta struct {
@@ -50,7 +49,7 @@ func (g *GPTScript) ListDatasets(ctx context.Context) ([]DatasetMeta, error) {
 	out, err := g.runBasicCommand(ctx, "datasets", datasetRequest{
 		Input:       "{}",
 		DatasetTool: g.globalOpts.DatasetTool,
-		Env:         append(g.globalOpts.Env, "GPTSCRIPT_WORKSPACE_ID="+os.Getenv("GPTSCRIPT_WORKSPACE_ID")),
+		Env:         g.globalOpts.Env,
 	})
 	if err != nil {
 		return nil, err
@@ -94,7 +93,7 @@ func (g *GPTScript) AddDatasetElements(ctx context.Context, datasetID string, el
 	return g.runBasicCommand(ctx, "datasets/add-elements", datasetRequest{
 		Input:       string(argsJSON),
 		DatasetTool: g.globalOpts.DatasetTool,
-		Env:         append(g.globalOpts.Env, "GPTSCRIPT_WORKSPACE_ID="+os.Getenv("GPTSCRIPT_WORKSPACE_ID")),
+		Env:         g.globalOpts.Env,
 	})
 }
 
@@ -110,7 +109,7 @@ func (g *GPTScript) ListDatasetElements(ctx context.Context, datasetID string) (
 	out, err := g.runBasicCommand(ctx, "datasets/list-elements", datasetRequest{
 		Input:       string(argsJSON),
 		DatasetTool: g.globalOpts.DatasetTool,
-		Env:         append(g.globalOpts.Env, "GPTSCRIPT_WORKSPACE_ID="+os.Getenv("GPTSCRIPT_WORKSPACE_ID")),
+		Env:         g.globalOpts.Env,
 	})
 	if err != nil {
 		return nil, err
@@ -136,7 +135,7 @@ func (g *GPTScript) GetDatasetElement(ctx context.Context, datasetID, elementNam
 	out, err := g.runBasicCommand(ctx, "datasets/get-element", datasetRequest{
 		Input:       string(argsJSON),
 		DatasetTool: g.globalOpts.DatasetTool,
-		Env:         append(g.globalOpts.Env, "GPTSCRIPT_WORKSPACE_ID="+os.Getenv("GPTSCRIPT_WORKSPACE_ID")),
+		Env:         g.globalOpts.Env,
 	})
 	if err != nil {
 		return DatasetElement{}, err

datasets_test.go

@@ -12,13 +12,17 @@ func TestDatasets(t *testing.T) {
 	workspaceID, err := g.CreateWorkspace(context.Background(), "directory")
 	require.NoError(t, err)
 
-	require.NoError(t, os.Setenv("GPTSCRIPT_WORKSPACE_ID", workspaceID))
+	client, err := NewGPTScript(GlobalOptions{
+		OpenAIAPIKey: os.Getenv("OPENAI_API_KEY"),
+		Env:          []string{"GPTSCRIPT_WORKSPACE_ID=" + workspaceID},
+	})
+	require.NoError(t, err)
 
 	defer func() {
 		_ = g.DeleteWorkspace(context.Background(), workspaceID)
 	}()
 
-	datasetID, err := g.CreateDatasetWithElements(context.Background(), []DatasetElement{
+	datasetID, err := client.CreateDatasetWithElements(context.Background(), []DatasetElement{
 		{
 			DatasetElementMeta: DatasetElementMeta{
 				Name:        "test-element-1",
@@ -33,7 +37,7 @@ func TestDatasets(t *testing.T) {
 	require.NoError(t, err)
 
 	// Add three more elements
-	_, err = g.AddDatasetElements(context.Background(), datasetID, []DatasetElement{
+	_, err = client.AddDatasetElements(context.Background(), datasetID, []DatasetElement{
 		{
 			DatasetElementMeta: DatasetElementMeta{
 				Name:        "test-element-2",
@@ -59,33 +63,33 @@ func TestDatasets(t *testing.T) {
 	require.NoError(t, err)
 
 	// Get the first element
-	element, err := g.GetDatasetElement(context.Background(), datasetID, "test-element-1")
+	element, err := client.GetDatasetElement(context.Background(), datasetID, "test-element-1")
 	require.NoError(t, err)
 	require.Equal(t, "test-element-1", element.Name)
 	require.Equal(t, "This is a test element 1", element.Description)
 	require.Equal(t, "This is the content 1", element.Contents)
 
 	// Get the third element
-	element, err = g.GetDatasetElement(context.Background(), datasetID, "test-element-3")
+	element, err = client.GetDatasetElement(context.Background(), datasetID, "test-element-3")
 	require.NoError(t, err)
 	require.Equal(t, "test-element-3", element.Name)
 	require.Equal(t, "This is a test element 3", element.Description)
 	require.Equal(t, "This is the content 3", element.Contents)
 
 	// Get the binary element
-	element, err = g.GetDatasetElement(context.Background(), datasetID, "binary-element")
+	element, err = client.GetDatasetElement(context.Background(), datasetID, "binary-element")
 	require.NoError(t, err)
 	require.Equal(t, "binary-element", element.Name)
 	require.Equal(t, "this element has binary contents", element.Description)
 	require.Equal(t, []byte("binary contents"), element.BinaryContents)
 
 	// List elements in the dataset
-	elements, err := g.ListDatasetElements(context.Background(), datasetID)
+	elements, err := client.ListDatasetElements(context.Background(), datasetID)
 	require.NoError(t, err)
 	require.Equal(t, 4, len(elements))
 
 	// List datasets
-	datasets, err := g.ListDatasets(context.Background())
+	datasets, err := client.ListDatasets(context.Background())
 	require.NoError(t, err)
 	require.Equal(t, 1, len(datasets))
 	require.Equal(t, datasetID, datasets[0].ID)