Merge pull request #2249 from dhermes/add-bigtable-emulator

Updating Bigtable stub creation to work with emulator.

Author: dhermes

google/cloud/bigtable/client.py

@@ -27,8 +27,11 @@
 """
 
 
+import os
+
 from pkg_resources import get_distribution
 
+from google.cloud._helpers import make_insecure_stub
 from google.cloud._helpers import make_secure_stub
 from google.cloud.bigtable._generated import bigtable_instance_admin_pb2
 from google.cloud.bigtable._generated import bigtable_pb2
@@ -40,6 +43,7 @@
 from google.cloud.client import _ClientFactoryMixin
 from google.cloud.client import _ClientProjectMixin
 from google.cloud.credentials import get_credentials
+from google.cloud.environment_vars import BIGTABLE_EMULATOR
 
 
 TABLE_ADMIN_HOST = 'bigtableadmin.googleapis.com'
@@ -74,8 +78,12 @@ def _make_data_stub(client):
  :rtype: :class:`._generated.bigtable_pb2.BigtableStub`
  :returns: A gRPC stub object.
  """
- return make_secure_stub(client.credentials, client.user_agent,
- bigtable_pb2.BigtableStub, DATA_API_HOST)
+ if client.emulator_host is None:
+ return make_secure_stub(client.credentials, client.user_agent,
+ bigtable_pb2.BigtableStub, DATA_API_HOST)
+ else:
+ return make_insecure_stub(bigtable_pb2.BigtableStub,
+ client.emulator_host)
 
 
 def _make_instance_stub(client):
@@ -87,10 +95,15 @@ def _make_instance_stub(client):
  :rtype: :class:`.bigtable_instance_admin_pb2.BigtableInstanceAdminStub`
  :returns: A gRPC stub object.
  """
- return make_secure_stub(
- client.credentials, client.user_agent,
- bigtable_instance_admin_pb2.BigtableInstanceAdminStub,
- INSTANCE_ADMIN_HOST)
+ if client.emulator_host is None:
+ return make_secure_stub(
+ client.credentials, client.user_agent,
+ bigtable_instance_admin_pb2.BigtableInstanceAdminStub,
+ INSTANCE_ADMIN_HOST)
+ else:
+ return make_insecure_stub(
+ bigtable_instance_admin_pb2.BigtableInstanceAdminStub,
+ client.emulator_host)
 
 
 def _make_operations_stub(client):
@@ -105,9 +118,13 @@ def _make_operations_stub(client):
  :rtype: :class:`._generated.operations_grpc_pb2.OperationsStub`
  :returns: A gRPC stub object.
  """
- return make_secure_stub(client.credentials, client.user_agent,
- operations_grpc_pb2.OperationsStub,
- OPERATIONS_API_HOST)
+ if client.emulator_host is None:
+ return make_secure_stub(client.credentials, client.user_agent,
+ operations_grpc_pb2.OperationsStub,
+ OPERATIONS_API_HOST)
+ else:
+ return make_insecure_stub(operations_grpc_pb2.OperationsStub,
+ client.emulator_host)
 
 
 def _make_table_stub(client):
