@@ -43,8 +43,16 @@ class ImageType(enum.Enum):
4343class ImageInfo :
4444 """Dataclass to hold image information."""
4545
46- id : str
47- name : str
46+ image_id : str
47+ image_name : str
48+
49+ def __eq__ (self , value : object ) -> bool :
50+ """
51+ Check if two ImageInfo objects are equal.
52+
53+ It only checks if the image_id is the same, since this should be the unique identifier.
54+ """
55+ return isinstance (value , ImageInfo ) and self .image_id == value .image_id
4856
4957
5058class BaseCloud (ABC ):
@@ -225,15 +233,15 @@ def clean(self) -> List[Exception]:
225233 exceptions .append (e )
226234 for image_info in self .created_images :
227235 try :
228- self .delete_image (image_id = image_info .id )
236+ self .delete_image (image_id = image_info .image_id )
229237 except Exception as e :
230238 exceptions .append (e )
231239 for image_info in self .preserved_images :
232240 # noop - just log that we're not cleaning up these images
233241 self ._log .info (
234242 "Preserved image %s [id:%s] is NOT being cleaned up." ,
235- image_info .name ,
236- image_info .id ,
243+ image_info .image_name ,
244+ image_info .image_id ,
237245 )
238246 return exceptions
239247
@@ -329,6 +337,18 @@ def _validate_tag(tag: str):
329337 raise InvalidTagNameError (tag = tag , rules_failed = rules_failed )
330338
331339 def _get_ssh_keys (self ) -> KeyPair :
340+ """
341+ Get the ssh key pair to use for the cloud instance.
342+
343+ If no key pair is provided in the config file, the default key pair
344+ will be used. The default key pair is the id_rsa or id_ed25519 key
345+ in the user's .ssh directory.
346+
347+ :raises PycloudlibError: if no public key path is provided and no default key is found
348+ :raises PycloudlibError: if the public key path provided in the config does not exist
349+
350+ :return: KeyPair object with the public and private key paths
351+ """
332352 user = getpass .getuser ()
333353 # check if id_rsa or id_ed25519 keys exist in the user's .ssh directory
334354 possible_default_keys = [
@@ -380,21 +400,48 @@ def _store_snapshot_info(
380400 :return: ImageInfo object with the snapshot information
381401 """
382402 image_info = ImageInfo (
383- id = snapshot_id ,
384- name = snapshot_name ,
403+ image_id = snapshot_id ,
404+ image_name = snapshot_name ,
385405 )
386406 if not keep_snapshot :
387407 self .created_images .append (image_info )
388408 self ._log .info (
389- "Created temporary snapshot %s [id:%s]" ,
390- image_info .name ,
391- image_info .id ,
409+ "Created temporary snapshot %s" ,
410+ image_info ,
392411 )
393412 else :
394413 self .preserved_images .append (image_info )
395414 self ._log .info (
396- "Created permanent snapshot %s [id:%s]" ,
397- image_info .name ,
398- image_info .id ,
415+ "Created permanent snapshot %s" ,
416+ image_info ,
399417 )
400418 return image_info
419+
420+ def _record_image_deletion (self , image_id : str ):
421+ """
422+ Record the deletion of an image.
423+
424+ This method should be called after an image is successfully deleted.
425+ It will remove the image from the list of created_images or preserved_images
426+ so that the cloud does not attempt to re-clean it up later. It will also log
427+ the deletion of the image.
428+
429+ :param image_id: ID of the image that was deleted
430+ """
431+ if match := [i for i in self .created_images if i .image_id == image_id ]:
432+ deleted_image = match [0 ]
433+ self .created_images .remove (deleted_image )
434+ self ._log .debug (
435+ "Snapshot %s has been deleted. Will no longer need to be cleaned up later." ,
436+ deleted_image ,
437+ )
438+ elif match := [i for i in self .preserved_images if i .image_id == image_id ]:
439+ deleted_image = match [0 ]
440+ self .preserved_images .remove (deleted_image )
441+ self ._log .debug (
442+ "Snapshot %s has been deleted. This snapshot was taken with keep=True, "
443+ "but since it has been manually deleted, it will not be preserved." ,
444+ deleted_image ,
445+ )
446+ else :
447+ self ._log .debug ("Deleted image %s" , image_id )
0 commit comments