[prism] Simpler io changes to add TestMain. (#27587)

Co-authored-by: lostluck <[email protected]>
apache · Jul 21, 2023 · 64552f7 · 64552f7
1 parent cf20649
commit 64552f7
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 7 deletions.
diff --git a/sdks/go/pkg/beam/io/bigtableio/bigtable_test.go b/sdks/go/pkg/beam/io/bigtableio/bigtable_test.go
@@ -23,8 +23,13 @@ import (
 
 	"cloud.google.com/go/bigtable"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
 )
 
+func TestMain(m *testing.M) {
+	ptest.Main(m)
+}
+
 func TestHashStringToInt(t *testing.T) {
 	equalVal := "equal"
 

diff --git a/sdks/go/pkg/beam/io/databaseio/database_test.go b/sdks/go/pkg/beam/io/databaseio/database_test.go
@@ -27,6 +27,10 @@ import (
 	_ "github.com/proullon/ramsql/driver"
 )
 
+func TestMain(m *testing.M) {
+	ptest.Main(m)
+}
+
 type Address struct {
 	Street        string
 	Street_number int

diff --git a/sdks/go/pkg/beam/io/fhirio/deidentify_test.go b/sdks/go/pkg/beam/io/fhirio/deidentify_test.go
@@ -24,6 +24,12 @@ import (
 	"google.golang.org/api/healthcare/v1"
 )
 
+func TestMain(m *testing.M) {
+	// TODO(https://github.com/apache/beam/issues/27547): Make tests compatible with portable runners.
+	// To work on this change, replace call with `ptest.Main(m)`
+	ptest.MainWithDefault(m, "direct")
+}
+
 func TestDeidentify_Error(t *testing.T) {
 	p, s := beam.NewPipelineWithRoot()
 	out := deidentify(s, "src", "dst", nil, requestReturnErrorFakeClient)

diff --git a/sdks/go/pkg/beam/io/fileio/match_test.go b/sdks/go/pkg/beam/io/fileio/match_test.go
@@ -27,6 +27,10 @@ import (
 	"github.com/google/go-cmp/cmp"
 )
 
+func TestMain(m *testing.M) {
+	ptest.Main(m)
+}
+
 type testFile struct {
 	filename string
 	data     []byte

diff --git a/sdks/go/pkg/beam/io/parquetio/parquetio_test.go b/sdks/go/pkg/beam/io/parquetio/parquetio_test.go
@@ -29,6 +29,10 @@ import (
 	"github.com/xitongsys/parquet-go/reader"
 )
 
+func TestMain(m *testing.M) {
+	ptest.Main(m)
+}
+
 type Student struct {
 	Name    string  `parquet:"name=name, type=BYTE_ARRAY, convertedtype=UTF8, encoding=PLAIN_DICTIONARY"`
 	Age     int32   `parquet:"name=age, type=INT32, encoding=PLAIN"`

diff --git a/sdks/go/pkg/beam/io/textio/textio_test.go b/sdks/go/pkg/beam/io/textio/textio_test.go
@@ -17,18 +17,26 @@
 package textio
 
 import (
-	"context"
 	"errors"
 	"os"
 	"path/filepath"
 	"testing"
 
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
 	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/io/filesystem/local"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
 )
 
+func TestMain(m *testing.M) {
+	ptest.Main(m)
+}
+
+func init() {
+	register.Function2x1(toKV)
+}
+
 const testDir = "../../../../data"
 
 var (
@@ -144,9 +152,7 @@ func TestReadSdf(t *testing.T) {
 	lines := ReadSdf(s, testFilePath)
 	passert.Count(s, lines, "NumLines", 1)
 
-	if _, err := beam.Run(context.Background(), "direct", p); err != nil {
-		t.Fatalf("Failed to execute job: %v", err)
-	}
+	ptest.RunAndValidate(t, p)
 }
 
 func TestReadAllSdf(t *testing.T) {
@@ -155,7 +161,5 @@ func TestReadAllSdf(t *testing.T) {
 	lines := ReadAllSdf(s, files)
 	passert.Count(s, lines, "NumLines", 1)
 
-	if _, err := beam.Run(context.Background(), "direct", p); err != nil {
-		t.Fatalf("Failed to execute job: %v", err)
-	}
+	ptest.RunAndValidate(t, p)
 }
diff --git a/sdks/go/pkg/beam/x/debug/head_test.go b/sdks/go/pkg/beam/x/debug/head_test.go
@@ -23,6 +23,10 @@ import (
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
 )
 
+func TestMain(m *testing.M) {
+	ptest.Main(m)
+}
+
 func TestHead(t *testing.T) {
 	p, s, sequence := ptest.CreateList([]int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
 	headSequence := Head(s, sequence, 5)