From b5fec01eb4fd13237b3da9f77737c57ea2d4319a Mon Sep 17 00:00:00 2001
From: Kenan Yildirim <kenan.yildirim@pangea.cloud>
Date: Fri, 3 May 2024 13:54:27 -0400
Subject: [PATCH] fixup! address PR comments

---
 pangea-sdk/v3/service/audit/api.go            |  6 ++--
 .../v3/service/audit/integration_test.go      | 28 +++++++++++++------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/pangea-sdk/v3/service/audit/api.go b/pangea-sdk/v3/service/audit/api.go
index 5addccfc..3d5e0cef 100644
--- a/pangea-sdk/v3/service/audit/api.go
+++ b/pangea-sdk/v3/service/audit/api.go
@@ -936,7 +936,7 @@ type RootOutput struct {
 type DownloadRequest struct {
 	pangea.BaseRequest
 
-	// ID returned by the search API.
+	// ID returned by the export API.
 	RequestID string `json:"request_id,omitempty"`
 
 	// ID returned by the search API.
@@ -1036,11 +1036,11 @@ type ExportRequest struct {
 	Format *DownloadFormat `json:"format,omitempty"`
 
 	// The start of the time range to perform the search on.
-	Start *time.Time `json:"start,omitempty"`
+	Start *pu.PangeaTimestamp `json:"start,omitempty"`
 
 	// The end of the time range to perform the search on. If omitted, then all
 	// records up to the latest will be searched.
-	End *time.Time `json:"end,omitempty"`
+	End *pu.PangeaTimestamp `json:"end,omitempty"`
 
 	// Specify the sort order of the response, "asc" or "desc".
 	Order *string `json:"order,omitempty"`
diff --git a/pangea-sdk/v3/service/audit/integration_test.go b/pangea-sdk/v3/service/audit/integration_test.go
index b3fc7287..edff60d2 100644
--- a/pangea-sdk/v3/service/audit/integration_test.go
+++ b/pangea-sdk/v3/service/audit/integration_test.go
@@ -1003,26 +1003,38 @@ func Test_Integration_LogStream(t *testing.T) {
 }
 
 func Test_Integration_Export_Download(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
 	cfg := auditIntegrationCfg(t)
 	client, _ := audit.New(cfg)
 
-	exportRes, err := client.Export(ctx, &audit.ExportRequest{Verbose: pangea.Bool(false)})
+	exportRes, err := client.Export(ctx, &audit.ExportRequest{
+		Start:   pangea.PangeaTime(pu.PangeaTimestamp(time.Now().UTC().Add(-12 * time.Hour))),
+		End:     pangea.PangeaTime(pu.PangeaTimestamp(time.Now().UTC())),
+		Verbose: pangea.Bool(false),
+	})
 	assert.NoError(t, err)
 	assert.NotNil(t, exportRes)
 	assert.Equal(t, "Accepted", pangea.StringValue(exportRes.Status))
 	assert.NotEmpty(t, exportRes.RequestID)
 
-	var pollResult *pangea.PangeaResponse[struct{}]
-	pollResponse, err := client.PollResultByID(ctx, *exportRes.RequestID, pollResult)
-	assert.Error(t, err)
-	assert.Nil(t, pollResponse)
+	retry := 0
+	for retry < 10 {
+		_, err := client.PollResultByID(ctx, *exportRes.RequestID, &audit.DownloadResult{})
+		if err == nil {
+			break
+		}
+
+		// Wait until result should be ready
+		time.Sleep(time.Duration(3 * time.Second))
+		retry++
+	}
 
 	downloadRes, err := client.DownloadResults(ctx, &audit.DownloadRequest{
 		RequestID: *exportRes.RequestID,
 	})
-	assert.Error(t, err)
-	assert.Nil(t, downloadRes)
+	assert.NoError(t, err)
+	assert.Equal(t, "Success", pangea.StringValue(downloadRes.Status))
+	assert.NotEmpty(t, downloadRes.Result.DestURL)
 }