X-Git-Url: https://osm.etsi.org/gitweb/?a=blobdiff_plain;f=tests%2Funit%2Ftest_model.py;h=222d881a6b5c16090fe6f3517e045009c9b72dc4;hb=c50c361a8b9a3bbf1a33f5659e492b481f065cd2;hp=67db5ae9e13bc2e0705a45c90e64cf11f320fcf6;hpb=a09dd5038d226fbcda54a79bdddc995054857198;p=osm%2FN2VC.git

diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py
index 67db5ae..222d881 100644
--- a/tests/unit/test_model.py
+++ b/tests/unit/test_model.py
@@ -1,6 +1,7 @@
 import unittest
 
 import mock
+import asynctest
 
 
 def _make_delta(entity, type_, data=None):
@@ -113,3 +114,42 @@ def test_get_series():
     assert model._get_series('~foo/ubuntu', entity) == 'xenial'
     assert model._get_series('ubuntu', entity) == 'xenial'
     assert model._get_series('cs:ubuntu', entity) == 'xenial'
+
+
+class TestContextManager(asynctest.TestCase):
+    @asynctest.patch('juju.model.Model.disconnect')
+    @asynctest.patch('juju.model.Model.connect_current')
+    async def test_normal_use(self, mock_connect, mock_disconnect):
+        from juju.model import Model
+
+        async with Model() as model:
+            self.assertTrue(isinstance(model, Model))
+
+        self.assertTrue(mock_connect.called)
+        self.assertTrue(mock_disconnect.called)
+
+    @asynctest.patch('juju.model.Model.disconnect')
+    @asynctest.patch('juju.model.Model.connect_current')
+    async def test_exception(self, mock_connect, mock_disconnect):
+        from juju.model import Model
+
+        class SomeException(Exception):
+            pass
+
+        with self.assertRaises(SomeException):
+            async with Model():
+                raise SomeException()
+
+        self.assertTrue(mock_connect.called)
+        self.assertTrue(mock_disconnect.called)
+
+    @asynctest.patch('juju.client.connection.JujuData.current_controller')
+    async def test_no_current_connection(self, mock_current_controller):
+        from juju.model import Model
+        from juju.errors import JujuConnectionError
+
+        mock_current_controller.return_value = ""
+
+        with self.assertRaises(JujuConnectionError):
+            async with Model():
+                pass