@@ -14,8 +14,21 @@ import (
14
14
15
15
"github.com/mozilla-services/stubattribution/attributioncode"
16
16
"github.com/mozilla-services/stubattribution/stubservice/backends"
17
+ "github.com/sirupsen/logrus"
18
+ "github.com/sirupsen/logrus/hooks/test"
19
+ "golang.org/x/exp/slices"
17
20
)
18
21
22
+ // testHook is a logrus hook that is registered globally in order to capture
23
+ // the messages logged during the execution of the tests. This is useful to
24
+ // verify some important log statements. We have to register a global hook
25
+ // because we use the global logrus instance in our code.
26
+ var testHook * test.Hook
27
+
28
+ func init () {
29
+ testHook = test .NewGlobal ()
30
+ }
31
+
19
32
func TestUniqueKey (t * testing.T ) {
20
33
f := func (url , code string ) bool {
21
34
key := uniqueKey (url , code )
@@ -137,15 +150,95 @@ func TestRedirectFull(t *testing.T) {
137
150
138
151
svc := NewStubService (
139
152
NewRedirectHandler (storage , server .URL + "/cdn/" , "" ),
140
- & attributioncode.Validator {})
153
+ & attributioncode.Validator {},
154
+ )
155
+
156
+ for _ , params := range []struct {
157
+ AttributionCode string
158
+ Referer string
159
+ ExpectedLocation string
160
+ ExpectedCode string
161
+ SkipDownloadLogChecks bool
162
+ ExpectedClientID string
163
+ ExpectedSessionID string
164
+ }{
165
+ {
166
+ AttributionCode : `campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=www.google.com` ,
167
+ Referer : "" ,
168
+ ExpectedLocation : `/cdn/builds/firefox-stub/en-US/win/` ,
169
+ ExpectedCode : `campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3Dwww.google.com` ,
170
+ },
171
+ {
172
+ AttributionCode : `campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=www.notinwhitelist.com` ,
173
+ Referer : "" ,
174
+ ExpectedLocation : `/cdn/builds/firefox-stub/en-US/win/` ,
175
+ ExpectedCode : `campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3D%2528other%2529` ,
176
+ },
177
+ {
178
+ // We expect the product to be prefixed in the location URL below because
179
+ // the attribution code contains data for RTAMO and the referer header
180
+ // contains the right value.
181
+ AttributionCode : `campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org` ,
182
+ Referer : `https://www.mozilla.org/` ,
183
+ ExpectedLocation : `/cdn/builds/rtamo-firefox-stub/en-US/win/` ,
184
+ ExpectedCode : `campaign%3Dfxa-cta-123%26content%3Drta%253Avalue%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org` ,
185
+ },
186
+ {
187
+ // Same as before with a `client_id`.
188
+ AttributionCode : `campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org&client_id=some-client-id` ,
189
+ Referer : `https://www.mozilla.org/` ,
190
+ ExpectedLocation : `/cdn/builds/rtamo-firefox-stub/en-US/win/` ,
191
+ ExpectedCode : `campaign%3Dfxa-cta-123%26content%3Drta%253Avalue%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org` ,
192
+ ExpectedClientID : "some-client-id" ,
193
+ },
194
+ {
195
+ // Same as before with a `visit_id`.
196
+ AttributionCode : `campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org&visit_id=some-visit-id` ,
197
+ Referer : `https://www.mozilla.org/` ,
198
+ ExpectedLocation : `/cdn/builds/rtamo-firefox-stub/en-US/win/` ,
199
+ ExpectedCode : `campaign%3Dfxa-cta-123%26content%3Drta%253Avalue%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org` ,
200
+ ExpectedClientID : "some-visit-id" ,
201
+ },
202
+ {
203
+ // Same as before with the addition of a `session_id`.
204
+ AttributionCode : `campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org&visit_id=some-visit-id&session_id=some-session-id` ,
205
+ Referer : `https://www.mozilla.org/` ,
206
+ ExpectedLocation : `/cdn/builds/rtamo-firefox-stub/en-US/win/` ,
207
+ ExpectedCode : `campaign%3Dfxa-cta-123%26content%3Drta%253Avalue%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org` ,
208
+ ExpectedClientID : "some-visit-id" ,
209
+ ExpectedSessionID : "some-session-id" ,
210
+ },
211
+ {
212
+ // We expect no prefix because the attribution data is not related to
213
+ // RTAMO.
214
+ AttributionCode : `campaign=some-campaign&content=not-for-rtamo&medium=referral&source=addons.mozilla.org` ,
215
+ Referer : `https://www.mozilla.org/` ,
216
+ ExpectedLocation : `/cdn/builds/firefox-stub/en-US/win/` ,
217
+ ExpectedCode : `campaign%3Dsome-campaign%26content%3Dnot-for-rtamo%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org` ,
218
+ },
219
+ {
220
+ // This should not return a modified installer because the referer is not
221
+ // the expected one.
222
+ AttributionCode : `campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org` ,
223
+ Referer : `https://example.org/` ,
224
+ ExpectedLocation : `\?lang=en-US&os=win&product=firefox-stub` ,
225
+ ExpectedCode : "" ,
226
+ SkipDownloadLogChecks : true ,
227
+ },
228
+ } {
229
+ testHook .Reset ()
230
+
231
+ expectedLocationRegexp := regexp .MustCompile (params .ExpectedLocation )
232
+ expectedCodeRegexp := regexp .MustCompile (params .ExpectedCode )
141
233
142
- runTest := func (attributionCode , referer string , expectedLocation string , expectedCode string ) {
143
- expectedCodeRegexp := regexp .MustCompile (expectedCode )
144
- expectedLocationRegexp := regexp .MustCompile (expectedLocation )
145
234
recorder := httptest .NewRecorder ()
146
- base64Code := base64 .URLEncoding .WithPadding ('.' ).EncodeToString ([]byte (attributionCode ))
147
- req := httptest .NewRequest ("GET" , `http://test/?product=firefox-stub&os=win&lang=en-US&attribution_code=` + url .QueryEscape (base64Code ), nil )
148
- req .Header .Set ("Referer" , referer )
235
+ base64Code := base64 .URLEncoding .WithPadding ('.' ).EncodeToString ([]byte (params .AttributionCode ))
236
+ req := httptest .NewRequest (
237
+ "GET" ,
238
+ `http://test/?product=firefox-stub&os=win&lang=en-US&attribution_code=` + url .QueryEscape (base64Code ),
239
+ nil ,
240
+ )
241
+ req .Header .Set ("Referer" , params .Referer )
149
242
svc .ServeHTTP (recorder , req )
150
243
151
244
location := recorder .HeaderMap .Get ("Location" )
@@ -169,53 +262,48 @@ func TestRedirectFull(t *testing.T) {
169
262
if err != nil {
170
263
t .Fatal ("could not read body" , err )
171
264
}
172
-
173
265
if len (bodyBytes ) != len (testFileBytes ) {
174
266
t .Error ("Returned file was not the same length as the original file" )
175
267
}
176
-
177
268
if ! expectedCodeRegexp .Match (bodyBytes ) {
178
269
t .Error ("Returned file did not contain attribution code" )
179
270
}
180
- }
181
271
182
- emptyReferer := ""
183
- runTest (
184
- `campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=www.google.com` ,
185
- emptyReferer ,
186
- `/cdn/builds/firefox-stub/en-US/win/` ,
187
- `campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3Dwww.google.com` ,
188
- )
189
- runTest (
190
- `campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=www.notinwhitelist.com` ,
191
- emptyReferer ,
192
- `/cdn/builds/firefox-stub/en-US/win/` ,
193
- `campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3D%2528other%2529` ,
194
- )
195
- // We expect the product to be prefixed in the location URL below because the
196
- // attribution code contains data for RTAMO and the referer header contains
197
- // the right value.
198
- runTest (
199
- `campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org` ,
200
- `https://www.mozilla.org/` ,
201
- `/cdn/builds/rtamo-firefox-stub/en-US/win/` ,
202
- `campaign%3Dfxa-cta-123%26content%3Drta%253Avalue%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org` ,
203
- )
204
- // We expect no prefix because the attribution data is not related to RTAMO.
205
- runTest (
206
- `campaign=some-campaign&content=not-for-rtamo&medium=referral&source=addons.mozilla.org` ,
207
- `https://www.mozilla.org/` ,
208
- `/cdn/builds/firefox-stub/en-US/win/` ,
209
- `campaign%3Dsome-campaign%26content%3Dnot-for-rtamo%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org` ,
210
- )
211
- // This should not return a modified installer because the referer is not the
212
- // expected one.
213
- runTest (
214
- `campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org` ,
215
- `https://example.org/` ,
216
- `\?lang=en-US&os=win&product=firefox-stub` ,
217
- `` ,
218
- )
272
+ if ! params .SkipDownloadLogChecks {
273
+ entries := testHook .AllEntries ()
274
+
275
+ idx := slices .IndexFunc (entries , func (e * logrus.Entry ) bool {
276
+ return e .Message == "Download Started"
277
+ })
278
+ if idx == - 1 {
279
+ t .Error ("Could not find Download Started log entry" )
280
+ }
281
+ downloadStarted := entries [idx ]
282
+
283
+ // This one should always be the last log entry.
284
+ downloadFinished := entries [len (entries )- 1 ]
285
+ if downloadFinished .Message != "Download Finished" {
286
+ t .Errorf ("Unexpected log message: %s" , downloadFinished .Message )
287
+ }
288
+
289
+ for _ , entry := range []* logrus.Entry {downloadStarted , downloadFinished } {
290
+ clientID := entry .Data ["client_id" ]
291
+ if clientID != params .ExpectedClientID {
292
+ t .Errorf ("Expected client_id: %s, got: %v" , params .ExpectedClientID , clientID )
293
+ }
294
+
295
+ visitID := entry .Data ["visit_id" ]
296
+ if visitID != params .ExpectedClientID {
297
+ t .Errorf ("Expected visit_id: %s, got: %v" , params .ExpectedClientID , visitID )
298
+ }
299
+
300
+ sessionID := entry .Data ["session_id" ]
301
+ if sessionID != params .ExpectedSessionID {
302
+ t .Errorf ("Expected session_id: %s, got: %v" , params .ExpectedSessionID , sessionID )
303
+ }
304
+ }
305
+ }
306
+ }
219
307
}
220
308
221
309
func TestDirectFull (t * testing.T ) {
@@ -245,14 +333,33 @@ func TestDirectFull(t *testing.T) {
245
333
246
334
svc := NewStubService (
247
335
NewDirectHandler (),
248
- & attributioncode.Validator {})
336
+ & attributioncode.Validator {},
337
+ )
249
338
250
- runTest := func (attributionCode , expectedCode string ) {
251
- expectedCodeRegexp := regexp .MustCompile (expectedCode )
252
- base64Code := base64 .URLEncoding .WithPadding ('.' ).EncodeToString ([]byte (attributionCode ))
253
- req := httptest .NewRequest ("GET" , `http://test/?product=firefox-stub&os=win&lang=en-US&attribution_code=` + url .QueryEscape (base64Code ), nil )
339
+ for _ , params := range []struct {
340
+ AttributionCode string
341
+ ExpectedCode string
342
+ }{
343
+ {
344
+ AttributionCode : `campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=www.google.com` ,
345
+ ExpectedCode : `campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3Dwww.google.com` ,
346
+ },
347
+ {
348
+ AttributionCode : `campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=notinthewhitelist.com` ,
349
+ ExpectedCode : `campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3D%2528other%2529` ,
350
+ },
351
+ } {
352
+ testHook .Reset ()
353
+
354
+ expectedCodeRegexp := regexp .MustCompile (params .ExpectedCode )
254
355
255
356
recorder := httptest .NewRecorder ()
357
+ base64Code := base64 .URLEncoding .WithPadding ('.' ).EncodeToString ([]byte (params .AttributionCode ))
358
+ req := httptest .NewRequest (
359
+ "GET" ,
360
+ `http://test/?product=firefox-stub&os=win&lang=en-US&attribution_code=` + url .QueryEscape (base64Code ),
361
+ nil ,
362
+ )
256
363
svc .ServeHTTP (recorder , req )
257
364
258
365
if recorder .Code != 200 {
@@ -263,27 +370,14 @@ func TestDirectFull(t *testing.T) {
263
370
if err != nil {
264
371
t .Fatal ("could not read body" , err )
265
372
}
266
-
267
373
if len (bodyBytes ) != len (testFileBytes ) {
268
374
t .Error ("Returned file was not the same length as the original file" )
269
375
}
270
376
271
377
if ! expectedCodeRegexp .Match (bodyBytes ) {
272
378
t .Error ("Returned file did not contain attribution code" )
273
379
}
274
- //if !bytes.Contains(bodyBytes, []byte(url.QueryEscape(expectedCode))) {
275
- //t.Error("Returned file did not contain attribution code")
276
- //}
277
380
}
278
-
279
- runTest (
280
- `campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=www.google.com` ,
281
- `campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3Dwww.google.com` ,
282
- )
283
- runTest (
284
- `campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=notinthewhitelist.com` ,
285
- `campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3D%2528other%2529` ,
286
- )
287
381
}
288
382
289
383
func TestStubServiceErrorCases (t * testing.T ) {
0 commit comments