content: # Copyright 2015 Google Inc. All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
  # in compliance with the License. You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software distributed under the License
  # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  # or implied. See the License for the specific language governing permissions and limitations under
  # the License.
  
  from __future__ import absolute_import
  from __future__ import unicode_literals
  from builtins import range
  import datetime as dt
  import collections
  import mock
  from oauth2client.client import AccessTokenCredentials
  import pandas
  import unittest
  
  # import Python so we can mock the parts we need to here.
  import IPython
  import IPython.core.magic
  
  
  IPython.core.magic.register_line_cell_magic = mock.Mock()
  IPython.core.magic.register_line_magic = mock.Mock()
  IPython.core.magic.register_cell_magic = mock.Mock()
  IPython.get_ipython = mock.Mock()
  
  
  import datalab.bigquery  # noqa: E402
  import datalab.context  # noqa: E402
  import datalab.utils.commands  # noqa: E402
  
  
  class TestCases(unittest.TestCase):
  
    @staticmethod
    def _get_expected_cols():
      cols = [
        {'type': 'number', 'id': 'Column1', 'label': 'Column1'},
        {'type': 'number', 'id': 'Column2', 'label': 'Column2'},
        {'type': 'string', 'id': 'Column3', 'label': 'Column3'},
        {'type': 'boolean', 'id': 'Column4', 'label': 'Column4'},
        {'type': 'number', 'id': 'Column5', 'label': 'Column5'},
        {'type': 'datetime', 'id': 'Column6', 'label': 'Column6'}
      ]
      return cols
  
    @staticmethod
    def _timestamp(d):
      return (d - dt.datetime(1970, 1, 1)).total_seconds()
  
    @staticmethod
    def _get_raw_rows():
      rows = [
        {'f': [
          {'v': 1}, {'v': 2}, {'v': '3'}, {'v': 'true'}, {'v': 0.0},
          {'v': TestCases._timestamp(dt.datetime(2000, 1, 1))}
        ]},
        {'f': [
          {'v': 11}, {'v': 12}, {'v': '13'}, {'v': 'false'}, {'v': 0.2},
          {'v': TestCases._timestamp(dt.datetime(2000, 1, 2))}
        ]},
        {'f': [
          {'v': 21}, {'v': 22}, {'v': '23'}, {'v': 'true'}, {'v': 0.3},
          {'v': TestCases._timestamp(dt.datetime(2000, 1, 3))}
        ]},
        {'f': [
          {'v': 31}, {'v': 32}, {'v': '33'}, {'v': 'false'}, {'v': 0.4},
          {'v': TestCases._timestamp(dt.datetime(2000, 1, 4))}
        ]},
        {'f': [
          {'v': 41}, {'v': 42}, {'v': '43'}, {'v': 'true'}, {'v': 0.5},
          {'v': TestCases._timestamp(dt.datetime(2000, 1, 5))}
        ]},
        {'f': [
          {'v': 51}, {'v': 52}, {'v': '53'}, {'v': 'true'}, {'v': 0.6},
          {'v': TestCases._timestamp(dt.datetime(2000, 1, 6))}
        ]}
      ]
      return rows
  
    @staticmethod
    def _get_expected_rows():
      rows = [
        {'c': [
          {'v': 1}, {'v': 2}, {'v': '3'}, {'v': True}, {'v': 0.0}, {'v': dt.datetime(2000, 1, 1)}
        ]},
        {'c': [
          {'v': 11}, {'v': 12}, {'v': '13'}, {'v': False}, {'v': 0.2}, {'v': dt.datetime(2000, 1, 2)}
        ]},
        {'c': [
          {'v': 21}, {'v': 22}, {'v': '23'}, {'v': True}, {'v': 0.3}, {'v': dt.datetime(2000, 1, 3)}
        ]},
        {'c': [
          {'v': 31}, {'v': 32}, {'v': '33'}, {'v': False}, {'v': 0.4}, {'v': dt.datetime(2000, 1, 4)}
        ]},
        {'c': [
          {'v': 41}, {'v': 42}, {'v': '43'}, {'v': True}, {'v': 0.5}, {'v': dt.datetime(2000, 1, 5)}
        ]},
        {'c': [
          {'v': 51}, {'v': 52}, {'v': '53'}, {'v': True}, {'v': 0.6}, {'v': dt.datetime(2000, 1, 6)}
        ]}
      ]
      return rows
  
    @staticmethod
    def _get_test_data_as_list_of_dicts():
      test_data = [
        {'Column1': 1, 'Column2': 2, 'Column3': '3',
         'Column4': True, 'Column5': 0.0, 'Column6': dt.datetime(2000, 1, 1)},
        {'Column1': 11, 'Column2': 12, 'Column3': '13',
         'Column4': False, 'Column5': 0.2, 'Column6': dt.datetime(2000, 1, 2)},
        {'Column1': 21, 'Column2': 22, 'Column3': '23',
         'Column4': True, 'Column5': 0.3, 'Column6': dt.datetime(2000, 1, 3)},
        {'Column1': 31, 'Column2': 32, 'Column3': '33',
         'Column4': False, 'Column5': 0.4, 'Column6': dt.datetime(2000, 1, 4)},
        {'Column1': 41, 'Column2': 42, 'Column3': '43',
         'Column4': True, 'Column5': 0.5, 'Column6': dt.datetime(2000, 1, 5)},
        {'Column1': 51, 'Column2': 52, 'Column3': '53',
         'Column4': True, 'Column5': 0.6, 'Column6': dt.datetime(2000, 1, 6)}
      ]
      # Use OrderedDicts to make testing the result easier.
      for i in range(0, len(test_data)):
        test_data[i] = collections.OrderedDict(sorted(list(test_data[i].items()), key=lambda t: t[0]))
  
      return test_data
  
    def test_get_data_from_list_of_dicts(self):
      self._test_get_data(TestCases._get_test_data_as_list_of_dicts(), TestCases._get_expected_cols(),
                          TestCases._get_expected_rows(), 6,
                          datalab.utils.commands._utils._get_data_from_list_of_dicts)
      self._test_get_data(TestCases._get_test_data_as_list_of_dicts(), TestCases._get_expected_cols(),
                          TestCases._get_expected_rows(), 6, datalab.utils.commands._utils.get_data)
  
    def test_get_data_from_list_of_lists(self):
      test_data = [
        [1, 2, '3', True, 0.0, dt.datetime(2000, 1, 1)],
        [11, 12, '13', False, 0.2, dt.datetime(2000, 1, 2)],
        [21, 22, '23', True, 0.3, dt.datetime(2000, 1, 3)],
        [31, 32, '33', False, 0.4, dt.datetime(2000, 1, 4)],
        [41, 42, '43', True, 0.5, dt.datetime(2000, 1, 5)],
        [51, 52, '53', True, 0.6, dt.datetime(2000, 1, 6)],
      ]
  
      self._test_get_data(test_data, TestCases._get_expected_cols(), TestCases._get_expected_rows(),
                          6, datalab.utils.commands._utils._get_data_from_list_of_lists)
      self._test_get_data(test_data, TestCases._get_expected_cols(), TestCases._get_expected_rows(),
                          6, datalab.utils.commands._utils.get_data)
  
    def test_get_data_from_dataframe(self):
      df = pandas.DataFrame(self._get_test_data_as_list_of_dicts())
      self._test_get_data(df, TestCases._get_expected_cols(), TestCases._get_expected_rows(), 6,
                          datalab.utils.commands._utils._get_data_from_dataframe)
      self._test_get_data(df, TestCases._get_expected_cols(), TestCases._get_expected_rows(), 6,
                          datalab.utils.commands._utils.get_data)
  
    @mock.patch('datalab.bigquery._api.Api.tabledata_list')
    @mock.patch('datalab.bigquery._table.Table.exists')
    @mock.patch('datalab.bigquery._api.Api.tables_get')
    @mock.patch('datalab.context._context.Context.default')
    def test_get_data_from_table(self, mock_context_default, mock_api_tables_get,
                                 mock_table_exists, mock_api_tabledata_list):
      data = TestCases._get_expected_rows()
      mock_context_default.return_value = TestCases._create_context()
      mock_api_tables_get.return_value = {
        'numRows': len(data),
        'schema': {
          'fields': [
            {'name': 'Column1', 'type': 'INTEGER'},
            {'name': 'Column2', 'type': 'INTEGER'},
            {'name': 'Column3', 'type': 'STRING'},
            {'name': 'Column4', 'type': 'BOOLEAN'},
            {'name': 'Column5', 'type': 'FLOAT'},
            {'name': 'Column6', 'type': 'TIMESTAMP'}
          ]
        }
      }
      mock_table_exists.return_value = True
      raw_data = self._get_raw_rows()
  
      def tabledata_list(*args, **kwargs):
        start_index = kwargs['start_index']
        max_results = kwargs['max_results']
        if max_results < 0:
          max_results = len(data)
        return {'rows': raw_data[start_index:start_index + max_results]}
  
      mock_api_tabledata_list.side_effect = tabledata_list
      t = datalab.bigquery.Table('foo.bar')
      self._test_get_data(t, TestCases._get_expected_cols(), TestCases._get_expected_rows(), 6,
                          datalab.utils.commands._utils._get_data_from_table)
      self._test_get_data(t, TestCases._get_expected_cols(), TestCases._get_expected_rows(), 6,
                          datalab.utils.commands._utils.get_data)
  
    def test_get_data_from_empty_list(self):
      self._test_get_data([], [], [], 0, datalab.utils.commands._utils.get_data)
  
    def test_get_data_from_malformed_list(self):
      with self.assertRaises(Exception) as error:
        self._test_get_data(['foo', 'bar'], [], [], 0, datalab.utils.commands._utils.get_data)
      self.assertEquals('To get tabular data from a list it must contain dictionaries or lists.',
                        str(error.exception))
  
    def _test_get_data(self, test_data, cols, rows, expected_count, fn):
      self.maxDiff = None
      data, count = fn(test_data)
      self.assertEquals(expected_count, count)
      self.assertEquals({'cols': cols, 'rows': rows}, data)
  
      # Test first_row. Note that count must be set in this case so we use a value greater than the
      # data set size.
      for first in range(0, 6):
        data, count = fn(test_data, first_row=first, count=10)
        self.assertEquals(expected_count, count)
        self.assertEquals({'cols': cols, 'rows': rows[first:]}, data)
  
      # Test first_row + count
      for first in range(0, 6):
        data, count = fn(test_data, first_row=first, count=2)
        self.assertEquals(expected_count, count)
        self.assertEquals({'cols': cols, 'rows': rows[first:first + 2]}, data)
  
      # Test subsets of columns
  
      # No columns
      data, count = fn(test_data, fields=[])
      self.assertEquals({'cols': [], 'rows': [{'c': []}] * expected_count}, data)
  
      # Single column
      data, count = fn(test_data, fields=['Column3'])
  
      if expected_count == 0:
        return
  
      self.assertEquals({'cols': [cols[2]],
                         'rows': [{'c': [row['c'][2]]} for row in rows]}, data)
  
      # Multi-columns
      data, count = fn(test_data, fields=['Column1', 'Column3', 'Column6'])
      self.assertEquals({'cols': [cols[0], cols[2], cols[5]],
                         'rows': [{'c': [row['c'][0], row['c'][2], row['c'][5]]} for row in rows]},
                        data)
  
      # Switch order
      data, count = fn(test_data, fields=['Column3', 'Column1'])
      self.assertEquals({'cols': [cols[2], cols[0]],
                         'rows': [{'c': [row['c'][2], row['c'][0]]} for row in rows]}, data)
  
      # Select all
      data, count = fn(test_data,
                       fields=['Column1', 'Column2', 'Column3', 'Column4', 'Column5', 'Column6'])
      self.assertEquals({'cols': cols, 'rows': rows}, data)
  
    @staticmethod
    def _create_api():
      context = TestCases._create_context()
      return datalab.bigquery._api.Api(context.credentials, context.project_id)
  
    @staticmethod
    def _create_context():
      project_id = 'test'
      creds = AccessTokenCredentials('test_token', 'test_ua')
      return datalab.context.Context(project_id, creds)
  