@@ -119,9 +136,15 @@ def _make_table_stub(client):
  :rtype: :class:`.bigtable_instance_admin_pb2.BigtableTableAdminStub`
  :returns: A gRPC stub object.
  """
- return make_secure_stub(client.credentials, client.user_agent,
- bigtable_table_admin_pb2.BigtableTableAdminStub,
- TABLE_ADMIN_HOST)
+ if client.emulator_host is None:
+ return make_secure_stub(
+ client.credentials, client.user_agent,
+ bigtable_table_admin_pb2.BigtableTableAdminStub,
+ TABLE_ADMIN_HOST)
+ else:
+ return make_insecure_stub(
+ bigtable_table_admin_pb2.BigtableTableAdminStub,
+ client.emulator_host)
 
 
 class Client(_ClientFactoryMixin, _ClientProjectMixin):
@@ -192,6 +215,7 @@ def __init__(self, project=None, credentials=None,
  pass
  self._credentials = credentials
  self.user_agent = user_agent
+ self.emulator_host = os.getenv(BIGTABLE_EMULATOR)
 
  # Create gRPC stubs for making requests.
  self._data_stub = _make_data_stub(self)

google/cloud/environment_vars.py

@@ -33,6 +33,9 @@
 PUBSUB_EMULATOR = 'PUBSUB_EMULATOR_HOST'
 """Environment variable defining host for Pub/Sub emulator."""
 
+BIGTABLE_EMULATOR = 'BIGTABLE_EMULATOR_HOST'
+"""Environment variable defining host for Bigtable emulator."""
+
 CREDENTIALS = 'GOOGLE_APPLICATION_CREDENTIALS'
 """Environment variable defining location of Google credentials."""
 

unit_tests/bigtable/test_client.py

@@ -22,7 +22,7 @@ def _callFUT(self, client):
  from google.cloud.bigtable.client import _make_data_stub
  return _make_data_stub(client)
 
- def test_it(self):
+ def test_without_emulator(self):
  from unit_tests._testing import _Monkey
  from google.cloud.bigtable import client as MUT
 
@@ -50,14 +50,39 @@ def mock_make_secure_stub(*args):
  ),
  ])
 
+ def test_with_emulator(self):
+ from unit_tests._testing import _Monkey
+ from google.cloud.bigtable import client as MUT
+
+ emulator_host = object()
+ client = _Client(None, None, emulator_host=emulator_host)
+
+ fake_stub = object()
+ make_insecure_stub_args = []
+
+ def mock_make_insecure_stub(*args):
+ make_insecure_stub_args.append(args)
+ return fake_stub
+
+ with _Monkey(MUT, make_insecure_stub=mock_make_insecure_stub):
+ result = self._callFUT(client)
+
+ self.assertIs(result, fake_stub)
+ self.assertEqual(make_insecure_stub_args, [
+ (
+ MUT.bigtable_pb2.BigtableStub,
+ emulator_host,
+ ),
+ ])
+
 
 class Test__make_instance_stub(unittest.TestCase):
 
  def _callFUT(self, client):
  from google.cloud.bigtable.client import _make_instance_stub
  return _make_instance_stub(client)
 
- def test_it(self):
+ def test_without_emulator(self):
  from unit_tests._testing import _Monkey
  from google.cloud.bigtable import client as MUT
 
@@ -85,14 +110,39 @@ def mock_make_secure_stub(*args):
  ),
  ])
 
+ def test_with_emulator(self):
+ from unit_tests._testing import _Monkey
+ from google.cloud.bigtable import client as MUT
+
+ emulator_host = object()
+ client = _Client(None, None, emulator_host=emulator_host)
+
+ fake_stub = object()
+ make_insecure_stub_args = []
+
+ def mock_make_insecure_stub(*args):
+ make_insecure_stub_args.append(args)
+ return fake_stub
+
+ with _Monkey(MUT, make_insecure_stub=mock_make_insecure_stub):
+ result = self._callFUT(client)
+
+ self.assertIs(result, fake_stub)
+ self.assertEqual(make_insecure_stub_args, [
+ (
+ MUT.bigtable_instance_admin_pb2.BigtableInstanceAdminStub,
+ emulator_host,
+ ),
+ ])
+
 
 class Test__make_operations_stub(unittest.TestCase):
 
  def _callFUT(self, client):
  from google.cloud.bigtable.client import _make_operations_stub
  return _make_operations_stub(client)
 
- def test_it(self):
+ def test_without_emulator(self):
  from unit_tests._testing import _Monkey
  from google.cloud.bigtable import client as MUT
 
@@ -120,14 +170,39 @@ def mock_make_secure_stub(*args):
  ),
  ])
 
+ def test_with_emulator(self):
+ from unit_tests._testing import _Monkey
+ from google.cloud.bigtable import client as MUT
+
+ emulator_host = object()
+ client = _Client(None, None, emulator_host=emulator_host)
+
+ fake_stub = object()
+ make_insecure_stub_args = []
+
+ def mock_make_insecure_stub(*args):
+ make_insecure_stub_args.append(args)
+ return fake_stub
+
+ with _Monkey(MUT, make_insecure_stub=mock_make_insecure_stub):
+ result = self._callFUT(client)
+
+ self.assertIs(result, fake_stub)
+ self.assertEqual(make_insecure_stub_args, [
+ (
+ MUT.operations_grpc_pb2.OperationsStub,
+ emulator_host,
+ ),
+ ])
+
 
 class Test__make_table_stub(unittest.TestCase):
 
  def _callFUT(self, client):
  from google.cloud.bigtable.client import _make_table_stub
  return _make_table_stub(client)
 
- def test_it(self):
+ def test_without_emulator(self):
  from unit_tests._testing import _Monkey
  from google.cloud.bigtable import client as MUT
 
@@ -155,6 +230,31 @@ def mock_make_secure_stub(*args):
  ),
  ])
 
+ def test_with_emulator(self):
+ from unit_tests._testing import _Monkey
+ from google.cloud.bigtable import client as MUT
+
+ emulator_host = object()
+ client = _Client(None, None, emulator_host=emulator_host)
+
+ fake_stub = object()
+ make_insecure_stub_args = []
+
+ def mock_make_insecure_stub(*args):
+ make_insecure_stub_args.append(args)
+ return fake_stub
+
+ with _Monkey(MUT, make_insecure_stub=mock_make_insecure_stub):
+ result = self._callFUT(client)
+
+ self.assertIs(result, fake_stub)
+ self.assertEqual(make_insecure_stub_args, [
+ (
+ MUT.bigtable_table_admin_pb2.BigtableTableAdminStub,
+ emulator_host,
+ ),
+ ])
+
 
 class TestClient(unittest.TestCase):
 
@@ -532,9 +632,10 @@ def __eq__(self, other):
 
 class _Client(object):
 
- def __init__(self, credentials, user_agent):
+ def __init__(self, credentials, user_agent, emulator_host=None):
  self.credentials = credentials
  self.user_agent = user_agent
+ self.emulator_host = emulator_host
 
 
 class _MakeStubMock(object):