@@ -164,3 +164,191 @@ func TestPipelinePushFailWithExtraRequirements(t *testing.T) {
164164 err = PipelinePush (t .Context (), "r8.im/user/test" , dir , apiClient , client , cfg )
165165 require .Error (t , err )
166166}
167+
168+ func TestPipelinePushSuccessWithBetaPatch (t * testing.T ) {
169+ // Setup mock web server for cog.replicate.com (token exchange)
170+ webServer := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
171+ switch r .URL .Path {
172+ case "/api/token/user" :
173+ // Mock token exchange response
174+ //nolint:gosec
175+ tokenResponse := `{
176+ "keys": {
177+ "cog": {
178+ "key": "test-api-token",
179+ "expires_at": "2024-12-31T23:59:59Z"
180+ }
181+ }
182+ }`
183+ w .WriteHeader (http .StatusOK )
184+ w .Write ([]byte (tokenResponse ))
185+ default :
186+ w .WriteHeader (http .StatusNotFound )
187+ }
188+ }))
189+ defer webServer .Close ()
190+
191+ // Setup mock API server for api.replicate.com (version and release endpoints)
192+ apiServer := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
193+ switch r .URL .Path {
194+ case "/v1/models/user/test/versions" :
195+ // Mock version creation response
196+ versionResponse := `{"id": "test-version-id"}`
197+ w .WriteHeader (http .StatusCreated )
198+ w .Write ([]byte (versionResponse ))
199+ case "/v1/models/user/test/releases" :
200+ // Mock release creation response - empty body with 204 status
201+ w .WriteHeader (http .StatusNoContent )
202+ default :
203+ w .WriteHeader (http .StatusNotFound )
204+ }
205+ }))
206+ defer apiServer .Close ()
207+
208+ cdnServer := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
209+ switch r .URL .Path {
210+ case "/requirements.txt" :
211+ // Mock requirements.txt response
212+ requirementsResponse := "mycustompackage==1.1.0b2"
213+ w .Header ().Add (EtagHeader , "a" )
214+ w .WriteHeader (http .StatusOK )
215+ w .Write ([]byte (requirementsResponse ))
216+ default :
217+ w .WriteHeader (http .StatusNotFound )
218+ }
219+ }))
220+ defer cdnServer .Close ()
221+
222+ webURL , err := url .Parse (webServer .URL )
223+ require .NoError (t , err )
224+ apiURL , err := url .Parse (apiServer .URL )
225+ require .NoError (t , err )
226+ cdnURL , err := url .Parse (cdnServer .URL )
227+ require .NoError (t , err )
228+
229+ t .Setenv (env .SchemeEnvVarName , webURL .Scheme )
230+ t .Setenv (env .WebHostEnvVarName , webURL .Host )
231+ t .Setenv (env .APIHostEnvVarName , apiURL .Host )
232+ t .Setenv (env .PipelinesRuntimeHostEnvVarName , cdnURL .Host )
233+
234+ dir := t .TempDir ()
235+
236+ // Create mock predict
237+ predictPyPath := filepath .Join (dir , "predict.py" )
238+ handle , err := os .Create (predictPyPath )
239+ require .NoError (t , err )
240+ handle .WriteString ("import cog" )
241+ handle .Close ()
242+ dockertest .MockCogConfig = "{\" build\" :{\" python_version\" :\" 3.12\" ,\" python_packages\" :[\" torch==2.5.0\" ,\" mycustompackage==1.1.0b2\" ],\" system_packages\" :[\" git\" ]},\" image\" :\" test\" ,\" predict\" :\" " + predictPyPath + ":Predictor\" }"
243+
244+ // Setup mock command
245+ command := dockertest .NewMockCommand ()
246+ client , err := cogHttp .ProvideHTTPClient (t .Context (), command )
247+ require .NoError (t , err )
248+ webClient := web .NewClient (command , client )
249+ apiClient := api .NewClient (command , client , webClient )
250+
251+ cfg := config .DefaultConfig ()
252+ requirementsPath := filepath .Join (dir , "requirements.txt" )
253+ handle , err = os .Create (requirementsPath )
254+ require .NoError (t , err )
255+ handle .WriteString ("mycustompackage==1.1.0b2" )
256+ handle .Close ()
257+ cfg .Build .PythonRequirements = filepath .Base (requirementsPath )
258+ err = PipelinePush (t .Context (), "r8.im/user/test" , dir , apiClient , client , cfg )
259+ require .NoError (t , err )
260+ }
261+
262+ func TestPipelinePushSuccessWithAlphaPatch (t * testing.T ) {
263+ // Setup mock web server for cog.replicate.com (token exchange)
264+ webServer := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
265+ switch r .URL .Path {
266+ case "/api/token/user" :
267+ // Mock token exchange response
268+ //nolint:gosec
269+ tokenResponse := `{
270+ "keys": {
271+ "cog": {
272+ "key": "test-api-token",
273+ "expires_at": "2024-12-31T23:59:59Z"
274+ }
275+ }
276+ }`
277+ w .WriteHeader (http .StatusOK )
278+ w .Write ([]byte (tokenResponse ))
279+ default :
280+ w .WriteHeader (http .StatusNotFound )
281+ }
282+ }))
283+ defer webServer .Close ()
284+
285+ // Setup mock API server for api.replicate.com (version and release endpoints)
286+ apiServer := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
287+ switch r .URL .Path {
288+ case "/v1/models/user/test/versions" :
289+ // Mock version creation response
290+ versionResponse := `{"id": "test-version-id"}`
291+ w .WriteHeader (http .StatusCreated )
292+ w .Write ([]byte (versionResponse ))
293+ case "/v1/models/user/test/releases" :
294+ // Mock release creation response - empty body with 204 status
295+ w .WriteHeader (http .StatusNoContent )
296+ default :
297+ w .WriteHeader (http .StatusNotFound )
298+ }
299+ }))
300+ defer apiServer .Close ()
301+
302+ cdnServer := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
303+ switch r .URL .Path {
304+ case "/requirements.txt" :
305+ // Mock requirements.txt response
306+ requirementsResponse := "mycustompackage==1.1.0b2"
307+ w .Header ().Add (EtagHeader , "a" )
308+ w .WriteHeader (http .StatusOK )
309+ w .Write ([]byte (requirementsResponse ))
310+ default :
311+ w .WriteHeader (http .StatusNotFound )
312+ }
313+ }))
314+ defer cdnServer .Close ()
315+
316+ webURL , err := url .Parse (webServer .URL )
317+ require .NoError (t , err )
318+ apiURL , err := url .Parse (apiServer .URL )
319+ require .NoError (t , err )
320+ cdnURL , err := url .Parse (cdnServer .URL )
321+ require .NoError (t , err )
322+
323+ t .Setenv (env .SchemeEnvVarName , webURL .Scheme )
324+ t .Setenv (env .WebHostEnvVarName , webURL .Host )
325+ t .Setenv (env .APIHostEnvVarName , apiURL .Host )
326+ t .Setenv (env .PipelinesRuntimeHostEnvVarName , cdnURL .Host )
327+
328+ dir := t .TempDir ()
329+
330+ // Create mock predict
331+ predictPyPath := filepath .Join (dir , "predict.py" )
332+ handle , err := os .Create (predictPyPath )
333+ require .NoError (t , err )
334+ handle .WriteString ("import cog" )
335+ handle .Close ()
336+ dockertest .MockCogConfig = "{\" build\" :{\" python_version\" :\" 3.12\" ,\" python_packages\" :[\" torch==2.5.0\" ,\" mycustompackage>=1.0\" ],\" system_packages\" :[\" git\" ]},\" image\" :\" test\" ,\" predict\" :\" " + predictPyPath + ":Predictor\" }"
337+
338+ // Setup mock command
339+ command := dockertest .NewMockCommand ()
340+ client , err := cogHttp .ProvideHTTPClient (t .Context (), command )
341+ require .NoError (t , err )
342+ webClient := web .NewClient (command , client )
343+ apiClient := api .NewClient (command , client , webClient )
344+
345+ cfg := config .DefaultConfig ()
346+ requirementsPath := filepath .Join (dir , "requirements.txt" )
347+ handle , err = os .Create (requirementsPath )
348+ require .NoError (t , err )
349+ handle .WriteString ("mycustompackage>=1.0" )
350+ handle .Close ()
351+ cfg .Build .PythonRequirements = filepath .Base (requirementsPath )
352+ err = PipelinePush (t .Context (), "r8.im/user/test" , dir , apiClient , client , cfg )
353+ require .NoError (t , err )
354+ }
0 commit comments