Make TraceWriter a Service object (#417)

PR-URL: #417

bin/run-test.sh

@@ -39,7 +39,7 @@ function run {
 }
 
 # Run test/coverage
-for test in test/*.js test/plugins/*.js ;
+for test in test/test-*.js test/plugins/*.js ;
 do
   if [[ ! $(node --version) =~ v0\.12\..* || ! "${test}" =~ .*trace\-koa\.js ]]
   then

src/trace-writer.js

@@ -17,7 +17,8 @@
 'use strict';
 
 var gcpMetadata = require('gcp-metadata');
-var util = require('@google-cloud/common').util;
+var common = require('@google-cloud/common');
+var util = require('util');
 var traceLabels = require('./trace-labels.js');
 var pjson = require('../package.json');
 var constants = require('./constants.js');
@@ -35,19 +36,22 @@ headers[constants.TRACE_AGENT_REQUEST_HEADER] = 1;
  *   authorization credentials.
  * @constructor
  */
-function TraceWriter(logger, config) {
- /** @private */
-  this.logger_ = logger;
+function TraceWriter(logger, options) {
+  options = options || {};
+
+  var serviceOptions = {
+    packageJson: pjson,
+    projectIdRequired: false,
+    baseUrl: 'https://cloudtrace.googleapis.com/v1',
+    scopes: SCOPES
+  };
+  common.Service.call(this, serviceOptions, options);
 
   /** @private */
-  this.config_ = config;
+  this.logger_ = logger;
 
-  /** @private {function} authenticated request function */
-  this.request_ = util.makeAuthenticatedRequestFactory({
-    scopes: SCOPES,
-    credentials: config.credentials,
-    keyFile: config.keyFilename
-  });
+  /** @private */
+  this.config_ = options;
 
   /** @private {Array<string>} stringified traces to be published */
   this.buffer_ = [];
@@ -95,6 +99,7 @@ function TraceWriter(logger, config) {
     });
   });
 }
