瀏覽代碼

HUE-9100 [hive] Adding a set of tests to the hiveserver2 lib

Preparation for more developed session creation testing.
Romain 6 年之前
父節點
當前提交
96d41105a3
共有 2 個文件被更改,包括 151 次插入5 次删除
  1. 149 3
      apps/beeswax/src/beeswax/server/hive_server2_lib_tests.py
  2. 2 2
      desktop/core/src/desktop/api2_tests.py

+ 149 - 3
apps/beeswax/src/beeswax/server/hive_server2_lib_tests.py

@@ -20,18 +20,164 @@ import logging
 import sys
 
 if sys.version_info[0] > 2:
-  from unittest.mock import patch, Mock
+  from unittest.mock import patch, Mock, MagicMock
 else:
-  from mock import patch, Mock
+  from mock import patch, Mock, MagicMock
 
 from nose.tools import assert_equal, assert_true
+from TCLIService.ttypes import TStatusCode
 
-from beeswax.server.hive_server2_lib import HiveServerTable
+from desktop.lib.django_test_util import make_logged_in_client
+from desktop.lib.test_utils import grant_access
+from useradmin.models import User
+
+from beeswax.models import Session
+from beeswax.server.hive_server2_lib import HiveServerTable, HiveServerClient
 
 
 LOG = logging.getLogger(__name__)
 
 
+class TestHiveServerClient():
+
+  def setUp(self):
+    self.client = make_logged_in_client(username="test_hive_server2_lib", groupname="default", recreate=True, is_superuser=False)
+    self.user = User.objects.get(username="test_hive_server2_lib")
+
+    grant_access(self.user.username, self.user.username, "beeswax")
+
+    self.query_server = {
+        'principal': 'hue',
+        'server_name': 'hive',
+        'QUERY_TIMEOUT_S': 60,
+        'auth_username': 'hue',
+        'auth_password': 'hue',
+        'use_sasl': True,
+        'server_host': 'localhost',
+        'server_port': 10000,
+    }
+
+  def test_open_session(self):
+    query = Mock(
+      get_query_statement=Mock(return_value=['SELECT 1']),
+      settings=[]
+    )
+
+    with patch('beeswax.server.hive_server2_lib.thrift_util.get_client') as get_client:
+      get_client.return_value = Mock(
+        OpenSession=Mock(
+          return_value=Mock(
+            status=Mock(
+              statusCode=TStatusCode.SUCCESS_STATUS
+            ),
+            configuration={},
+            sessionHandle=Mock(
+              sessionId=Mock(
+                secret=b'1',
+                guid=b'1'
+              )
+            ),
+            serverProtocolVersion=11
+          )
+        ),
+        get_coordinator_host=Mock(return_value='hive-host')
+      )
+      session_count = Session.objects.filter(owner=self.user, application=self.query_server['server_name']).count()
+
+      # Send open session
+      session = HiveServerClient(self.query_server, self.user).open_session(self.user)
+
+      assert_equal(
+        session_count + 1,  # +1 as setUp resets the user which deletes cascade the sessions
+        Session.objects.filter(owner=self.user, application=self.query_server['server_name']).count()
+      )
+      assert_equal(
+        session.guid,
+        Session.objects.get_session(self.user, self.query_server['server_name']).guid.encode()
+      )
+
+  def test_explain(self):
+    query = Mock(
+      get_query_statement=Mock(return_value=['SELECT 1']),
+      settings=[]
+    )
+
+    with patch('beeswax.server.hive_server2_lib.thrift_util.get_client') as get_client:
+      get_client.return_value = Mock(
+        OpenSession=Mock(
+          return_value=Mock(
+            status=Mock(
+              statusCode=TStatusCode.SUCCESS_STATUS
+            ),
+            configuration={},
+            sessionHandle=Mock(
+              sessionId=Mock(
+                secret=b'1',
+                guid=b'1'
+              )
+            ),
+            serverProtocolVersion=11
+          )
+        ),
+        ExecuteStatement=Mock(
+          return_value=Mock(
+            status=Mock(
+              statusCode=TStatusCode.SUCCESS_STATUS
+            ),
+          )
+        ),
+        FetchResults=Mock(
+          return_value=Mock(
+            status=Mock(
+              statusCode=TStatusCode.SUCCESS_STATUS
+            ),
+            results=Mock(
+              columns=[
+                # Dump of `EXPLAIN SELECT 1`
+                Mock(stringVal=Mock(values=['Plan optimized by CBO.', '', 'Stage-0', '	  Fetch Operator', '5	    limit:-1' ], nulls='')),
+              ]
+            ),
+            schema=Mock(
+              columns=[
+                Mock(columnName='Explain'),
+              ]
+            )
+          )
+        ),
+        GetResultSetMetadata=Mock(
+          return_value=Mock(
+            status=Mock(
+              statusCode=TStatusCode.SUCCESS_STATUS
+            ),
+            results=Mock(
+              columns=[
+                Mock(stringVal=Mock(values=['Explain', ], nulls='')),  # Fake but ok
+              ]
+            ),
+            schema=Mock(
+              columns=[
+                Mock(columnName='primitiveEntry 7'),
+              ]
+            )
+          )
+        ),
+        get_coordinator_host=Mock(return_value='hive-host')
+      )
+      session_count = Session.objects.filter(owner=self.user, application=self.query_server['server_name']).count()
+
+      # Send explain
+      explain = HiveServerClient(self.query_server, self.user).explain(query)
+
+      assert_equal(
+        [['Plan optimized by CBO.'], [''], ['Stage-0'], ['	  Fetch Operator'], ['5	    limit:-1']],
+        list(explain.rows())
+      )
+      assert_equal(
+        session_count + 1,
+        Session.objects.filter(owner=self.user, application=self.query_server['server_name']).count()
+      )
+
+
 class TestHiveServerTable():
 
   def test_cols_impala(self):

+ 2 - 2
desktop/core/src/desktop/api2_tests.py

@@ -22,11 +22,11 @@ import re
 
 from nose.tools import assert_true, assert_false, assert_equal, assert_not_equal, assert_raises
 
-from useradmin.models import get_default_user_group
+from useradmin.models import get_default_user_group, User
 
 from desktop.lib.django_test_util import make_logged_in_client
 from desktop.lib.test_utils import grant_access
-from desktop.models import Document2, User
+from desktop.models import Document2
 
 
 class TestApi2(object):