99from splitgraph .exceptions import (
1010 IncompleteObjectDownloadError ,
1111 IncompleteObjectUploadError ,
12+ ObjectCacheError ,
1213)
1314from splitgraph .hooks .s3 import S3ExternalObjectHandler
1415from splitgraph .hooks .s3_server import (
@@ -108,6 +109,10 @@ def test_s3_push_pull(
108109 )
109110
110111
112+ class FlakyHandlerError (Exception ):
113+ pass
114+
115+
111116def _flaky_handler (incomplete = False ):
112117 class FlakyExternalObjectHandler (S3ExternalObjectHandler ):
113118 """
@@ -124,7 +129,7 @@ def download_objects(self, objects, remote_engine):
124129 Downloads just the first object and then fails.
125130 """
126131 super ().download_objects (objects [:1 ], remote_engine )
127- ex = Exception ("Something bad happened." )
132+ ex = FlakyHandlerError ("Something bad happened." )
128133 if incomplete :
129134 raise IncompleteObjectDownloadError (reason = ex , successful_objects = [objects [0 ][0 ]])
130135 else :
@@ -135,7 +140,7 @@ def upload_objects(self, objects, remote_engine):
135140 Uploads just the first object and then fails.
136141 """
137142 super ().upload_objects (objects [:1 ], remote_engine )
138- ex = Exception ("Something bad happened." )
143+ ex = FlakyHandlerError ("Something bad happened." )
139144 if incomplete :
140145 raise IncompleteObjectUploadError (
141146 reason = ex , successful_objects = objects [:1 ], successful_object_urls = objects [:1 ]
@@ -163,7 +168,7 @@ def test_push_upload_error(
163168 "splitgraph.hooks.external_objects._EXTERNAL_OBJECT_HANDLERS" ,
164169 {"S3" : _flaky_handler (incomplete = interrupted )},
165170 ):
166- with pytest .raises (Exception ):
171+ with pytest .raises (FlakyHandlerError ):
167172 PG_MNT .push (remote_repository = unprivileged_pg_repo , handler = "S3" , handler_options = {})
168173
169174 assert head not in unprivileged_pg_repo .images
@@ -232,7 +237,7 @@ def test_pull_download_error(local_engine_empty, unprivileged_pg_repo, clean_min
232237 "splitgraph.hooks.external_objects._EXTERNAL_OBJECT_HANDLERS" ,
233238 {"S3" : _flaky_handler (interrupted )},
234239 ):
235- with pytest .raises (Exception ):
240+ with pytest .raises (FlakyHandlerError if interrupted else ObjectCacheError ):
236241 clone (unprivileged_pg_repo , local_repository = PG_MNT , download_all = True )
237242
238243 # Check that the pull succeeded (repository registered locally) but the objects
0 commit comments