Skip to content

Commit fe4f127

Browse files
cutecatfanncutecatfann
committed
Fixing test override
1 parent 88f828f commit fe4f127

File tree

4 files changed

+48
-35
lines changed

4 files changed

+48
-35
lines changed

‎pkg/workingset/dcr_cleanup_test.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ func TestCleanupOrphanedDCREntries_DeletesOrphanedAndRevoked(t *testing.T) {
8888
mock := newMockDCRClient("server-a", "server-b")
8989

9090
// Remove server-b (not in any profile, not authorized) → should be deleted
91-
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-b"}, nil)
91+
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-b"})
9292

9393
assert.Equal(t, []string{"server-b"}, mock.deleted)
9494
assert.True(t, mock.registered["server-a"], "server-a should not be deleted")
@@ -111,7 +111,7 @@ func TestCleanupOrphanedDCREntries_SkipsAuthorized(t *testing.T) {
111111
mock := newMockDCRClient("server-a").withAuthorized("server-a")
112112

113113
// server-a is not in any profile but IS authorized → should NOT be deleted
114-
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-a"}, nil)
114+
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-a"})
115115

116116
assert.Empty(t, mock.deleted)
117117
assert.True(t, mock.registered["server-a"])
@@ -142,7 +142,7 @@ func TestCleanupOrphanedDCREntries_SkipsInUse(t *testing.T) {
142142
mock := newMockDCRClient("server-a")
143143

144144
// server-a is still in profile-2, should NOT be deleted
145-
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-a"}, nil)
145+
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-a"})
146146

147147
assert.Empty(t, mock.deleted)
148148
assert.True(t, mock.registered["server-a"])
@@ -165,7 +165,7 @@ func TestCleanupOrphanedDCREntries_SkipsNoDCREntry(t *testing.T) {
165165
mock := newMockDCRClient()
166166

167167
// server-x has no DCR entry → GetDCRClient returns error → skip delete
168-
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-x"}, nil)
168+
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-x"})
169169

170170
assert.Empty(t, mock.deleted)
171171
}
@@ -194,7 +194,7 @@ func TestCleanupOrphanedDCREntries_MultipleServers(t *testing.T) {
194194
mock := newMockDCRClient("server-a", "server-b", "server-c").withAuthorized("server-c")
195195

196196
// Remove server-a (still in profile), server-b (revoked), server-c (authorized)
197-
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-a", "server-b", "server-c"}, nil)
197+
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-a", "server-b", "server-c"})
198198

199199
// Only server-b should be deleted (server-a in profile, server-c authorized)
200200
assert.Equal(t, []string{"server-b"}, mock.deleted)
@@ -235,7 +235,7 @@ func TestCleanupOrphanedDCREntries_ServerInDifferentProfile(t *testing.T) {
235235

236236
// Cleanup after removing server-a and server-b from profile-1
237237
// server-b is still in profile-2, so only server-a should be deleted
238-
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-a", "server-b"}, nil)
238+
doCleanupOrphanedDCREntries(ctx, dao, mock, []string{"server-a", "server-b"})
239239

240240
assert.Equal(t, []string{"server-a"}, mock.deleted)
241241
assert.True(t, mock.registered["server-b"])

‎pkg/workingset/oauth.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
// Desktop registration functions. Tests can override these to avoid
1313
// requiring a running Docker Desktop backend.
1414
var (
15+
isCEModeFunc = oauth.IsCEMode
1516
registerWithSnapshotFunc = oauth.RegisterProviderWithSnapshot
1617
registerForDynamicDiscoveryFunc = oauth.RegisterProviderForDynamicDiscovery
1718
)
@@ -26,7 +27,7 @@ var (
2627
// happens during the authorize command instead.
2728
func RegisterOAuthProvidersForServers(ctx context.Context, servers []Server) {
2829
// CE mode: all OAuth is Gateway-owned, skip Desktop registration entirely.
29-
if oauth.IsCEMode() {
30+
if isCEModeFunc() {
3031
return
3132
}
3233

‎pkg/workingset/oauth_test.go‎

Lines changed: 25 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,28 @@ import (
1010
"github.com/docker/mcp-gateway/pkg/oauth"
1111
)
1212

13+
// mockDesktopMode overrides the package-level function vars so the test
14+
// runs as if Docker Desktop is present. Call the returned cleanup func
15+
// (or use t.Cleanup) to restore the originals.
16+
func mockDesktopMode(t *testing.T) {
17+
t.Helper()
18+
oldCE := isCEModeFunc
19+
oldSnapshot := registerWithSnapshotFunc
20+
oldDiscovery := registerForDynamicDiscoveryFunc
21+
22+
isCEModeFunc = func() bool { return false }
23+
24+
t.Cleanup(func() {
25+
isCEModeFunc = oldCE
26+
registerWithSnapshotFunc = oldSnapshot
27+
registerForDynamicDiscoveryFunc = oldDiscovery
28+
})
29+
}
30+
1331
func TestRegisterOAuthProvidersForServers_CEModeSkipsAll(t *testing.T) {
14-
t.Setenv("DOCKER_MCP_USE_CE", "true")
32+
oldCE := isCEModeFunc
33+
isCEModeFunc = func() bool { return true }
34+
t.Cleanup(func() { isCEModeFunc = oldCE })
1535

1636
// Track registration calls — none should happen in CE mode.
1737
var snapshotCalls, discoveryCalls int
@@ -60,7 +80,7 @@ func TestRegisterOAuthProvidersForServers_CEModeSkipsAll(t *testing.T) {
6080
}
6181

6282
func TestRegisterOAuthProvidersForServers_NilSnapshotSkipped(t *testing.T) {
63-
t.Setenv("DOCKER_MCP_USE_CE", "false")
83+
mockDesktopMode(t)
6484

6585
var snapshotCalls int
6686
registerWithSnapshotFunc = func(_ context.Context, _, _ string) error {
@@ -70,10 +90,6 @@ func TestRegisterOAuthProvidersForServers_NilSnapshotSkipped(t *testing.T) {
7090
registerForDynamicDiscoveryFunc = func(_ context.Context, _, _ string) error {
7191
return nil
7292
}
73-
t.Cleanup(func() {
74-
registerWithSnapshotFunc = oauth.RegisterProviderWithSnapshot
75-
registerForDynamicDiscoveryFunc = oauth.RegisterProviderForDynamicDiscovery
76-
})
7793

7894
servers := []Server{
7995
{Type: ServerTypeRemote, Snapshot: nil},
@@ -85,7 +101,7 @@ func TestRegisterOAuthProvidersForServers_NilSnapshotSkipped(t *testing.T) {
85101
}
86102

87103
func TestRegisterOAuthProvidersForServers_ExplicitOAuthRegistered(t *testing.T) {
88-
t.Setenv("DOCKER_MCP_USE_CE", "false")
104+
mockDesktopMode(t)
89105

90106
var registeredServers []string
91107
registerWithSnapshotFunc = func(_ context.Context, serverName, _ string) error {
@@ -95,10 +111,6 @@ func TestRegisterOAuthProvidersForServers_ExplicitOAuthRegistered(t *testing.T)
95111
registerForDynamicDiscoveryFunc = func(_ context.Context, _, _ string) error {
96112
return nil
97113
}
98-
t.Cleanup(func() {
99-
registerWithSnapshotFunc = oauth.RegisterProviderWithSnapshot
100-
registerForDynamicDiscoveryFunc = oauth.RegisterProviderForDynamicDiscovery
101-
})
102114

103115
servers := []Server{
104116
{
@@ -121,7 +133,7 @@ func TestRegisterOAuthProvidersForServers_ExplicitOAuthRegistered(t *testing.T)
121133
}
122134

123135
func TestRegisterOAuthProvidersForServers_DynamicDiscovery(t *testing.T) {
124-
t.Setenv("DOCKER_MCP_USE_CE", "false")
136+
mockDesktopMode(t)
125137

126138
var discoveredServers []string
127139
registerWithSnapshotFunc = func(_ context.Context, _, _ string) error {
@@ -131,10 +143,6 @@ func TestRegisterOAuthProvidersForServers_DynamicDiscovery(t *testing.T) {
131143
discoveredServers = append(discoveredServers, serverName)
132144
return nil
133145
}
134-
t.Cleanup(func() {
135-
registerWithSnapshotFunc = oauth.RegisterProviderWithSnapshot
136-
registerForDynamicDiscoveryFunc = oauth.RegisterProviderForDynamicDiscovery
137-
})
138146

139147
servers := []Server{
140148
{
@@ -155,7 +163,7 @@ func TestRegisterOAuthProvidersForServers_DynamicDiscovery(t *testing.T) {
155163
}
156164

157165
func TestRegisterOAuthProvidersForServers_NonRemoteSkipped(t *testing.T) {
158-
t.Setenv("DOCKER_MCP_USE_CE", "false")
166+
mockDesktopMode(t)
159167

160168
var snapshotCalls, discoveryCalls int
161169
registerWithSnapshotFunc = func(_ context.Context, _, _ string) error {
@@ -166,10 +174,6 @@ func TestRegisterOAuthProvidersForServers_NonRemoteSkipped(t *testing.T) {
166174
discoveryCalls++
167175
return nil
168176
}
169-
t.Cleanup(func() {
170-
registerWithSnapshotFunc = oauth.RegisterProviderWithSnapshot
171-
registerForDynamicDiscoveryFunc = oauth.RegisterProviderForDynamicDiscovery
172-
})
173177

174178
servers := []Server{
175179
{

‎pkg/workingset/server.go‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -183,10 +183,23 @@ func CleanupOrphanedDCREntries(ctx context.Context, dao db.DAO, serverNames []st
183183
if oauth.IsCEMode() {
184184
return
185185
}
186-
doCleanupOrphanedDCREntries(ctx, dao, desktop.NewAuthClient(), serverNames, communityServers)
186+
187+
// Filter out servers where Gateway owns OAuth — their DCR entries
188+
// are not managed by Desktop, so there is nothing to clean up.
189+
filtered := make([]string, 0, len(serverNames))
190+
for _, name := range serverNames {
191+
if !oauth.ShouldUseGatewayOAuth(ctx, communityServers[name]) {
192+
filtered = append(filtered, name)
193+
}
194+
}
195+
if len(filtered) == 0 {
196+
return
197+
}
198+
199+
doCleanupOrphanedDCREntries(ctx, dao, desktop.NewAuthClient(), filtered)
187200
}
188201

189-
func doCleanupOrphanedDCREntries(ctx context.Context, dao db.DAO, client dcrClient, serverNames []string, communityServers map[string]bool) {
202+
func doCleanupOrphanedDCREntries(ctx context.Context, dao db.DAO, client dcrClient, serverNames []string) {
190203
allSets, err := dao.ListWorkingSets(ctx)
191204
if err != nil {
192205
log.Logf("Warning: Failed to list profiles for DCR cleanup: %v", err)
@@ -204,11 +217,6 @@ func doCleanupOrphanedDCREntries(ctx context.Context, dao db.DAO, client dcrClie
204217
}
205218

206219
for _, name := range serverNames {
207-
// Skip servers where Gateway owns OAuth — their DCR entries are
208-
// not managed by Desktop, so there is nothing to clean up.
209-
if oauth.ShouldUseGatewayOAuth(ctx, communityServers[name]) {
210-
continue
211-
}
212220
if inUse[name] {
213221
continue
214222
}

0 commit comments

Comments
 (0)