swh/model/tests/test_hypothesis_strategies.py

@@ -20,19 +20,12 @@ from swh.model.hypothesis_strategies import (
     skipped_contents,
     snapshots,
 )
-from swh.model.model import TargetType
+from swh.model.model import ModelObjectType, SnapshotTargetType
 
 target_types = ("content", "directory", "revision", "release", "snapshot", "alias")
-all_but_skipped_content = (
-    "origin",
-    "origin_visit",
-    "origin_visit_status",
-    "snapshot",
-    "release",
-    "revision",
-    "directory",
-    "content",
-)
+all_but_skipped_content = {
+    o_t for o_t in ModelObjectType if o_t != ModelObjectType.SKIPPED_CONTENT
+}
 
 
 @given(objects(blacklist_types=()))
@@ -44,25 +37,35 @@ def test_generation(obj_type_and_obj):
 @given(objects(split_content=False))
 def test_generation_merged_content(obj_type_and_obj):
     # we should never generate a "skipped_content" here
-    assert obj_type_and_obj[0] != "skipped_content"
+    assert obj_type_and_obj[0] != ModelObjectType.SKIPPED_CONTENT
 
 
 @given(objects(split_content=True, blacklist_types=all_but_skipped_content))
 def test_generation_split_content(obj_type_and_obj):
     # we should only generate "skipped_content"
-    assert obj_type_and_obj[0] == "skipped_content"
+    assert obj_type_and_obj[0] == ModelObjectType.SKIPPED_CONTENT
 
 
-@given(objects(blacklist_types=("origin_visit", "directory")))
+@given(
+    objects(
+        blacklist_types={
+            ModelObjectType.DIRECTORY,
+            ModelObjectType.ORIGIN_VISIT,
+        }
+    )
+)
 def test_generation_blacklist(obj_type_and_obj):
-    assert obj_type_and_obj[0] not in ("origin_visit", "directory")
+    assert obj_type_and_obj[0] not in {
+        ModelObjectType.DIRECTORY,
+        ModelObjectType.ORIGIN_VISIT,
+    }
 
 
 def assert_nested_dict(obj):
     """Tests the object is a nested dict and contains no more class
     from swh.model.model."""
     if isinstance(obj, dict):
-        for (key, value) in obj.items():
+        for key, value in obj.items():
             assert isinstance(key, (str, bytes)), key
             assert_nested_dict(value)
     elif isinstance(obj, tuple):
@@ -78,7 +81,7 @@ def assert_nested_dict(obj):
 def test_dicts_generation(obj_type_and_obj):
     (obj_type, object_) = obj_type_and_obj
     assert_nested_dict(object_)
-    if obj_type == "content":
+    if obj_type == ModelObjectType.CONTENT:
         COMMON_KEYS = set(DEFAULT_ALGORITHMS) | {"length", "status", "ctime"}
         if object_["status"] == "visible":
             assert set(object_) <= COMMON_KEYS | {"data"}
@@ -88,9 +91,9 @@ def test_dicts_generation(obj_type_and_obj):
             assert set(object_) <= COMMON_KEYS | {"data"}
         else:
             assert False, object_
-    elif obj_type == "release":
+    elif obj_type == ModelObjectType.RELEASE:
         assert object_["target_type"] in target_types
-    elif obj_type == "snapshot":
+    elif obj_type == ModelObjectType.SNAPSHOT:
         for branch in object_["branches"].values():
             assert branch is None or branch["target_type"] in target_types
 
@@ -106,18 +109,28 @@ def test_datetimes(dt):
 @given(object_dicts(split_content=False))
 def test_dicts_generation_merged_content(obj_type_and_obj):
     # we should never generate a "skipped_content" here
-    assert obj_type_and_obj[0] != "skipped_content"
+    assert obj_type_and_obj[0] != ModelObjectType.SKIPPED_CONTENT
 
 
 @given(object_dicts(split_content=True, blacklist_types=all_but_skipped_content))
 def test_dicts_generation_split_content(obj_type_and_obj):
     # we should only generate "skipped_content"
-    assert obj_type_and_obj[0] == "skipped_content"
+    assert obj_type_and_obj[0] == ModelObjectType.SKIPPED_CONTENT
 
 
-@given(object_dicts(blacklist_types=("release", "content")))
+@given(
+    object_dicts(
+        blacklist_types={
+            ModelObjectType.CONTENT,
+            ModelObjectType.RELEASE,
+        }
+    )
+)
 def test_dicts_generation_blacklist(obj_type_and_obj):
-    assert obj_type_and_obj[0] not in ("release", "content")
+    assert obj_type_and_obj[0] not in {
+        ModelObjectType.CONTENT,
+        ModelObjectType.RELEASE,
+    }
 
 
 @given(objects())
@@ -126,7 +139,7 @@ def test_model_to_dicts(obj_type_and_obj):
     object_type = object_.object_type
     obj_dict = object_.to_dict()
     assert_nested_dict(obj_dict)
-    if object_type in ("content", "skipped_content"):
+    if object_type in {ModelObjectType.CONTENT, ModelObjectType.SKIPPED_CONTENT}:
         COMMON_KEYS = set(DEFAULT_ALGORITHMS) | {"length", "status"}
         if object_.ctime is not None:
             COMMON_KEYS |= {"ctime"}
@@ -139,9 +152,9 @@ def test_model_to_dicts(obj_type_and_obj):
             assert set(obj_dict) == COMMON_KEYS | {"data"}
         else:
             assert False, obj_dict
-    elif object_type == "release":
+    elif object_type == ModelObjectType.RELEASE:
         assert obj_dict["target_type"] in target_types
-    elif object_type == "snapshot":
+    elif object_type == ModelObjectType.RELEASE:
         for branch in obj_dict["branches"].values():
             assert branch is None or branch["target_type"] in target_types
 
@@ -163,7 +176,6 @@ _max_snp_size = 100
 @given(snapshots(min_size=_min_snp_size, max_size=_max_snp_size))
 @settings(max_examples=1)
 def test_snapshots_strategy(snapshot):
-
     branches = snapshot.branches
 
     assert len(branches) >= _min_snp_size
@@ -174,7 +186,7 @@ def test_snapshots_strategy(snapshot):
     # check snapshot integrity
     for name, branch in branches.items():
         assert branch is None or branch.target_type.value in target_types
-        if branch is not None and branch.target_type == TargetType.ALIAS:
+        if branch is not None and branch.target_type == SnapshotTargetType.ALIAS:
             aliases.append(name)
             assert branch.target in branches
 
@@ -184,7 +196,7 @@ def test_snapshots_strategy(snapshot):
         current_alias = alias
         while (
             branches[current_alias] is not None
-            and branches[current_alias].target_type == TargetType.ALIAS
+            and branches[current_alias].target_type == SnapshotTargetType.ALIAS
         ):
             assert branches[current_alias].target not in processed_alias
             processed_alias.add(current_alias)