Skip to content

Commit 2a2dbef

Browse files
leafoleafo
committed
add async and injected tests to extract handler tests
1 parent ade35cd commit 2a2dbef

File tree

1 file changed

+186
-0
lines changed

1 file changed

+186
-0
lines changed

‎zipserver/extract_handler_test.go‎

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,20 @@
11
package zipserver
22

33
import (
4+
"archive/zip"
5+
"bytes"
6+
"context"
7+
"encoding/json"
48
"fmt"
9+
"net/http"
10+
"net/http/httptest"
511
"net/url"
12+
"strings"
613
"testing"
14+
"time"
715

816
"github.com/stretchr/testify/assert"
17+
"github.com/stretchr/testify/require"
918
)
1019

1120
func Test_Limits(t *testing.T) {
@@ -76,3 +85,180 @@ func Test_LimitsWithHtmlFooter(t *testing.T) {
7685
assert.NoError(t, err)
7786
assert.EqualValues(t, "", el.HtmlFooter)
7887
}
88+
89+
// createTestZip creates an in-memory zip file with the given entries
90+
func createTestZip(t *testing.T, entries map[string][]byte) []byte {
91+
var buf bytes.Buffer
92+
zw := zip.NewWriter(&buf)
93+
94+
for name, data := range entries {
95+
writer, err := zw.CreateHeader(&zip.FileHeader{
96+
Name: name,
97+
UncompressedSize64: uint64(len(data)),
98+
})
99+
require.NoError(t, err)
100+
_, err = writer.Write(data)
101+
require.NoError(t, err)
102+
}
103+
104+
require.NoError(t, zw.Close())
105+
return buf.Bytes()
106+
}
107+
108+
func TestExtractHandler_AsyncCallback(t *testing.T) {
109+
withGoogleCloudStorage(t, func(storage Storage, config *Config) {
110+
targetName := "mem-target-extract-async"
111+
defer ClearNamedMemStorage(targetName)
112+
113+
config.ExtractionThreads = 4
114+
config.AsyncNotificationTimeout = Duration(10 * time.Second)
115+
config.StorageTargets = []StorageConfig{{
116+
Name: targetName,
117+
Type: Mem,
118+
Bucket: "target-bucket",
119+
}}
120+
121+
oldConfig := globalConfig
122+
globalConfig = config
123+
defer func() { globalConfig = oldConfig }()
124+
125+
ctx := context.Background()
126+
127+
zipData := createTestZip(t, map[string][]byte{
128+
"index.html": []byte("<html><body>Hello</body></html>"),
129+
"other.txt": []byte("other content"),
130+
})
131+
132+
// Upload zip to primary storage (GCS)
133+
_, err := storage.PutFile(ctx, config.Bucket, "zipserver_test/async_test.zip", bytes.NewReader(zipData), PutOptions{
134+
ContentType: "application/zip",
135+
})
136+
require.NoError(t, err)
137+
138+
// Set up callback server to capture the async notification
139+
callbackReceived := make(chan url.Values, 1)
140+
callbackServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
141+
require.NoError(t, r.ParseForm())
142+
callbackReceived <- r.Form
143+
w.WriteHeader(http.StatusOK)
144+
}))
145+
defer callbackServer.Close()
146+
147+
// Make request to extract handler with async callback
148+
form := url.Values{}
149+
form.Set("key", "zipserver_test/async_test.zip")
150+
form.Set("prefix", "zipserver_test/async_extracted")
151+
form.Set("target", targetName)
152+
form.Set("async", callbackServer.URL)
153+
form.Set("html_footer", "<script>injected</script>")
154+
155+
req := httptest.NewRequest(http.MethodPost, "/extract", strings.NewReader(form.Encode()))
156+
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
157+
w := httptest.NewRecorder()
158+
159+
err = extractHandler(w, req)
160+
require.NoError(t, err)
161+
162+
// Verify immediate response indicates async processing
163+
var immediateResp struct {
164+
Processing bool
165+
Async bool
166+
}
167+
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &immediateResp))
168+
assert.True(t, immediateResp.Processing)
169+
assert.True(t, immediateResp.Async)
170+
171+
// Wait for callback
172+
select {
173+
case callbackData := <-callbackReceived:
174+
assert.Equal(t, "true", callbackData.Get("Success"))
175+
176+
// Find the index.html entry and verify Injected is set
177+
foundIndex := false
178+
foundOther := false
179+
for i := 1; i <= 10; i++ {
180+
key := callbackData.Get(fmt.Sprintf("ExtractedFiles[%d][Key])", i))
181+
if key == "" {
182+
break
183+
}
184+
injected := callbackData.Get(fmt.Sprintf("ExtractedFiles[%d][Injected])", i))
185+
186+
if strings.HasSuffix(key, "/index.html") {
187+
foundIndex = true
188+
assert.Equal(t, "true", injected, "index.html should have Injected=true")
189+
} else if strings.HasSuffix(key, "/other.txt") {
190+
foundOther = true
191+
assert.Empty(t, injected, "other.txt should not have Injected field")
192+
}
193+
}
194+
assert.True(t, foundIndex, "should have found index.html in callback")
195+
assert.True(t, foundOther, "should have found other.txt in callback")
196+
197+
case <-time.After(10 * time.Second):
198+
t.Fatal("timeout waiting for async callback")
199+
}
200+
})
201+
}
202+
203+
func TestExtractHandler_SyncWithInjected(t *testing.T) {
204+
withGoogleCloudStorage(t, func(storage Storage, config *Config) {
205+
targetName := "mem-target-extract-sync"
206+
defer ClearNamedMemStorage(targetName)
207+
208+
config.ExtractionThreads = 4
209+
config.StorageTargets = []StorageConfig{{
210+
Name: targetName,
211+
Type: Mem,
212+
Bucket: "target-bucket",
213+
}}
214+
215+
oldConfig := globalConfig
216+
globalConfig = config
217+
defer func() { globalConfig = oldConfig }()
218+
219+
ctx := context.Background()
220+
221+
zipData := createTestZip(t, map[string][]byte{
222+
"index.html": []byte("<html><body>Hello</body></html>"),
223+
"other.txt": []byte("other content"),
224+
})
225+
226+
// Upload zip to primary storage (GCS)
227+
_, err := storage.PutFile(ctx, config.Bucket, "zipserver_test/sync_test.zip", bytes.NewReader(zipData), PutOptions{
228+
ContentType: "application/zip",
229+
})
230+
require.NoError(t, err)
231+
232+
// Make sync request (no async parameter)
233+
form := url.Values{}
234+
form.Set("key", "zipserver_test/sync_test.zip")
235+
form.Set("prefix", "zipserver_test/sync_extracted")
236+
form.Set("target", targetName)
237+
form.Set("html_footer", "<script>injected</script>")
238+
239+
req := httptest.NewRequest(http.MethodPost, "/extract", strings.NewReader(form.Encode()))
240+
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
241+
w := httptest.NewRecorder()
242+
243+
err = extractHandler(w, req)
244+
require.NoError(t, err)
245+
246+
// Parse sync response
247+
var resp struct {
248+
Success bool
249+
ExtractedFiles []ExtractedFile
250+
}
251+
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
252+
assert.True(t, resp.Success)
253+
assert.Len(t, resp.ExtractedFiles, 2)
254+
255+
// Verify Injected field
256+
for _, f := range resp.ExtractedFiles {
257+
if strings.HasSuffix(f.Key, "/index.html") {
258+
assert.True(t, f.Injected, "index.html should have Injected=true")
259+
} else {
260+
assert.False(t, f.Injected, "other.txt should have Injected=false")
261+
}
262+
}
263+
})
264+
}

0 commit comments

Comments
 (0)