+util.inherits(TraceWriter, common.Service);
 
 TraceWriter.prototype.stop = function() {
   this.isActive = false;
@@ -244,16 +249,19 @@ TraceWriter.prototype.flushBuffer_ = function(projectId) {
  * @param {string} json The stringified json representation of the queued traces.
  */
 TraceWriter.prototype.publish_ = function(projectId, json) {
+  // TODO(ofrobots): assert.ok(this.config_.project), and stop accepting
+  // projectId as an argument.
   var that = this;
   var uri = 'https://cloudtrace.googleapis.com/v1/projects/' +
     projectId + '/traces';
 
-  this.request_({
+  var options = {
     method: 'PATCH',
     uri: uri,
     body: json,
     headers: headers
-  }, function(err, response, body) {
+  };
+  that.request(options, function(err, body, response) {
     if (err) {
       that.logger_.error('TraceWriter: error: ',
         ((response && response.statusCode) || '') + '\n' + err.stack);

test/nocks.js

@@ -0,0 +1,60 @@
+/**
+ * Copyright 2016 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.
+ */
+'use strict';
+
+var nock = require('nock');
+
+// In the future _=>true.
+function accept() {
+  return true;
+}
+
+function nockOAuth2(validator) {
+  validator = validator || accept;
+  return nock('https://accounts.google.com')
+      .post('/o/oauth2/token', validator)
+      .once()
+      .reply(200, {
+        refresh_token: 'hello',
+        access_token: 'goodbye',
+        expiry_date: new Date(9999, 1, 1)
+      });
+}
+
+function nockProjectId(reply) {
+  return nock('http://metadata.google.internal')
+    .get('/computeMetadata/v1/project/project-id')
+    .once()
+    .reply(reply);
+}
+
+function nockPatchTraces(project, validator, reply, withError) {
+  validator = validator || accept;
+  var scope = nock('https://cloudtrace.googleapis.com')
+      .intercept('/v1/projects/' + project + '/traces', 'PATCH', validator);
+  if (withError) {
+    scope = scope.replyWithError(reply);
+  } else {
+    scope = scope.reply(reply || 200);
+  }
+  return scope;
+}
+
+module.exports = {
+  oauth2: nockOAuth2,
+  patchTraces: nockPatchTraces,
+  projectId: nockProjectId,
+};

test/plugins/common.js

@@ -232,7 +232,7 @@ function installNoopTraceWriter(agent) {
 }
 
 function avoidTraceWriterAuth(agent) {
-  agent.private_().traceWriter.request_ = request;
+  agent.private_().traceWriter.request = request;
 }
 
 function stopAgent(agent) {

test/test-trace-writer.js

@@ -17,95 +17,140 @@
 'use strict';
 
 var assert = require('assert');
+var fakeCredentials = require('./fixtures/gcloud-credentials.json');
 var nock = require('nock');
-var cls = require('../src/cls.js');
-var common = require('./plugins/common.js');
-var trace = require('..');
+var nocks = require('./nocks.js');
+var Service = require('@google-cloud/common').Service;
+var traceLabels = require('../src/trace-labels.js');
 
 nock.disableNetConnect();
 
-var uri = 'https://cloudtrace.googleapis.com';
-var path = '/v1/projects/0/traces';
+var PROJECT = 'fake-project';
+var DEFAULT_DELAY = 200;
 
-process.env.GCLOUD_PROJECT = 0;
 
-var queueSpans = function(n, agent) {
-  for (var i = 0; i < n; i++) {
-    common.runInTransaction(agent, function(end) {
-      end();
-    });
-  }
+var fakeLogger = {
+  warn: function() {},
+  info: function() {},
+  error: function() {},
+  debug: function() {}
 };
 
-describe('tracewriter publishing', function() {
-
-  it('should publish when queue fills', function(done) {
-    var buf;
-    var scope = nock(uri)
-        .intercept(path, 'PATCH', function(body) {
-          assert.equal(JSON.stringify(body.traces), JSON.stringify(buf));
-          return true;
-        }).reply(200);
-    var agent = trace.start({
-      bufferSize: 2,
-      samplingRate: 0,
-      forceNewAgent_: true
-    });
-    common.avoidTraceWriterAuth(agent);
-    cls.getNamespace().run(function() {
-      queueSpans(2, agent);
-      buf = common.getTraces(agent);
-      setTimeout(function() {
-        scope.done();
+function createFakeSpan(name) {
+  // creates a fake span.
+  return {
+    trace: {
+      spans: [
+        {
+          name: name,
+          startTime: 'fake startTime',
+          endTime: '',
+          closed_: false,
+          labels_: {},
+          close: function() { this.closed_ = true; },
+        }
+      ]
+    },
+    labels_: {},
+    addLabel: function(k, v) { this.labels_[k] = v; },
+  };
+}
+
+describe('TraceWriter', function() {
+  var TraceWriter = require('../src/trace-writer.js');
+
+  it('should be a Service instance', function() {
+    var writer = new TraceWriter(fakeLogger, {projectId: 'fake project'});
+    assert.ok(writer instanceof Service);
+  });
+
+  describe('projectId', function() {
+    it('should request project from metadata if not locally available',
+       function(done) {
+         var scope = nocks.projectId('from metadata');
+         // the constructor should fetch the projectId.
+         new TraceWriter(fakeLogger);
+         setTimeout(function() {
+           assert.ok(scope.isDone());
+           done();
+         }, DEFAULT_DELAY);
+       });
+  });
+
+  describe('writeSpan', function(done) {
+
+    it('should close spans, add defaultLabels and queue', function(done) {
+      var writer =
+          new TraceWriter(fakeLogger, {projectId: PROJECT, bufferSize: 4});
+      var spanData = createFakeSpan('fake span');
+      writer.queueTrace_ = function(trace) {
+        assert.ok(trace && trace.spans && trace.spans[0]);
+        var span = trace.spans[0];
+        assert.strictEqual(span.name, 'fake span');
+        assert.ok(span.closed_);
+        assert.ok(spanData.labels_[traceLabels.AGENT_DATA]);
+        // TODO(ofrobots): check serviceContext labels as well.
         done();
-      }, 80);
+      };
+
+      // TODO(ofrobots): the delay is needed to allow async initialization of
+      // labels.
+      setTimeout(function() { writer.writeSpan(spanData); }, DEFAULT_DELAY);
     });
   });
 
-  it('should publish after timeout', function(done) {
-    var buf;
-    var scope = nock(uri)
-        .intercept(path, 'PATCH', function(body) {
-          assert.equal(JSON.stringify(body.traces), JSON.stringify(buf));
-          return true;
-        }).reply(200);
-    var agent = trace.start({
-      flushDelaySeconds: 0.01,
-      samplingRate: -1,
-      forceNewAgent_: true
+  describe('publish', function() {
+    it('should submit a PATCH request to the API', function(done) {
+      nocks.oauth2();
+      var scope = nocks.patchTraces(PROJECT);
+
+      var writer = new TraceWriter(
+          fakeLogger, {projectId: PROJECT, credentials: fakeCredentials});
+      writer.publish_(PROJECT, '{"valid": "json"}');
+      setTimeout(function() {
+        assert.ok(scope.isDone());
+        done();
+      }, DEFAULT_DELAY);
     });
-    common.avoidTraceWriterAuth(agent);
-    cls.getNamespace().run(function() {
-      queueSpans(1, agent);
-      buf = common.getTraces(agent);
+
+    it('should drop on server error', function(done) {
+      var MESSAGE = {valid: 'json'};
+      nocks.oauth2();
+      var scope = nocks.patchTraces(PROJECT, null, 'Simulated Network Error',
+                                    true /* withError */);
+
+      var writer = new TraceWriter(
+          fakeLogger, {projectId: PROJECT, credentials: fakeCredentials});
+      writer.publish_(PROJECT, JSON.stringify(MESSAGE));
       setTimeout(function() {
-        scope.done();
+        assert.ok(scope.isDone());
+        assert.equal(writer.buffer_.length, 0);
         done();
-      }, 20);
+      }, DEFAULT_DELAY);
     });
   });
 
-  it('should drop on server error', function(done) {
-    var buf;
-    var scope = nock(uri)
-        .intercept(path, 'PATCH', function(body) {
-          assert.equal(JSON.stringify(body.traces), JSON.stringify(buf));
-          return true;
-        }).replyWithError('Simulated Network Error');
-    var agent = trace.start({
-      bufferSize: 2,
-      samplingRate: -1,
-      forceNewAgent_: true
+  describe('publishing', function() {
+    it('should publish when the queue fills', function(done) {
+      var writer = new TraceWriter(
+          fakeLogger,
+          {projectId: PROJECT, bufferSize: 4, flushDelaySeconds: 3600});
+      writer.publish_ = function() { done(); };
+      for (var i = 0; i < 4; i++) {
+        writer.writeSpan(createFakeSpan(i));
+      }
     });
-    common.avoidTraceWriterAuth(agent);
-    cls.getNamespace().run(function() {
-      queueSpans(2, agent);
-      buf = common.getTraces(agent);
+
+    it('should publish after timeout', function(done) {
+      var published = false;
+      var writer = new TraceWriter(
+          fakeLogger, {projectId: PROJECT, flushDelaySeconds: 0.01});
+      writer.publish_ = function() { published = true; };
+      writer.writeSpan(createFakeSpan('fake span'));
       setTimeout(function() {
-        scope.done();
+        assert.ok(published);
         done();
-      }, 20);
+      }, DEFAULT_DELAY);
     });
   });
-
 });