Add run_tests.sh script for generating mock and run tests

go/README.md

@@ -64,7 +64,13 @@ Here we're querying for a dataset called `NYC-taxi-trips`, in a source called `S
 
 ## Tests
 
-This example Go client also has unit tests that can be run using the following command:
+To run the tests, you'll need a flight client mock class. This class is generated using mockgen. To aid in this process,
+we created a script that generates the mock class and runs all tests. You can run the script using the following command:
+```bash
+./run_tests.sh 
+```
+
+If the mock class is already generated you can alternatively use the following command:
 ```go
 go test -v
 ```

go/example.go

@@ -13,14 +13,13 @@
 package main
 
 import (
+	"arrow-flight-client-example/interfaces"
 	"context"
 	"crypto/tls"
 	"fmt"
-	"github.com/apache/arrow-go/v18/arrow"
 	"log"
 	"net"
 
-	"arrow-flight-client-example/interfaces"
 	"github.com/apache/arrow-go/v18/arrow/flight"
 	flightgen "github.com/apache/arrow-go/v18/arrow/flight/gen/flight"
 	"github.com/apache/arrow-go/v18/arrow/memory"
@@ -50,14 +49,7 @@ Options:
   --certs=<path>      Path to trusted certificates for encrypted connection.
   --project_id=<project_id>   Dremio project ID`
 
-type RecordReader interface {
-	Next() bool
-	Record() arrow.Record
-	Err() error
-	Release()
-}
-
-func WrapRecordReader(stream flight.FlightService_DoGetClient) (RecordReader, error) {
+func WrapRecordReader(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
 	return flight.NewRecordReader(stream)
 }
 
@@ -67,7 +59,7 @@ func main() {
 		log.Fatalf("error parsing arguments: %v", err)
 	}
 
-	var config interfaces.FlightConfig
+	var config FlightConfig
 	if err := args.Bind(&config); err != nil {
 		log.Fatalf("error binding arguments: %v", err)
 	}
@@ -114,8 +106,8 @@ func main() {
 	}
 }
 
-func run(config interfaces.FlightConfig, client flight.Client,
-	readerCreator func(flight.FlightService_DoGetClient) (RecordReader, error),
+func run(config FlightConfig, client flight.Client,
+	readerCreator func(flight.FlightService_DoGetClient) (interfaces.RecordReader, error),
 ) error {
 
 	// Two WLM settings can be provided upon initial authentication with the dremio

go/example_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"arrow-flight-client-example/interfaces"
 	"bytes"
 	"context"
 	"github.com/apache/arrow-go/v18/arrow"
@@ -16,7 +17,6 @@ import (
 	"github.com/golang/mock/gomock"
 
 	"arrow-flight-client-example/implementations"
-	"arrow-flight-client-example/interfaces"
 )
 
 func TestUsernamePassAuth(t *testing.T) {
@@ -30,7 +30,7 @@ func TestUsernamePassAuth(t *testing.T) {
 		Return(context.Background(), nil).
 		Times(1)
 
-	config := interfaces.FlightConfig{
+	config := FlightConfig{
 		User:  "testuser",
 		Pass:  "testpass",
 		Query: "",
@@ -58,7 +58,7 @@ func TestPATAuth(t *testing.T) {
 		Return(nil, nil).
 		Times(1)
 
-	config := interfaces.FlightConfig{
+	config := FlightConfig{
 		Pat:       "testpat",
 		ProjectID: "testproject",
 	}
@@ -138,13 +138,13 @@ func TestRun(t *testing.T) {
 		Return(mockStream, nil).
 		Times(1)
 
-	config := interfaces.FlightConfig{
+	config := FlightConfig{
 		User:  "testuser",
 		Pass:  "testpass",
 		Query: "SELECT * FROM test",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
 		return implementations.NewMockRecordReader([]arrow.Record{record}), nil
 	}
 
@@ -225,13 +225,13 @@ func TestRunWithPAT(t *testing.T) {
 		Return(nil, nil).
 		Times(1)
 
-	config := interfaces.FlightConfig{
+	config := FlightConfig{
 		Pat:       "test_pat_token",
 		Query:     "SELECT * FROM test",
 		ProjectID: "test_project_id",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
 		return implementations.NewMockRecordReader([]arrow.Record{record}), nil
 	}
 
@@ -302,12 +302,12 @@ func TestRunWithPATNoProjectID(t *testing.T) {
 		Return(mockStream, nil).
 		Times(1)
 
-	config := interfaces.FlightConfig{
+	config := FlightConfig{
 		Pat:   "test_pat_token",
 		Query: "SELECT * FROM test",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
 		return implementations.NewMockRecordReader([]arrow.Record{record}), nil
 	}
 
@@ -329,13 +329,13 @@ func TestInvalidCredentials(t *testing.T) {
 		Return(context.Background(), expectedErr).
 		Times(1)
 
-	config := interfaces.FlightConfig{
+	config := FlightConfig{
 		User:  "dremio",
 		Pass:  "dremio12",
 		Query: "SELECT 1",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
 		t.Fatal("Reader creator should not be called due to authentication failure")
 		return nil, nil
 	}
@@ -367,13 +367,13 @@ func TestInvalidHost(t *testing.T) {
 		Return(context.Background(), expectedErr).
 		Times(1)
 
-	config := interfaces.FlightConfig{
+	config := FlightConfig{
 		User:  "dremio",
 		Pass:  "dremio12",
 		Query: "SELECT 1",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
 		t.Fatal("Reader creator should not be called due to authentication failure")
 		return nil, nil
 	}
@@ -405,13 +405,13 @@ func TestInvalidPort(t *testing.T) {
 		Return(context.Background(), expectedErr).
 		Times(1)
 
-	config := interfaces.FlightConfig{
+	config := FlightConfig{
 		User:  "dremio",
 		Pass:  "dremio12",
 		Query: "SELECT 1",
 	}
 
-	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (RecordReader, error) {
+	mockReaderCreator := func(stream flight.FlightService_DoGetClient) (interfaces.RecordReader, error) {
 		t.Fatal("Reader creator should not be called due to authentication failure")
 		return nil, nil
 	}

go/interfaces/flight_config.go → go/flight_config.go

@@ -1,4 +1,4 @@
-package interfaces
+package main
 
 type FlightConfig struct {
 	Host      string

go/implementations/mock_record_reader.go

@@ -1,9 +1,12 @@
 package implementations
 
 import (
+	"arrow-flight-client-example/interfaces"
 	"github.com/apache/arrow-go/v18/arrow"
 )
 
+var _ interfaces.RecordReader = (*MockRecordReader)(nil)
+
 // MockRecordReader for testing purposes
 type MockRecordReader struct {
 	records      []arrow.Record

go/interfaces/record_reader.go

@@ -0,0 +1,10 @@
+package interfaces
+
+import "github.com/apache/arrow-go/v18/arrow"
+
+type RecordReader interface {
+	Next() bool
+	Record() arrow.Record
+	Err() error
+	Release()
+}