diff --git a/aiida/backends/tests/work/work_chain.py b/aiida/backends/tests/work/work_chain.py
index 8756820eb5..88bf7aead2 100644
--- a/aiida/backends/tests/work/work_chain.py
+++ b/aiida/backends/tests/work/work_chain.py
@@ -114,7 +114,7 @@ def test_dict(self):
         self.assertEqual(wc.ctx['new_attr'], 5)
 
         del wc.ctx['new_attr']
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(KeyError):
             wc.ctx['new_attr']