diff --git a/src/runtime/python/pypgf.c b/src/runtime/python/pypgf.c
index 939ad44ff..b40747978 100644
--- a/src/runtime/python/pypgf.c
+++ b/src/runtime/python/pypgf.c
@@ -295,7 +295,12 @@ PGF_categoryProbability(PGFObject *self, PyObject *args)
         return NULL;
     }
 
-    return PyFloat_FromDouble((double)prob);
+    double dprob = (double) prob;
+    if (dprob == INFINITY) {
+        PyErr_Format(PyExc_KeyError, "category '%s' is not defined", s);
+        return NULL;
+    }
+    return PyFloat_FromDouble(dprob);
 }
 
 static PyObject *
@@ -315,7 +320,12 @@ PGF_functionProbability(PGFObject *self, PyObject *args)
         return NULL;
     }
 
-    return PyFloat_FromDouble((double)prob);
+    double dprob = (double) prob;
+    if (dprob == INFINITY) {
+        PyErr_Format(PyExc_KeyError, "function '%s' is not defined", s);
+        return NULL;
+    }
+    return PyFloat_FromDouble(dprob);
 }
 
 static PyObject *
@@ -331,7 +341,8 @@ PGF_exprProbability(PGFObject *self, PyObject *args)
         return NULL;
     }
 
-    return PyFloat_FromDouble((double)prob);
+    double dprob = (double) prob;
+    return PyFloat_FromDouble(dprob);
 }
 
 static PyGetSetDef PGF_getseters[] = {
diff --git a/src/runtime/python/tests/test_transactions.py b/src/runtime/python/tests/test_transactions.py
index 00cfaab00..f2c4a6182 100644
--- a/src/runtime/python/tests/test_transactions.py
+++ b/src/runtime/python/tests/test_transactions.py
@@ -66,19 +66,18 @@ def test_original_function_type(gr1):
         gr1.functionType("foo")
 
 def test_original_function_prob(gr1):
-    # with pytest.raises(KeyError):
-    #     gr1.functionProbability("foo")
-    assert gr1.functionProbability("foo") == float('inf')
+    with pytest.raises(KeyError):
+        gr1.functionProbability("foo")
+    # assert gr1.functionProbability("foo") == float('inf')
 
-@pytest.mark.skip(reason="failing")
 def test_original_category_prob(gr1):
-    # with pytest.raises(KeyError):
-    #     gr1.categoryProbability("Q")
-    assert gr1.categoryProbability("Q") == float('inf')
+    with pytest.raises(KeyError):
+        gr1.categoryProbability("Q")
+    # assert gr1.categoryProbability("Q") == float('inf')
 
 def test_original_expr_prob(gr1):
     # with pytest.raises(KeyError):
-    #     gr1.functionProbability("foo")
+    #     gr1.exprProbability(ExprFun("foo"))
     assert gr1.exprProbability(ExprFun("foo")) == float('inf')
 
 # gr2