Esempio n. 1
0
class HiveQueryEngineTestCase(unittest.TestCase):

    @patch.object(HiveQueryEngine, '_connect', return_value=MagicMock())
    def setUp(self, connect):
        self.hive = HiveQueryEngine('1/XXX', 'https://api.treasuredata.com/', 'sample_datasets', True)
        self.assertTrue(connect.called)

    def test_user_agent(self):
        hive_no_header = HiveQueryEngine('1/XXX', 'https://api.treasuredata.com/', 'sample_datasets', False)
        self.assertEquals(hive_no_header.create_header('foo'), '')

        ua = self.hive.user_agent
        self.assertEquals(ua, 'pytd/%s (tdclient/%s)' % (__version__, tdclient.__version__))

    def test_create_header(self):
        ua = self.hive.user_agent

        header = self.hive.create_header()
        self.assertEquals(header, "-- client: {0}\n".format(ua))

        header = self.hive.create_header('foo')
        self.assertEquals(header, "-- client: {0}\n-- foo\n".format(ua))

        header = self.hive.create_header(['foo', 'bar'])
        self.assertEquals(header, "-- client: {0}\n-- foo\n-- bar\n".format(ua))

    def test_cursor(self):
        self.hive.cursor()
        self.assertTrue(self.hive.engine.cursor.called)

    def test_close(self):
        self.hive.close()
        self.assertTrue(self.hive.engine.close.called)
Esempio n. 2
0
class HiveQueryEngineTestCase(unittest.TestCase):
    @patch.object(HiveQueryEngine, "_connect", return_value=MagicMock())
    def setUp(self, connect):
        self.hive = HiveQueryEngine("1/XXX", "https://api.treasuredata.com/",
                                    "sample_datasets", True)
        self.assertTrue(connect.called)
        self.assertEqual(self.hive.executed, None)

    def test_user_agent(self):
        hive_no_header = HiveQueryEngine("1/XXX",
                                         "https://api.treasuredata.com/",
                                         "sample_datasets", False)
        self.assertEqual(hive_no_header.create_header("foo"), "")

        ua = self.hive.user_agent
        self.assertEqual(
            ua, "pytd/{0} (tdclient/{1})".format(__version__,
                                                 tdclient.__version__))

    def test_create_header(self):
        ua = self.hive.user_agent

        header = self.hive.create_header()
        self.assertEqual(header, "-- client: {0}\n".format(ua))

        header = self.hive.create_header("foo")
        self.assertEqual(header, "-- client: {0}\n-- foo\n".format(ua))

        header = self.hive.create_header(["foo", "bar"])
        self.assertEqual(header, "-- client: {0}\n-- foo\n-- bar\n".format(ua))

    def test_cursor(self):
        self.hive.cursor()
        self.assertTrue(self.hive.engine.cursor.called)

        # the `force_tdclient` flag has no effect for HiveQueryEngine
        self.hive.cursor(force_tdclient=False)
        self.assertTrue(self.hive.engine.cursor.called)

    def test_close(self):
        self.hive.close()
        self.assertTrue(self.hive.engine.close.called)