Test the new exceptions

4 files changed, 12 insertions, 11 deletions

diff --git a/src/etcd/tests/integration/test_election.py b/src/etcd/tests/integration/test_election.py
index 76f397b..46ea2bb 100644
--- a/src/etcd/tests/integration/test_election.py
+++ b/src/etcd/tests/integration/test_election.py
@@ -31,4 +31,4 @@ class TestElection(test_simple.EtcdIntegrationTest):
         e.set('/mysql', name='foo', ttl=1)
         time.sleep(2)
         self.assertRaises(etcd.EtcdException, e.get, '/mysql')
-        self.assertRaises(KeyError, e.delete, '/mysql', name='foo')
+        self.assertRaises(etcd.EtcdKeyNotFound, e.delete, '/mysql', name='foo')
diff --git a/src/etcd/tests/integration/test_simple.py b/src/etcd/tests/integration/test_simple.py
index 50530e9..185aad3 100644
--- a/src/etcd/tests/integration/test_simple.py
+++ b/src/etcd/tests/integration/test_simple.py
@@ -74,7 +74,7 @@ class TestSimple(EtcdIntegrationTest):
         try:
             get_result = self.client.get('/test_set')
             assert False
-        except KeyError as e:
+        except etcd.EtcdKeyNotFound as e:
             pass
 
         self.assertFalse('/test_set' in self.client)
@@ -100,7 +100,7 @@ class TestSimple(EtcdIntegrationTest):
         try:
             get_result = self.client.get('/test_set')
             assert False
-        except KeyError as e:
+        except etcd.EtcdKeyNotFound as e:
             pass
 
     def test_update(self):
@@ -140,7 +140,7 @@ class TestErrors(EtcdIntegrationTest):
         """ INTEGRATION: try to write  value to an existing directory """
 
         self.client.set('/directory/test-key', 'test-value')
-        self.assertRaises(KeyError, self.client.set, '/directory', 'test-value')
+        self.assertRaises(etcd.EtcdNotFile, self.client.set, '/directory', 'test-value')
 
     def test_test_and_set(self):
         """ INTEGRATION: try test_and_set operation """
@@ -159,7 +159,8 @@ class TestErrors(EtcdIntegrationTest):
         `prevExist=True` should fail """
         self.client.write('/mydir', None, dir=True)
 
-        self.assertRaises(KeyError, self.client.write, '/mydir', None, dir=True)
+        self.assertRaises(etcd.EtcdNotFile, self.client.write, '/mydir', None, dir=True)
+        self.assertRaises(etcd.EtcdAlreadyExist, self.client.write, '/mydir', None, dir=True, prevExist=False)
 
 
 class TestClusterFunctions(EtcdIntegrationTest):
diff --git a/src/etcd/tests/unit/test_old_request.py b/src/etcd/tests/unit/test_old_request.py
index e449b7d..9367ebd 100644
--- a/src/etcd/tests/unit/test_old_request.py
+++ b/src/etcd/tests/unit/test_old_request.py
@@ -171,7 +171,7 @@ class TestClientRequest(unittest.TestCase):
     def test_not_in(self):
         """ Can check if key is not in client """
         client = etcd.Client()
-        client.get = mock.Mock(side_effect=KeyError())
+        client.get = mock.Mock(side_effect=etcd.EtcdKeyNotFound())
         result = '/testkey' not in client
         self.assertEquals(True, result)
 
@@ -307,8 +307,8 @@ class TestClientApiExecutor(unittest.TestCase):
         try:
             client.api_execute('/v2/keys/testkey', client._MGET)
             assert False
-        except KeyError as e:
-            self.assertEquals(str(e), "'message : cause'")
+        except etcd.EtcdKeyNotFound as e:
+            self.assertEquals(str(e), 'message : cause')
 
     def test_put(self):
         """ http put request """
@@ -357,8 +357,8 @@ class TestClientApiExecutor(unittest.TestCase):
         try:
             client.api_execute('/v2/keys/testkey', client._MPUT, payload)
             self.fail()
-        except KeyError as e:
-            self.assertEquals("'message : cause'", str(e))
+        except etcd.EtcdNotFile as e:
+            self.assertEquals('message : cause', str(e))
 
     def test_get_error_unknown(self):
         """ http get error request unknown """
diff --git a/src/etcd/tests/unit/test_request.py b/src/etcd/tests/unit/test_request.py
index 7103fa1..df1ae57 100644
--- a/src/etcd/tests/unit/test_request.py
+++ b/src/etcd/tests/unit/test_request.py
@@ -282,7 +282,7 @@ class TestClientApiInterface(TestClientApiBase):
 
     def test_not_in(self):
         """ Can check if key is not in client """
-        self._mock_exception(KeyError, 'Key not Found : /testKey')
+        self._mock_exception(etcd.EtcdKeyNotFound, 'Key not Found : /testKey')
         self.assertTrue('/testey' not in self.client)
 
     def test_in(self):