1 files changed, 5 insertions, 25 deletions
diff --git a/tests/contrib/test_sqlalchemy.py b/tests/contrib/test_sqlalchemy.py
index 068aa92..421f516 100644
--- a/tests/contrib/test_sqlalchemy.py
+++ b/tests/contrib/test_sqlalchemy.py
@@ -13,12 +13,11 @@
 # limitations under the License.
 
 import datetime
-import unittest
 
-import mock
 import sqlalchemy
 import sqlalchemy.ext.declarative
 import sqlalchemy.orm
+import unittest2
 
 import oauth2client
 import oauth2client.client
@@ -37,7 +36,7 @@ class DummyModel(Base):
         oauth2client.contrib.sqlalchemy.CredentialsType)
 
 
-class TestSQLAlchemyStorage(unittest.TestCase):
+class TestSQLAlchemyStorage(unittest2.TestCase):
     def setUp(self):
         engine = sqlalchemy.create_engine('sqlite://')
         Base.metadata.create_all(engine)
@@ -67,8 +66,7 @@ class TestSQLAlchemyStorage(unittest.TestCase):
         self.assertEqual(result.token_uri, self.credentials.token_uri)
         self.assertEqual(result.user_agent, self.credentials.user_agent)
 
-    @mock.patch('oauth2client.client.OAuth2Credentials.set_store')
-    def test_get(self, set_store):
+    def test_get(self):
         session = self.session()
         credentials_storage = oauth2client.contrib.sqlalchemy.Storage(
             session=session,
@@ -77,21 +75,7 @@ class TestSQLAlchemyStorage(unittest.TestCase):
             key_value=1,
             property_name='credentials',
         )
-        # No credentials stored
         self.assertIsNone(credentials_storage.get())
-
-        # Invalid credentials stored
-        session.add(DummyModel(
-            key=1,
-            credentials=oauth2client.client.Credentials(),
-        ))
-        session.commit()
-        bad_credentials = credentials_storage.get()
-        self.assertIsInstance(bad_credentials, oauth2client.client.Credentials)
-        set_store.assert_not_called()
-
-        # Valid credentials stored
-        session.query(DummyModel).filter_by(key=1).delete()
         session.add(DummyModel(
             key=1,
             credentials=self.credentials,
@@ -99,20 +83,16 @@ class TestSQLAlchemyStorage(unittest.TestCase):
         session.commit()
 
         self.compare_credentials(credentials_storage.get())
-        set_store.assert_called_with(credentials_storage)
 
     def test_put(self):
         session = self.session()
-        storage = oauth2client.contrib.sqlalchemy.Storage(
+        oauth2client.contrib.sqlalchemy.Storage(
             session=session,
             model_class=DummyModel,
             key_name='key',
             key_value=1,
             property_name='credentials',
-        )
-        # Store invalid credentials first to verify overwriting
-        storage.put(oauth2client.client.Credentials())
-        storage.put(self.credentials)
+        ).put(self.credentials)
         session.commit()
 
         entity = session.query(DummyModel).filter_by(key=1).first()