diff --git a/tests/__init__.py b/tests/__init__.py index 1dbba0836..e5937e75d 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -109,4 +109,42 @@ def cleanup_patch_test(model: Base, payload: dict, original_data: Base) -> None: session.commit() +def retrieve_latest_polymorphic_table_record( + target_record: Base, + polymorphic_relationship: str, + polymorphic_type: str, +) -> Base: + """ + Retrieve the latest record from a polymorphic table. This function assumes that the + parent class has the correct mixin to support retrieval via an attribute. This + requires end_date to be None + + Parameters: + ---------- + target_record : Base + The parent record from which to retrieve the polymorphic child record. + polymorphic_relationship : str + The name of the relationship attribute on the parent record that corresponds to the polymorphic table. + polymorphic_type : str + The specific type of the polymorphic record to retrieve (e.g., 'Use Status' or 'Monitoring Status' for StatusHistory). + latest : bool, optional + If True, retrieves the latest record based on start_date. Defaults to True. + """ + if polymorphic_relationship == "permissions": + type_field = "permission_type" + elif polymorphic_relationship == "status_history": + type_field = "status_type" + + polymorphic_records = getattr(target_record, polymorphic_relationship) + type_polymorphic_records = [ + r + for r in polymorphic_records + if getattr(r, type_field) == polymorphic_type and r.end_date is None + ] + sorted_type_polymorphic_records = sorted( + type_polymorphic_records, key=lambda r: r.start_date, reverse=True + ) + return sorted_type_polymorphic_records[0] + + # ============= EOF =============================================