use check for tests

Author: matthias-kloeckner

postgres_exporter_test.go

@@ -86,66 +86,66 @@ func (s *FunctionalSuite) TestSemanticVersionColumnDiscard(c *C) {
 	}
 }
 
-func TestEnvironmentSettingWithSecretsFiles(t *testing.T) {
+// test read username and password from file
+func (s *FunctionalSuite) TestEnvironmentSettingWithSecretsFiles(c *C) {
 
 	err := os.Setenv("DATA_SOURCE_USER_FILE", "./tests/username_file")
-	if err != nil {
-		t.Errorf("DATA_SOURCE_USER_FILE could not be set")
-	}
+	c.Assert(err, IsNil)
+	defer UnsetEnvironment(c, "DATA_SOURCE_USER_FILE")
 
 	err = os.Setenv("DATA_SOURCE_PASS_FILE", "./tests/userpass_file")
-	if err != nil {
-		t.Errorf("DATA_SOURCE_PASS_FILE could not be set")
-	}
+	c.Assert(err, IsNil)
+	defer UnsetEnvironment(c, "DATA_SOURCE_PASS_FILE")
 
 	err = os.Setenv("DATA_SOURCE_URI", "localhost:5432/?sslmode=disable")
-	if err != nil {
-		t.Errorf("DATA_SOURCE_URI could not be set")
-	}
+	c.Assert(err, IsNil)
+	defer UnsetEnvironment(c, "DATA_SOURCE_URI")
 
 	var expected = "postgresql://custom_username:custom_password@localhost:5432/?sslmode=disable"
 
 	dsn := getDataSource()
 	if dsn != expected {
-		t.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, expected)
+		c.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, expected)
 	}
 }
 
-func TestEnvironmentSettingWithDns(t *testing.T) {
+// test read DATA_SOURCE_NAME from environment
+func (s *FunctionalSuite) TestEnvironmentSettingWithDns(c *C) {
 
 	envDsn := "postgresql://user:password@localhost:5432/?sslmode=enabled"
 	err := os.Setenv("DATA_SOURCE_NAME", envDsn)
-	if err != nil {
-		t.Errorf("DATA_SOURCE_NAME could not be set")
-	}
+	c.Assert(err, IsNil)
+	defer UnsetEnvironment(c, "DATA_SOURCE_NAME")
 
 	dsn := getDataSource()
 	if dsn != envDsn {
-		t.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, envDsn)
+		c.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, envDsn)
 	}
 }
 
-// test DATA_SOURCE_NAME is used even if username and password environment wariables are set
-func TestEnvironmentSettingWithDnsAndSecrets(t *testing.T) {
+// test DATA_SOURCE_NAME is used even if username and password environment variables are set
+func (s *FunctionalSuite) TestEnvironmentSettingWithDnsAndSecrets(c *C) {
 
 	envDsn := "postgresql://userDsn:passwordDsn@localhost:55432/?sslmode=disabled"
 	err := os.Setenv("DATA_SOURCE_NAME", envDsn)
-	if err != nil {
-		t.Errorf("DATA_SOURCE_NAME could not be set")
-	}
+	c.Assert(err, IsNil)
+	defer UnsetEnvironment(c, "DATA_SOURCE_NAME")
 
 	err = os.Setenv("DATA_SOURCE_USER_FILE", "./tests/username_file")
-	if err != nil {
-		t.Errorf("DATA_SOURCE_USER_FILE could not be set")
-	}
+	c.Assert(err, IsNil)
+	defer UnsetEnvironment(c, "DATA_SOURCE_USER_FILE")
 
 	err = os.Setenv("DATA_SOURCE_PASS", "envUserPass")
-	if err != nil {
-		t.Errorf("DATA_SOURCE_PASS could not be set")
-	}
+	c.Assert(err, IsNil)
+	defer UnsetEnvironment(c, "DATA_SOURCE_PASS")
 
 	dsn := getDataSource()
 	if dsn != envDsn {
-		t.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, envDsn)
+		c.Errorf("Expected Username to be read from file. Found=%v, expected=%v", dsn, envDsn)
 	}
 }
+
+func UnsetEnvironment(c *C, d string) {
+	err := os.Unsetenv(d)
+	c.Assert(err, IsNil)
+}