Be more pythonic in a test

firedrakeproject · Jan 4, 2018 · ec599b6 · ec599b6
1 parent f90e46b
commit ec599b6
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/test/unit/test_fiat.py b/test/unit/test_fiat.py
@@ -371,8 +371,7 @@ def test_nodality_tabulate(element):
 
     # Check nodality
     for j, x in enumerate(nodes_coords):
-        table = element.tabulate(0, (x,))
-        basis = table[list(table.keys())[0]]
+        basis, = element.tabulate(0, (x,)).values()
         for i in range(len(basis)):
             assert np.isclose(basis[i], 1.0 if i == j else 0.0)
 
@@ -414,14 +413,12 @@ def test_facet_nodality_tabulate(element):
 
             # Map dof coordinates to reference element due to
             # HdivTrace interface peculiarity
-            ref_coords = map_to_reference_facet((coords,), vertices, facet)
-            assert len(ref_coords) == 1
-            nodes_coords.append((facet, ref_coords[0]))
+            ref_coords, = map_to_reference_facet((coords,), vertices, facet)
+            nodes_coords.append((facet, ref_coords))
 
     # Check nodality
     for j, (facet, x) in enumerate(nodes_coords):
-        table = element.tabulate(0, (x,), entity=(facet_dim, facet))
-        basis = table[list(table.keys())[0]]
+        basis, = element.tabulate(0, (x,), entity=(facet_dim, facet)).values()
         for i in range(len(basis)):
             assert np.isclose(basis[i], 1.0 if i == j else 0.0)