From 9b4d2ba4b92b92075f5c1b8688be26e58dd6e0c2 Mon Sep 17 00:00:00 2001
From: Robert Burke <lostluck@users.noreply.github.com>
Date: Fri, 21 Jul 2023 07:37:12 -0700
Subject: [PATCH] [prism] Simpler io changes to add TestMain. (#27587)

Co-authored-by: lostluck <13907733+lostluck@users.noreply.github.com>
---
 .../go/pkg/beam/io/bigtableio/bigtable_test.go |  5 +++++
 .../go/pkg/beam/io/databaseio/database_test.go |  4 ++++
 sdks/go/pkg/beam/io/fhirio/deidentify_test.go  |  6 ++++++
 sdks/go/pkg/beam/io/fileio/match_test.go       |  4 ++++
 .../go/pkg/beam/io/parquetio/parquetio_test.go |  4 ++++
 sdks/go/pkg/beam/io/textio/textio_test.go      | 18 +++++++++++-------
 sdks/go/pkg/beam/x/debug/head_test.go          |  4 ++++
 7 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/sdks/go/pkg/beam/io/bigtableio/bigtable_test.go b/sdks/go/pkg/beam/io/bigtableio/bigtable_test.go
index 4d2dc1b33380..2f41f4ff615e 100644
--- 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
index 1876f5701215..b93d5c9da727 100644
--- 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
index 10f281cd1ed6..caa5b88d7c83 100644
--- 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
index 5bc849e5057e..69e17e9181a4 100644
--- 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
index f5f966ab1693..1cceefcef46b 100644
--- 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
index 10b0f5f4b1b5..ff4579b0db8d 100644
--- a/sdks/go/pkg/beam/io/textio/textio_test.go
+++ b/sdks/go/pkg/beam/io/textio/textio_test.go
@@ -17,7 +17,6 @@
 package textio
 
 import (
-	"context"
 	"errors"
 	"os"
 	"path/filepath"
@@ -25,10 +24,19 @@ import (
 
 	"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
index 8aa5b41545da..5903768d9f7e 100644
--- 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)