diff --git a/tests/storage/test_matrix.py b/tests/storage/test_matrix.py
index d1c4f83..c709214 100644
--- a/tests/storage/test_matrix.py
+++ b/tests/storage/test_matrix.py
@@ -11,20 +11,22 @@ class TestMatrix:
         assert type(matrix) is stormpy.storage.SparseMatrix
         assert matrix.nr_rows == model.nr_states
         assert matrix.nr_columns == model.nr_states
-        assert matrix.nr_entries == 27 #model.nr_transitions
+        assert matrix.nr_entries == 20
+        assert matrix.nr_entries == model.nr_transitions
         for e in matrix:
             assert e.value() == 0.5 or e.value() == 0 or (e.value() == 1 and e.column > 6)
-    
+
     def test_backward_matrix(self):
         model = stormpy.parse_explicit_model(get_example_path("dtmc", "die.tra"), get_example_path("dtmc", "die.lab"))
         matrix = model.backward_transition_matrix
         assert type(matrix) is stormpy.storage.SparseMatrix
         assert matrix.nr_rows == model.nr_states
         assert matrix.nr_columns == model.nr_states
-        assert matrix.nr_entries == 20 #model.nr_transitions
+        assert matrix.nr_entries == 20
+        assert matrix.nr_entries == model.nr_transitions
         for e in matrix:
             assert e.value() == 0.5 or e.value() == 0 or (e.value() == 1 and e.column > 6)
-    
+
     def test_change_sparse_matrix(self):
         model = stormpy.parse_explicit_model(get_example_path("dtmc", "die.tra"), get_example_path("dtmc", "die.lab"))
         matrix = model.transition_matrix
@@ -38,7 +40,7 @@ class TestMatrix:
         for e in matrix:
             assert e.value() == i
             i += 0.1
-    
+
     def test_change_sparse_matrix_modelchecking(self):
         import stormpy.logic
         model = stormpy.parse_explicit_model(get_example_path("dtmc", "die.tra"), get_example_path("dtmc", "die.lab"))
@@ -51,7 +53,7 @@ class TestMatrix:
         result = stormpy.model_checking(model, formulas[0])
         resValue = result.at(model.initial_states[0])
         assert math.isclose(resValue, 0.16666666666666663)
-        
+
         # Change probabilities
         i = 0
         for e in matrix:
@@ -67,7 +69,7 @@ class TestMatrix:
         result = stormpy.model_checking(model, formulas[0])
         resValue = result.at(model.initial_states[0])
         assert math.isclose(resValue, 0.06923076923076932)
-        
+
         # Change probabilities again
         for state in model.states:
             for action in state.actions:
@@ -80,7 +82,7 @@ class TestMatrix:
         result = stormpy.model_checking(model, formulas[0])
         resValue = result.at(model.initial_states[0])
         assert math.isclose(resValue, 0.3555555555555556)
-    
+
     def test_change_parametric_sparse_matrix_modelchecking(self):
         import stormpy.logic
 
@@ -100,7 +102,7 @@ class TestMatrix:
         result = stormpy.model_checking(model, formulas[0])
         ratFunc = result.at(initial_state)
         assert len(ratFunc.gather_variables()) > 0
-        
+
         # Change probabilities
         two_pol = stormpy.RationalRF(2)
         two_pol = stormpy.FactorizedPolynomial(two_pol)
diff --git a/tests/storage/test_state.py b/tests/storage/test_state.py
index 25a1a10..56375a8 100644
--- a/tests/storage/test_state.py
+++ b/tests/storage/test_state.py
@@ -55,17 +55,17 @@ class TestState:
                 assert action.id == 0 or action.id == 1
 
     def test_transitions_dtmc(self):
-        transitions_orig = [(0, 0, 0), (0, 1, 0.5), (0, 2, 0.5), (1, 1, 0), (1, 3, 0.5), (1, 4, 0.5), 
-                            (2, 2, 0), (2, 5, 0.5), (2, 6, 0.5), (3, 1, 0.5), (3, 3, 0), (3, 7, 0.5),
-                            (4, 4, 0), (4, 8, 0.5), (4, 9, 0.5), (5, 5, 0), (5, 10, 0.5), (5, 11, 0.5),
-                            (6, 2, 0.5), (6, 6, 0), (6, 12, 0.5), (7, 7, 1), (8, 8, 1),
+        transitions_orig = [(0, 1, 0.5), (0, 2, 0.5), (1, 3, 0.5), (1, 4, 0.5),
+                            (2, 5, 0.5), (2, 6, 0.5), (3, 1, 0.5), (3, 7, 0.5),
+                            (4, 8, 0.5), (4, 9, 0.5), (5, 10, 0.5), (5, 11, 0.5),
+                            (6, 2, 0.5), (6, 12, 0.5), (7, 7, 1), (8, 8, 1),
                             (9, 9, 1), (10, 10, 1), (11, 11, 1), (12, 12, 1)
                            ]
         model = stormpy.parse_explicit_model(get_example_path("dtmc", "die.tra"), get_example_path("dtmc", "die.lab"))
         i = 0
         for state in model.states:
             for action in state.actions:
-                assert (state.id < 7 and len(action.transitions) == 3) or (state.id >= 7 and len(action.transitions) == 1)
+                assert (state.id < 7 and len(action.transitions) == 2) or (state.id >= 7 and len(action.transitions) == 1)
                 for transition in action.transitions:
                     transition_orig = transitions_orig[i]
                     i += 1
@@ -85,10 +85,10 @@ class TestState:
             assert i == 1 or i == 2
 
     def test_row_iterator(self):
-        transitions_orig = [(0, 0, 0), (0, 1, 0.5), (0, 2, 0.5), (1, 1, 0), (1, 3, 0.5), (1, 4, 0.5),
-                            (2, 2, 0), (2, 5, 0.5), (2, 6, 0.5), (3, 1, 0.5), (3, 3, 0), (3, 7, 0.5),
-                            (4, 4, 0), (4, 8, 0.5), (4, 9, 0.5), (5, 5, 0), (5, 10, 0.5), (5, 11, 0.5),
-                            (6, 2, 0.5), (6, 6, 0), (6, 12, 0.5), (7, 7, 1), (8, 8, 1),
+        transitions_orig = [(0, 1, 0.5), (0, 2, 0.5), (1, 3, 0.5), (1, 4, 0.5),
+                            (2, 5, 0.5), (2, 6, 0.5), (3, 1, 0.5), (3, 7, 0.5),
+                            (4, 8, 0.5), (4, 9, 0.5), (5, 10, 0.5), (5, 11, 0.5),
+                            (6, 2, 0.5), (6, 12, 0.5), (7, 7, 1), (8, 8, 1),
                             (9, 9, 1), (10, 10, 1), (11, 11, 1), (12, 12, 1)
                            ]
         model = stormpy.parse_explicit_model(get_example_path("dtmc", "die.tra"), get_example_path("dtmc", "die.lab"))