alice.pds.demo.boltless.dev
This repository has no description
2.8 kB
130 lines
1package artifactstore
2
3import (
4 "bytes"
5 "context"
6 "io"
7 "os"
8 "strings"
9 "sync"
10 "testing"
11
12 "github.com/aws/aws-sdk-go-v2/service/s3"
13)
14
15type mockS3Client struct {
16 mu sync.Mutex
17 store map[string][]byte
18}
19
20func newMockS3Client() *mockS3Client {
21 return &mockS3Client{
22 store: make(map[string][]byte),
23 }
24}
25
26func (m *mockS3Client) PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
27 m.mu.Lock()
28 defer m.mu.Unlock()
29 b, err := io.ReadAll(params.Body)
30 if err != nil {
31 return nil, err
32 }
33 key := *params.Bucket + "/" + *params.Key
34 m.store[key] = b
35 return &s3.PutObjectOutput{}, nil
36}
37
38func (m *mockS3Client) GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
39 m.mu.Lock()
40 defer m.mu.Unlock()
41 key := *params.Bucket + "/" + *params.Key
42 data, ok := m.store[key]
43 if !ok {
44 return nil, os.ErrNotExist
45 }
46 return &s3.GetObjectOutput{
47 Body: io.NopCloser(bytes.NewReader(data)),
48 }, nil
49}
50
51func TestDiskStore(t *testing.T) {
52 tempDir := t.TempDir()
53 store, err := NewDiskStore(tempDir)
54 if err != nil {
55 t.Fatal(err)
56 }
57
58 ctx := context.Background()
59 ref := "logs/test.log"
60 content := "hello world log content"
61
62 if err := store.Put(ctx, ref, strings.NewReader(content)); err != nil {
63 t.Fatalf("Put failed: %v", err)
64 }
65
66 rc, err := store.Open(ctx, ref)
67 if err != nil {
68 t.Fatalf("Open failed: %v", err)
69 }
70 defer rc.Close()
71
72 got, err := io.ReadAll(rc)
73 if err != nil {
74 t.Fatalf("ReadAll failed: %v", err)
75 }
76 if string(got) != content {
77 t.Fatalf("got content %q, want %q", string(got), content)
78 }
79}
80
81func TestDiskStoreTraversalProtection(t *testing.T) {
82 tempDir := t.TempDir()
83 store, err := NewDiskStore(tempDir)
84 if err != nil {
85 t.Fatal(err)
86 }
87
88 ctx := context.Background()
89 badRef := "../outside"
90
91 err = store.Put(ctx, badRef, strings.NewReader("bad"))
92 if err == nil {
93 t.Fatal("expected error putting file outside diskDir, got nil")
94 }
95
96 _, err = store.Open(ctx, badRef)
97 if err == nil {
98 t.Fatal("expected error opening file outside diskDir, got nil")
99 }
100}
101
102func TestS3Store(t *testing.T) {
103 mock := newMockS3Client()
104 store, err := NewS3Store(mock, "mybucket")
105 if err != nil {
106 t.Fatal(err)
107 }
108
109 ctx := context.Background()
110 ref := "logs/run1.log"
111 content := "s3 log payload"
112
113 if err := store.Put(ctx, ref, strings.NewReader(content)); err != nil {
114 t.Fatalf("Put to S3 failed: %v", err)
115 }
116
117 rc, err := store.Open(ctx, ref)
118 if err != nil {
119 t.Fatalf("Open from S3 failed: %v", err)
120 }
121 defer rc.Close()
122
123 got, err := io.ReadAll(rc)
124 if err != nil {
125 t.Fatalf("ReadAll failed: %v", err)
126 }
127 if string(got) != content {
128 t.Fatalf("got content %q, want %q", string(got), content)
129 }
130}