11package containers
22
33import (
4+ "bytes"
45 "context"
56 "fmt"
7+ "io"
8+ "os"
9+ "path/filepath"
10+ "strings"
611 "sync/atomic"
712 "testing"
813 "time"
@@ -11,6 +16,7 @@ import (
1116 "github.com/aws/aws-sdk-go-v2/config"
1217 "github.com/aws/aws-sdk-go-v2/credentials"
1318 "github.com/aws/aws-sdk-go-v2/service/s3"
19+ "github.com/aws/aws-sdk-go-v2/service/s3/types"
1420 "github.com/stretchr/testify/require"
1521 "github.com/testcontainers/testcontainers-go"
1622 "github.com/testcontainers/testcontainers-go/wait"
@@ -79,6 +85,139 @@ func (lc *LocalstackTestContainer) MakeS3Connection(ctx context.Context, t *test
7985 return client , bucketName
8086}
8187
88+ // createS3Client creates an S3 client connected to the test container
89+ func (lc * LocalstackTestContainer ) createS3Client (ctx context.Context ) (* s3.Client , error ) {
90+ cfg , err := config .LoadDefaultConfig (ctx ,
91+ config .WithRegion ("us-east-1" ),
92+ config .WithCredentialsProvider (credentials .NewStaticCredentialsProvider ("test" , "test" , "" )),
93+ )
94+ if err != nil {
95+ return nil , fmt .Errorf ("failed to load AWS config: %w" , err )
96+ }
97+
98+ client := s3 .NewFromConfig (cfg , func (o * s3.Options ) {
99+ o .BaseEndpoint = aws .String (lc .Endpoint )
100+ o .UsePathStyle = true
101+ })
102+
103+ return client , nil
104+ }
105+
106+ // GetFile downloads a file from S3
107+ func (lc * LocalstackTestContainer ) GetFile (ctx context.Context , bucketName , objectKey , localPath string ) error {
108+ client , err := lc .createS3Client (ctx )
109+ if err != nil {
110+ return fmt .Errorf ("failed to create S3 client: %w" , err )
111+ }
112+
113+ localDir := filepath .Dir (localPath )
114+ if err := os .MkdirAll (localDir , 0o750 ); err != nil {
115+ return fmt .Errorf ("failed to create local directory %s: %w" , localDir , err )
116+ }
117+
118+ if ! strings .HasPrefix (filepath .Clean (localPath ), filepath .Clean (localDir )) {
119+ return fmt .Errorf ("localPath %s attempts to escape from directory %s" , localPath , localDir )
120+ }
121+
122+ // get object from S3
123+ output , err := client .GetObject (ctx , & s3.GetObjectInput {
124+ Bucket : aws .String (bucketName ),
125+ Key : aws .String (objectKey ),
126+ })
127+ if err != nil {
128+ return fmt .Errorf ("failed to get object %s from bucket %s: %w" , objectKey , bucketName , err )
129+ }
130+ defer output .Body .Close ()
131+
132+ // create local file
133+ file , err := os .OpenFile (localPath , os .O_CREATE | os .O_WRONLY | os .O_TRUNC , 0o600 )
134+ if err != nil {
135+ return fmt .Errorf ("failed to create local file %s: %w" , localPath , err )
136+ }
137+ defer file .Close ()
138+
139+ // copy object content to local file
140+ if _ , err := io .Copy (file , output .Body ); err != nil {
141+ return fmt .Errorf ("failed to copy content from S3 object to local file: %w" , err )
142+ }
143+
144+ return nil
145+ }
146+
147+ // SaveFile uploads a file to S3
148+ func (lc * LocalstackTestContainer ) SaveFile (ctx context.Context , localPath , bucketName , objectKey string ) error {
149+ client , err := lc .createS3Client (ctx )
150+ if err != nil {
151+ return fmt .Errorf ("failed to create S3 client: %w" , err )
152+ }
153+
154+ if ! strings .HasPrefix (filepath .Clean (localPath ), filepath .Clean (filepath .Dir (localPath ))) {
155+ return fmt .Errorf ("localPath %s attempts to escape from its directory" , localPath )
156+ }
157+
158+ // read local file
159+ fileData , err := os .ReadFile (localPath )
160+ if err != nil {
161+ return fmt .Errorf ("failed to read local file %s: %w" , localPath , err )
162+ }
163+
164+ // upload to S3
165+ _ , err = client .PutObject (ctx , & s3.PutObjectInput {
166+ Bucket : aws .String (bucketName ),
167+ Key : aws .String (objectKey ),
168+ Body : bytes .NewReader (fileData ),
169+ })
170+ if err != nil {
171+ return fmt .Errorf ("failed to upload file to S3: %w" , err )
172+ }
173+
174+ return nil
175+ }
176+
177+ // ListFiles lists objects in an S3 bucket, optionally with a prefix
178+ func (lc * LocalstackTestContainer ) ListFiles (ctx context.Context , bucketName , prefix string ) ([]types.Object , error ) {
179+ client , err := lc .createS3Client (ctx )
180+ if err != nil {
181+ return nil , fmt .Errorf ("failed to create S3 client: %w" , err )
182+ }
183+
184+ // list objects
185+ input := & s3.ListObjectsV2Input {
186+ Bucket : aws .String (bucketName ),
187+ }
188+
189+ // add prefix if provided
190+ if prefix != "" {
191+ input .Prefix = aws .String (prefix )
192+ }
193+
194+ output , err := client .ListObjectsV2 (ctx , input )
195+ if err != nil {
196+ return nil , fmt .Errorf ("failed to list objects in bucket %s: %w" , bucketName , err )
197+ }
198+
199+ return output .Contents , nil
200+ }
201+
202+ // DeleteFile deletes an object from S3
203+ func (lc * LocalstackTestContainer ) DeleteFile (ctx context.Context , bucketName , objectKey string ) error {
204+ client , err := lc .createS3Client (ctx )
205+ if err != nil {
206+ return fmt .Errorf ("failed to create S3 client: %w" , err )
207+ }
208+
209+ // delete object
210+ _ , err = client .DeleteObject (ctx , & s3.DeleteObjectInput {
211+ Bucket : aws .String (bucketName ),
212+ Key : aws .String (objectKey ),
213+ })
214+ if err != nil {
215+ return fmt .Errorf ("failed to delete object %s from bucket %s: %w" , objectKey , bucketName , err )
216+ }
217+
218+ return nil
219+ }
220+
82221// Close terminates the container
83222func (lc * LocalstackTestContainer ) Close (ctx context.Context ) error {
84223 return lc .Container .Terminate (ctx )
0 commit comments