@@ -2,14 +2,17 @@ package files_test
22
33import (
44"context"
5+ "sync"
56"sync/atomic"
67"testing"
78"time"
89
910"github.com/google/uuid"
1011"github.com/prometheus/client_golang/prometheus"
1112"github.com/spf13/afero"
13+ "github.com/stretchr/testify/assert"
1214"github.com/stretchr/testify/require"
15+ "go.uber.org/mock/gomock"
1316"golang.org/x/sync/errgroup"
1417
1518"cdr.dev/slog/sloggers/slogtest"
@@ -18,13 +21,70 @@ import (
1821"github.com/coder/coder/v2/coderd/database"
1922"github.com/coder/coder/v2/coderd/database/dbauthz"
2023"github.com/coder/coder/v2/coderd/database/dbgen"
24+ "github.com/coder/coder/v2/coderd/database/dbmock"
2125"github.com/coder/coder/v2/coderd/database/dbtestutil"
2226"github.com/coder/coder/v2/coderd/files"
2327"github.com/coder/coder/v2/coderd/rbac"
2428"github.com/coder/coder/v2/coderd/rbac/policy"
2529"github.com/coder/coder/v2/testutil"
2630)
2731
32+ // TestCancelledFetch runs 2 Acquire calls. The first fails with a ctx.Canceled
33+ // error. The second call should ignore the first error and try to fetch the file
34+ // again, which should succeed.
35+ func TestCancelledFetch (t * testing.T ) {
36+ t .Parallel ()
37+
38+ fileID := uuid .New ()
39+ rdy := make (chan struct {})
40+ dbM := dbmock .NewMockStore (gomock .NewController (t ))
41+
42+ // First call should fail
43+ dbM .EXPECT ().GetFileByID (gomock .Any (),gomock .Any ()).DoAndReturn (func (mTx context.Context ,fileID uuid.UUID ) (database.File ,error ) {
44+ // Wait long enough for the second call to be queued up.
45+ <- rdy
46+ return database.File {},context .Canceled
47+ })
48+
49+ // Second call should succeed
50+ dbM .EXPECT ().GetFileByID (gomock .Any (),gomock .Any ()).DoAndReturn (func (mTx context.Context ,fileID uuid.UUID ) (database.File ,error ) {
51+ return database.File {
52+ ID :fileID ,
53+ Data :make ([]byte ,100 ),
54+ },nil
55+ })
56+
57+ //nolint:gocritic // Unit testing
58+ ctx := dbauthz .AsFileReader (testutil .Context (t ,testutil .WaitShort ))
59+ cache := files .NewFromStore (dbM ,prometheus .NewRegistry (),& coderdtest.FakeAuthorizer {})
60+
61+ var wg sync.WaitGroup
62+ wg .Add (2 )
63+
64+ // First call that will fail
65+ go func () {
66+ _ ,err := cache .Acquire (ctx ,fileID )
67+ assert .ErrorIs (t ,err ,context .Canceled )
68+ wg .Done ()
69+ }()
70+
71+ // Second call, that should succeed
72+ go func () {
73+ fs ,err := cache .Acquire (ctx ,fileID )
74+ assert .NoError (t ,err )
75+ if fs != nil {
76+ fs .Close ()
77+ }
78+ wg .Done ()
79+ }()
80+
81+ // We need that second Acquire call to be queued up
82+ time .Sleep (testutil .IntervalFast )
83+
84+ close (rdy )
85+ wg .Wait ()
86+ }
87+
2888// nolint:paralleltest,tparallel // Serially testing is easier
2989func TestCacheRBAC (t * testing.T ) {
3090t .Parallel ()