Function TestS3Model

runtime/ai/developer_agent_test.go:111–159

(t *testing.T)

}

func TestS3Model(t *testing.T) {
	// Setup a basic empty project
	rt, instanceID := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{
		AIConnector:    "openai",
		TestConnectors: []string{"s3"}, // Add environment variables for the test S3 connector
		Files: map[string]string{
			"rill.yaml": `
olap_connector: duckdb
`,
			"connectors/duckdb.yaml": `
type: connector
driver: duckdb
managed: true
`,
		},
	})
	testruntime.RequireReconcileState(t, rt, instanceID, 2, 0, 0)

	// Initialize eval
	s := newEval(t, rt, instanceID)

	// Ask it to build a DuckDB model for CSV files in S3
	var res *ai.RouterAgentResult
	_, err := s.CallTool(t.Context(), ai.RoleUser, ai.RouterAgentName, &res, ai.RouterAgentArgs{
		Prompt: "I have some CSV files in S3. Can you create a connector for S3 and a DuckDB model that loads the data at s3://integration-test.rilldata.com/glob_test/y=*/*.csv? I've already added environment variables for S3 access. Please proceed without asking clarifying questions.",
	})
	require.NoError(t, err)
	require.Equal(t, ai.DeveloperAgentName, res.Agent)

	// Check that it doesn't have any parse or reconcile errors.
	testruntime.RequireReconcileState(t, rt, instanceID, -1, 0, 0)

	// Check there's a DuckDB and S3 connector
	ctrl, err := rt.Controller(t.Context(), instanceID)
	require.NoError(t, err)
	connectors, err := ctrl.List(t.Context(), runtime.ResourceKindConnector, "", false)
	require.NoError(t, err)
	var connectorNames []string
	for _, c := range connectors {
		connectorNames = append(connectorNames, c.Meta.Name.Name)
	}
	require.Contains(t, connectorNames, "duckdb")
	require.Contains(t, connectorNames, "s3")

	// Check there's a model created
	models, err := ctrl.List(t.Context(), runtime.ResourceKindModel, "", false)
	require.NoError(t, err)
	require.Greater(t, len(models), 0)
}

func TestS3Introspection(t *testing.T) {
	// Setup a project with the S3 connector