From f6ca078a7963ada75cfbe6ef7934305353122dd6 Mon Sep 17 00:00:00 2001
From: Dave Gramlich <callmehiphop@gmail.com>
Date: Wed, 24 Aug 2016 15:55:10 -0400
Subject: [PATCH] common: allow grpc service baseUrl override

---
 packages/common/src/grpc-service.js  | 20 +++++----
 packages/common/test/grpc-service.js | 62 ++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 7 deletions(-)

diff --git a/packages/common/src/grpc-service.js b/packages/common/src/grpc-service.js
index 3696ad39405a..3cfd995ab540 100644
--- a/packages/common/src/grpc-service.js
+++ b/packages/common/src/grpc-service.js
@@ -173,12 +173,18 @@ function GrpcService(config, options) {
     protoServices[service] = googleProtoFiles[service][apiVersion];
   }
 
-  for (var protoServiceName in protoServices) {
-    var protoService = this.loadProtoFile_(
-      protoServices[protoServiceName], config);
+  var self = this;
 
-    this.protos[protoServiceName] = protoService;
-  }
+  Object.keys(protoServices).forEach(function(name) {
+    var protoConfig = protoServices[name];
+    var service = self.loadProtoFile_(protoConfig, config);
+
+    self.protos[name] = service;
+
+    if (protoConfig.baseUrl) {
+      service.baseUrl = protoConfig.baseUrl;
+    }
+  });
 }
 
 nodeutil.inherits(GrpcService, Service);
@@ -677,7 +683,7 @@ GrpcService.prototype.loadProtoFile_ = function(protoConfig, config) {
   var apiVersion = protoConfig.apiVersion || config.apiVersion;
   var service = dotProp.get(services.google, serviceName);
 
-  return service[apiVersion];
+  return service[apiVersion] || service;
 };
 
 /**
@@ -701,7 +707,7 @@ GrpcService.prototype.getService_ = function(protoOpts) {
 
   if (!service) {
     service = new proto[protoOpts.service](
-      this.baseUrl,
+      proto.baseUrl || this.baseUrl,
       this.grpcCredentials
     );
 
diff --git a/packages/common/test/grpc-service.js b/packages/common/test/grpc-service.js
index 19fd9e5a1ed8..9d4622046900 100644
--- a/packages/common/test/grpc-service.js
+++ b/packages/common/test/grpc-service.js
@@ -307,6 +307,30 @@ describe('GrpcService', function() {
       );
     });
 
+    it('should store the baseUrl properly', function() {
+      var fakeBaseUrl = 'a.googleapis.com';
+
+      grpcLoadOverride = function() {
+        return MOCK_GRPC_API;
+      };
+
+      var config = extend(true, {}, CONFIG, {
+        protoServices: {
+          CustomServiceName: {
+            path: '../file/path.proto',
+            baseUrl: fakeBaseUrl
+          }
+        }
+      });
+
+      var grpcService = new GrpcService(config, OPTIONS);
+
+      assert.strictEqual(
+        grpcService.protos.CustomServiceName.baseUrl,
+        fakeBaseUrl
+      );
+    });
+
     it('should not run in the gcloud sandbox environment', function() {
       global.GCLOUD_SANDBOX_ENV = {};
       var grpcService = new GrpcService();
@@ -1361,6 +1385,26 @@ describe('GrpcService', function() {
       var service = grpcService.loadProtoFile_(fakeProtoConfig, fakeMainConfig);
       assert.strictEqual(service, fakeServices.google.FakeService.v1);
     });
+
+    it('should return the services object if invalid version', function() {
+      var fakeProtoConfig = {
+        path: '/root/dir/path',
+        service: 'FakeService',
+        apiVersion: null
+      };
+
+      var fakeMainConfig = {
+        service: 'OtherFakeService',
+        apiVersion: 'v2'
+      };
+
+      grpcLoadOverride = function(pathOpts, type, grpOpts) {
+        return fakeServices;
+      };
+
+      var service = grpcService.loadProtoFile_(fakeProtoConfig, fakeMainConfig);
+      assert.strictEqual(service, fakeServices.google.FakeService);
+    });
   });
 
   describe('getService_', function() {
@@ -1423,5 +1467,23 @@ describe('GrpcService', function() {
       var cachedService = grpcService.activeServiceMap_.get('Service');
       assert.strictEqual(cachedService, fakeService);
     });
+
+    it('should use the baseUrl override if applicable', function() {
+      var fakeBaseUrl = 'a.googleapis.com';
+      var fakeService = {};
+
+      grpcService.protos = {
+        Service: {
+          baseUrl: fakeBaseUrl,
+          Service: function(baseUrl, grpcCredentials) {
+            assert.strictEqual(baseUrl, fakeBaseUrl);
+            return fakeService;
+          }
+        }
+      };
+
+      var service = grpcService.getService_({ service: 'Service' });
+      assert.strictEqual(service, fakeService);
+    });
   });
 });