outputs/postgresql: fix go test detection for pguint extension

phemmer · Feb 18, 2022 · 0e511d6 · 0e511d6
1 parent ddb5575
commit 0e511d6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 16 deletions.
diff --git a/plugins/outputs/postgresql/postgresql_test.go b/plugins/outputs/postgresql/postgresql_test.go
@@ -678,14 +678,12 @@ func TestWrite_UnsignedIntegers(t *testing.T) {
 	p := newPostgresqlTest(t)
 	p.UseUint8 = true
 	_ = p.Init()
-	require.NoError(t, p.Connect())
-
-	row := p.db.QueryRow(ctx, "SELECT count(*) FROM pg_extension WHERE extname='uint'")
-	var n int
-	require.NoError(t, row.Scan(&n))
-	if n == 0 {
-		t.Skipf("pguint extension is not installed")
-		t.SkipNow()
+	if err := p.Connect(); err != nil {
+		if strings.Contains(err.Error(), "retreiving OID for uint8 data type") {
+			t.Skipf("pguint extension is not installed")
+			t.SkipNow()
+		}
+		require.NoError(t, err)
 	}
 
 	metrics := []telegraf.Metric{

diff --git a/plugins/outputs/postgresql/table_manager_test.go b/plugins/outputs/postgresql/table_manager_test.go
@@ -127,14 +127,12 @@ func TestTableManager_MatchSource_UnsignedIntegers(t *testing.T) {
 	p := newPostgresqlTest(t)
 	p.UseUint8 = true
 	_ = p.Init()
-	require.NoError(t, p.Connect())
-
-	row := p.db.QueryRow(ctx, "SELECT count(*) FROM pg_extension WHERE extname='uint'")
-	var n int
-	require.NoError(t, row.Scan(&n))
-	if n == 0 {
-		t.Skipf("pguint extension is not installed")
-		t.SkipNow()
+	if err := p.Connect(); err != nil {
+		if strings.Contains(err.Error(), "retreiving OID for uint8 data type") {
+			t.Skipf("pguint extension is not installed")
+			t.SkipNow()
+		}
+		require.NoError(t, err)
 	}
 
 	metrics := []telegraf.Metric{