Skip to content

Commit 26cc885

Browse files
committed
Fixes unit tests
1 parent c51a6b0 commit 26cc885

File tree

2 files changed

+9
-11
lines changed

2 files changed

+9
-11
lines changed

ads/common/extended_enum.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,10 @@ def __contains__(cls, item: object) -> bool:
4343
# For non-string items (e.g., int), do a direct membership check
4444
return item in attr_values
4545

46+
def __iter__(cls):
47+
# Make the class iterable by returning an iterator over its values
48+
return iter(cls.values())
49+
4650
def values(cls) -> tuple:
4751
"""
4852
Gets the tuple of class attribute values, excluding private or special

ads/model/datascience_model.py

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -144,10 +144,7 @@ def from_dict(cls, data: Dict) -> "ModelBackupSetting":
144144
return cls(
145145
is_backup_enabled=data.get("is_backup_enabled"),
146146
backup_region=data.get("backup_region"),
147-
customer_notification_type=CustomerNotificationType(
148-
data.get("customer_notification_type")
149-
)
150-
or None,
147+
customer_notification_type=data.get("customer_notification_type") or None,
151148
)
152149

153150
def to_json(self) -> str:
@@ -227,10 +224,7 @@ def from_dict(cls, data: Dict) -> "ModelRetentionSetting":
227224
return cls(
228225
archive_after_days=data.get("archive_after_days"),
229226
delete_after_days=data.get("delete_after_days"),
230-
customer_notification_type=CustomerNotificationType(
231-
data.get("customer_notification_type")
232-
)
233-
or None,
227+
customer_notification_type=data.get("customer_notification_type") or None,
234228
)
235229

236230
def to_json(self) -> str:
@@ -322,9 +316,9 @@ def to_dict(self) -> Dict:
322316
def from_dict(cls, data: Dict) -> "ModelRetentionOperationDetails":
323317
"""Constructs retention operation details from a dictionary."""
324318
return cls(
325-
archive_state=SettingStatus(data.get("archive_state")) or None,
319+
archive_state=data.get("archive_state") or None,
326320
archive_state_details=data.get("archive_state_details"),
327-
delete_state=SettingStatus(data.get("delete_state")) or None,
321+
delete_state=data.get("delete_state") or None,
328322
delete_state_details=data.get("delete_state_details"),
329323
time_archival_scheduled=data.get("time_archival_scheduled"),
330324
time_deletion_scheduled=data.get("time_deletion_scheduled"),
@@ -405,7 +399,7 @@ def to_dict(self) -> Dict:
405399
def from_dict(cls, data: Dict) -> "ModelBackupOperationDetails":
406400
"""Constructs backup operation details from a dictionary."""
407401
return cls(
408-
backup_state=SettingStatus(data.get("backup_state")) or None,
402+
backup_state=data.get("backup_state") or None,
409403
backup_state_details=data.get("backup_state_details"),
410404
time_last_backup=data.get("time_last_backup"),
411405
)

0 commit comments

Comments
 (0)