simplify ugrid test code

test/test_ugrid.py

@@ -45,32 +45,6 @@ def test_generic_io():
     helpers.generic_io("test.0.lb8.ugrid")
 
 
-def first(iterable, condition=lambda x: True):
-    """
-    Returns the first item in the `iterable` that
-    satisfies the `condition`.
-
-    If the condition is not given, returns the first item of
-    the iterable.
-
-    returns None if no item satysfing the condition is found.
-
-    original version : https://stackoverflow.com/a/35513376
-
-    >>> first( (1,2,3), condition=lambda x: x % 2 == 0)
-    2
-    >>> first(range(3, 100))
-    3
-    >>> first( () ) is None
-    True
-    """
-
-    try:
-        return next(x for x in iterable if condition(x))
-    except StopIteration:
-        return None
-
-
 # sphere_mixed.1.lb8.ugrid and hch_strct.4.lb8.ugrid created
 # using the codes from http://cfdbooks.com
 @pytest.mark.skipif(sys.version_info < (3, 6), reason="requires Python 3.6 or higher")
@@ -107,40 +81,49 @@ def test_reference_file(
     assert mesh.points.shape[0] == ref_num_points
     assert mesh.points.shape[1] == 3
 
+    ugrid_meshio_id = {
+        "triangle": None,
+        "quad": None,
+        "tetra": None,
+        "pyramid": None,
+        "wedge": None,
+        "hexahedron": None,
+    }
+
+    for i, (key, data) in enumerate(mesh.cells):
+        if key in ugrid_meshio_id:
+            ugrid_meshio_id[key] = i
+
     # validate element counts
     if ref_num_triangle > 0:
-        c = first(mesh.cells, lambda c: c.type == "triangle")
-        assert c is not None
+        c = mesh.cells[ugrid_meshio_id["triangle"]]
         assert c.data.shape == (ref_num_triangle, 3)
     else:
-        assert first(mesh.cells, lambda c: c.type == "triangle") is None
+        assert ugrid_meshio_id["triangle"] is None
 
     if ref_num_quad > 0:
-        c = first(mesh.cells, lambda c: c.type == "quad")
-        assert c is not None
+        c = mesh.cells[ugrid_meshio_id["quad"]]
         assert c.data.shape == (ref_num_quad, 4)
     else:
-        assert first(mesh.cells, lambda c: c.type == "quad") is None
+        assert ugrid_meshio_id["quad"] is None
 
     if ref_num_tet > 0:
-        assert mesh.cells[1].type == "tetra"
+        c = mesh.cells[ugrid_meshio_id["tetra"]]
         assert mesh.cells[1].data.shape == (ref_num_tet, 4)
     else:
-        assert first(mesh.cells, lambda c: c.type == "tetra") is None
+        assert ugrid_meshio_id["tetra"] is None
 
     if ref_num_wedge > 0:
-        c = first(mesh.cells, lambda c: c.type == "wedge")
-        assert c is not None
+        c = mesh.cells[ugrid_meshio_id["wedge"]]
         assert c.data.shape == (ref_num_wedge, 6)
     else:
-        assert first(mesh.cells, lambda c: c.type == "wedge") is None
+        assert ugrid_meshio_id["wedge"] is None
 
     if ref_num_hex > 0:
-        c = first(mesh.cells, lambda c: c.type == "hexahedron")
-        assert c is not None
+        c = mesh.cells[ugrid_meshio_id["hexahedron"]]
         assert c.data.shape == (ref_num_hex, 8)
     else:
-        first(mesh.cells, lambda c: c.type == "hexahedron") is None
+        assert ugrid_meshio_id["hexahedron"] is None
 
     # validate boundary tags
 
@@ -239,18 +222,28 @@ def test_area(filename, area_tria_ref, area_quad_ref, accuracy):
     filename = os.path.join(this_dir, "meshes", "ugrid", filename)
 
     mesh = meshio.read(filename)
-
-    tria = first(mesh.cells, lambda c: c.type == "triangle")
-    assert tria is not None
+    ugrid_meshio_id = {
+        "triangle": None,
+        "quad": None,
+        "tetra": None,
+        "pyramid": None,
+        "wedge": None,
+        "hexahedron": None,
+    }
+
+    for i, (key, data) in enumerate(mesh.cells):
+        if key in ugrid_meshio_id:
+            ugrid_meshio_id[key] = i
+
+    tria = mesh.cells[ugrid_meshio_id["triangle"]]
     total_tri_area = 0
     for _cell in tria.data:
         cell = numpy.array([mesh.points[i] for i in _cell])
         a = _triangle_area(cell)
         total_tri_area += a
     assert numpy.isclose(total_tri_area, area_tria_ref, accuracy)
 
-    quad = first(mesh.cells, lambda c: c.type == "quad")
-    assert quad is not None
+    quad = mesh.cells[ugrid_meshio_id["quad"]]
     total_quad_area = 0
     for _cell in quad.data:
         cell = numpy.array([mesh.points[i] for i in _cell])