from cartoframes.data import Dataset
from cartoframes.data.clients.auth_api_client import AuthAPIClient
from cartoframes.auth import Credentials
from cartoframes.lib.context.api_context import APIContext


try:
    from unittest.mock import patch
except ImportError:
    from mock import patch

TOKEN_MOCK = '1234'


class APIKeyMock():
    def __init__(self, name):
        self.name = name
        self.token = TOKEN_MOCK
        self.type = None
        self.created_at = None
        self.updated_at = None
        self.grants = None


class APIKeyManagerMock():
    def create(self, name, **kwargs):
        return APIKeyMock(name)


class TestSQLClient(object):
    @patch('cartoframes.data.clients.auth_api_client._get_api_key_manager')
    def test_instantiation(self, get_api_key_manager_mock):
        get_api_key_manager_mock.return_value = APIKeyManagerMock()

        auth_api_client = AuthAPIClient()

        assert isinstance(auth_api_client, AuthAPIClient)
        assert isinstance(auth_api_client._api_key_manager, APIKeyManagerMock)

    @patch.object(APIContext, 'get_schema')
    @patch('cartoframes.data.clients.auth_api_client._get_api_key_manager')
    def test_create_api_key(self, get_api_key_manager_mock, get_schema_mock):
        get_api_key_manager_mock.return_value = APIKeyManagerMock()

        dataset = Dataset('fake_table', credentials=Credentials('fakeuser'))
        api_key_name = 'fake_name'

        auth_api_client = AuthAPIClient()
        token = auth_api_client.create_api_key([dataset], api_key_name)

        assert token == TOKEN_MOCK

    @patch.object(APIContext, 'get_schema')
    @patch('cartoframes.data.clients.auth_api_client._get_api_key_manager')
    def test_create_api_key_several_datasets(self, get_api_key_manager_mock, get_schema_mock):
        get_api_key_manager_mock.return_value = APIKeyManagerMock()

        dataset = Dataset('fake_table', credentials=Credentials('fakeuser'))
        api_key_name = 'fake_name'

        auth_api_client = AuthAPIClient()
        token = auth_api_client.create_api_key([dataset, dataset, dataset], api_key_name)

        assert token == TOKEN_MOCK

@@ -14,20 +14,19 @@
 mes.
-data
+viz
  import 
 Data
@@ -21,23 +21,22 @@
  import 
-Dataset
+Source
 %0Afrom ca
@@ -49,367 +49,379 @@
 mes.
-data.clients.auth_api_client import AuthAPIClient%0Afrom cartoframes.auth import Credentials%0Afrom cartoframes.lib.context.api_context import APIContext%0A%0A%0Atry:%0A    from unittest.mock import patch%0Aexcept ImportError:%0A    from mock import patch%0A%0ATOKEN_MOCK = '1234'%0A%0A%0Aclass APIKeyMock():%0A    def __init__(self, name):%0A        self.name = name%0A        self.token = 
+auth import Credentials%0Afrom cartoframes.io.context import ContextManager%0Afrom cartoframes.data.clients.auth_api_client import AuthAPIClient%0A%0Afrom ..mocks.api_key_mock import APIKeyManagerMock%0A%0ATOKEN_MOCK = '1234'%0A%0A%0Adef setup_mocks(mocker):%0A    mocker.patch(%0A        'cartoframes.data.clients.auth_api_client._get_api_key_manager',%0A        return_value=APIKeyManagerMock(
 TOKE
@@ -430,220 +430,179 @@
 MOCK
+))
 %0A    
-    self.type = None%0A        self.created_at = None%0A        self.updated_at = None%0A        self.grants = None%0A%0A%0Aclass APIKeyManagerMock():%0A    def create(self, name, **kwargs):%0A        return APIKeyMock(
+mocker.patch.object(ContextManager, 'compute_query')%0A    mocker.patch.object(ContextManager, 'get_schema')%0A    mocker.patch.object(ContextManager, 'get_table_
 name
+s'
 )%0A%0A%0A
@@ -615,11 +615,15 @@
 Test
-SQL
+AuthAPI
 Clie
@@ -638,84 +638,8 @@
 t):%0A
-    @patch('cartoframes.data.clients.auth_api_client._get_api_key_manager')%0A
     
@@ -671,32 +671,14 @@
 lf, 
-get_api_key_manager_
 mock
+er
 ):%0A 
@@ -688,66 +688,26 @@
     
-get_api_key_manager_mock.return_value = APIKeyManagerMock(
+setup_mocks(mocker
 )%0A%0A 
@@ -894,199 +894,44 @@
     
-@patch.object(APIContext, 'get_schema')%0A    @patch('cartoframes.data.clients.auth_api_client._get_api_key_manager')%0A    def test_create_api_key(self, get_api_key_manager_mock, get_schema_mock
+def test_create_api_key(self, mocker
 ):%0A 
@@ -941,94 +941,52 @@
     
-get_api_key_manager_mock.return_value = APIKeyManagerMock()%0A%0A        dataset = Dataset
+setup_mocks(mocker)%0A%0A        source = Source
 ('fa
@@ -1159,23 +1159,22 @@
 pi_key(%5B
-dataset
+source
 %5D, api_k
@@ -1223,128 +1223,8 @@
 CK%0A%0A
-    @patch.object(APIContext, 'get_schema')%0A    @patch('cartoframes.data.clients.auth_api_client._get_api_key_manager')%0A
     
@@ -1259,23 +1259,22 @@
 ral_
-dataset
+source
 s(self, 
 get_
@@ -1273,49 +1273,14 @@
 lf, 
-get_api_key_manager_mock, get_schema_
 mock
+er
 ):%0A 
@@ -1290,66 +1290,26 @@
     
-get_api_key_manager_mock.return_value = APIKeyManagerMock(
+setup_mocks(mocker
 )%0A%0A 
@@ -1319,25 +1319,23 @@
     
-dataset = Dataset
+source = Source
 ('fa
@@ -1512,33 +1512,30 @@
 ey(%5B
-dataset, dataset, dataset
+source, source, source
 %5D, a

Update unit auth tests