From efe4879ef382945696798e78bbc91a8fdeb6cbcf Mon Sep 17 00:00:00 2001
From: Sebastian Rivera <sebastian.rivera@hashicorp.com>
Date: Mon, 6 Jun 2022 14:02:53 -0400
Subject: [PATCH] Improve audit trail interface by using client token

---
 audit_trail.go                  | 12 +++++++-----
 audit_trail_integration_test.go | 18 ++++++++----------
 helper_test.go                  | 16 ++++++++++++++++
 3 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/audit_trail.go b/audit_trail.go
index 6515cfa83..72e498f40 100644
--- a/audit_trail.go
+++ b/audit_trail.go
@@ -15,13 +15,14 @@ import (
 var _ AuditTrails = (*auditTrails)(nil)
 
 // AuditTrails describes all the audit event related methods that the Terraform
-// Cloud API supports. **Note:** These methods require an organization token for
-// an organization in the Business tier and are only available in Terraform Cloud.
+// Cloud API supports.
+// **Note:** These methods require the client to be configured with an organization token for
+// an organization in the Business tier. Furthermore, these methods are only available in Terraform Cloud.
 //
 // TFC API Docs: https://www.terraform.io/cloud-docs/api-docs/audit-trails
 type AuditTrails interface {
 	// Read all the audit events in an organization.
-	List(ctx context.Context, orgToken string, options *AuditTrailListOptions) (*AuditTrailList, error)
+	List(ctx context.Context, options *AuditTrailListOptions) (*AuditTrailList, error)
 }
 
 // auditTrails implements AuditTrails
@@ -78,14 +79,15 @@ type AuditTrailListOptions struct {
 }
 
 // List all the audit events in an organization.
-func (s *auditTrails) List(ctx context.Context, orgToken string, options *AuditTrailListOptions) (*AuditTrailList, error) {
+func (s *auditTrails) List(ctx context.Context, options *AuditTrailListOptions) (*AuditTrailList, error) {
 	u, err := s.client.baseURL.Parse("/api/v2/organization/audit-trail")
 	if err != nil {
 		return nil, err
 	}
 
 	headers := make(http.Header)
-	headers.Set("Authorization", "Bearer "+orgToken)
+	headers.Set("User-Agent", _userAgent)
+	headers.Set("Authorization", "Bearer "+s.client.token)
 	headers.Set("Content-Type", "application/json")
 
 	if options != nil {
diff --git a/audit_trail_integration_test.go b/audit_trail_integration_test.go
index d4ac34660..b0152fc16 100644
--- a/audit_trail_integration_test.go
+++ b/audit_trail_integration_test.go
@@ -15,25 +15,23 @@ import (
 func TestAuditTrailsList(t *testing.T) {
 	skipIfEnterprise(t)
 
-	client := testClient(t)
+	userClient := testClient(t)
 	ctx := context.Background()
 
-	org, orgCleanup := createOrganization(t, client)
+	org, orgCleanup := createOrganization(t, userClient)
 	t.Cleanup(orgCleanup)
 
-	upgradeOrganizationSubscription(t, client, org)
+	auditTrailClient := testAuditTrailClient(t, userClient, org)
 
-	orgToken, orgTokenCleanup := createOrganizationToken(t, client, org)
-	t.Cleanup(orgTokenCleanup)
 	// First let's generate some audit events in this test organization
-	_, wkspace1Cleanup := createWorkspace(t, client, org)
+	_, wkspace1Cleanup := createWorkspace(t, userClient, org)
 	t.Cleanup(wkspace1Cleanup)
 
-	_, wkspace2Cleanup := createWorkspace(t, client, org)
+	_, wkspace2Cleanup := createWorkspace(t, userClient, org)
 	t.Cleanup(wkspace2Cleanup)
 
 	t.Run("with no specified timeframe", func(t *testing.T) {
-		atl, err := client.AuditTrails.List(ctx, orgToken.Token, nil)
+		atl, err := auditTrailClient.AuditTrails.List(ctx, nil)
 		require.NoError(t, err)
 		require.Greater(t, len(atl.Items), 0)
 
@@ -75,10 +73,10 @@ func TestAuditTrailsList(t *testing.T) {
 		time.Sleep(1 * time.Second)
 
 		// Let's create an event that is sent to the audit log
-		_, wsCleanup := createWorkspace(t, client, org)
+		_, wsCleanup := createWorkspace(t, userClient, org)
 		t.Cleanup(wsCleanup)
 
-		atl, err := client.AuditTrails.List(ctx, orgToken.Token, &AuditTrailListOptions{
+		atl, err := auditTrailClient.AuditTrails.List(ctx, &AuditTrailListOptions{
 			Since: since,
 			ListOptions: &ListOptions{
 				PageNumber: 1,
diff --git a/helper_test.go b/helper_test.go
index 3d04c0b9f..c8e46cada 100644
--- a/helper_test.go
+++ b/helper_test.go
@@ -60,6 +60,22 @@ func testClient(t *testing.T) *Client {
 	return client
 }
 
+func testAuditTrailClient(t *testing.T, userClient *Client, org *Organization) *Client {
+	upgradeOrganizationSubscription(t, userClient, org)
+
+	orgToken, orgTokenCleanup := createOrganizationToken(t, userClient, org)
+	t.Cleanup(orgTokenCleanup)
+
+	client, err := NewClient(&Config{
+		Token: orgToken.Token,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return client
+}
+
 func fetchTestAccountDetails(t *testing.T, client *Client) *TestAccountDetails {
 	if _testAccountDetails == nil {
 		_testAccountDetails = FetchTestAccountDetails(t, client)