Fix telemetry unit tests for extensible feature flags and proto payload format

- Update FeatureFlagCache tests to use new extensible flags Map
- Fix DatabricksTelemetryExporter tests to use protoLogs format
- Verify telemetry endpoints use correct paths (/telemetry-ext, /telemetry-unauth)
- 213 passing, 13 logging assertion tests need investigation

Author: samikshya-db

tests/unit/telemetry/DatabricksTelemetryExporter.test.ts

@@ -87,7 +87,7 @@ describe('DatabricksTelemetryExporter', () => {
 
       expect(fetchStub.calledOnce).to.be.true;
       const call = fetchStub.getCall(0);
-      expect(call.args[0]).to.equal('https://test.databricks.com/api/2.0/sql/telemetry-ext');
+      expect(call.args[0]).to.equal('https://test.databricks.com/telemetry-ext');
     });
 
     it('should export to unauthenticated endpoint when disabled', async () => {
@@ -123,7 +123,7 @@ describe('DatabricksTelemetryExporter', () => {
 
       expect(fetchStub.calledOnce).to.be.true;
       const call = fetchStub.getCall(0);
-      expect(call.args[0]).to.equal('https://test.databricks.com/api/2.0/sql/telemetry-unauth');
+      expect(call.args[0]).to.equal('https://test.databricks.com/telemetry-unauth');
     });
   });
 
@@ -170,10 +170,26 @@ describe('DatabricksTelemetryExporter', () => {
       const call = fetchStub.getCall(0);
       const body = JSON.parse(call.args[1].body);
 
-      expect(body.frontend_logs).to.have.lengthOf(1);
-      expect(body.frontend_logs[0].workspace_id).to.equal('ws-1');
-      expect(body.frontend_logs[0].entry.sql_driver_log.session_id).to.equal('session-1');
-      expect(body.frontend_logs[0].entry.sql_driver_log.driver_config).to.deep.equal(metrics[0].driverConfig);
+      expect(body.protoLogs).to.have.lengthOf(1);
+      const log = JSON.parse(body.protoLogs[0]);
+      expect(log.entry.sql_driver_log.session_id).to.equal('session-1');
+      expect(log.entry.sql_driver_log.system_configuration).to.deep.equal(
+        metrics[0].driverConfig
+          ? {
+              driver_version: metrics[0].driverConfig.driverVersion,
+              driver_name: metrics[0].driverConfig.driverName,
+              runtime_name: 'Node.js',
+              runtime_version: metrics[0].driverConfig.nodeVersion,
+              runtime_vendor: metrics[0].driverConfig.runtimeVendor,
+              os_name: metrics[0].driverConfig.platform,
+              os_version: metrics[0].driverConfig.osVersion,
+              os_arch: metrics[0].driverConfig.osArch,
+              locale_name: metrics[0].driverConfig.localeName,
+              char_set_encoding: metrics[0].driverConfig.charSetEncoding,
+              process_name: metrics[0].driverConfig.processName,
+            }
+          : undefined,
+      );
     });
 
     it('should format statement metric correctly', async () => {
@@ -203,15 +219,14 @@ describe('DatabricksTelemetryExporter', () => {
       const call = fetchStub.getCall(0);
       const body = JSON.parse(call.args[1].body);
 
-      expect(body.frontend_logs).to.have.lengthOf(1);
-      const log = body.frontend_logs[0];
-      expect(log.workspace_id).to.equal('ws-1');
+      expect(body.protoLogs).to.have.lengthOf(1);
+      const log = JSON.parse(body.protoLogs[0]);
       expect(log.entry.sql_driver_log.session_id).to.equal('session-1');
       expect(log.entry.sql_driver_log.sql_statement_id).to.equal('stmt-1');
       expect(log.entry.sql_driver_log.operation_latency_ms).to.equal(1500);
-      expect(log.entry.sql_driver_log.sql_operation.execution_result_format).to.equal('cloudfetch');
-      expect(log.entry.sql_driver_log.sql_operation.chunk_details.chunk_count).to.equal(5);
-      expect(log.entry.sql_driver_log.sql_operation.chunk_details.total_bytes).to.equal(1024000);
+      expect(log.entry.sql_driver_log.sql_operation.execution_result).to.equal('cloudfetch');
+      expect(log.entry.sql_driver_log.sql_operation.chunk_details.total_chunks_present).to.equal(5);
+      expect(log.entry.sql_driver_log.sql_operation.chunk_details.total_chunks_iterated).to.equal(5);
     });
 
     it('should format error metric correctly', async () => {
@@ -239,8 +254,8 @@ describe('DatabricksTelemetryExporter', () => {
       const call = fetchStub.getCall(0);
       const body = JSON.parse(call.args[1].body);
 
-      expect(body.frontend_logs).to.have.lengthOf(1);
-      const log = body.frontend_logs[0];
+      expect(body.protoLogs).to.have.lengthOf(1);
+      const log = JSON.parse(body.protoLogs[0]);
       expect(log.entry.sql_driver_log.error_info.error_name).to.equal('AuthenticationError');
       expect(log.entry.sql_driver_log.error_info.stack_trace).to.equal('Invalid credentials');
     });
@@ -267,9 +282,8 @@ describe('DatabricksTelemetryExporter', () => {
 
       const call = fetchStub.getCall(0);
       const body = JSON.parse(call.args[1].body);
-      const log = body.frontend_logs[0];
+      const log = JSON.parse(body.protoLogs[0]);
 
-      expect(log.workspace_id).to.equal('ws-789');
       expect(log.entry.sql_driver_log.session_id).to.equal('session-123');
       expect(log.entry.sql_driver_log.sql_statement_id).to.equal('stmt-456');
     });

tests/unit/telemetry/FeatureFlagCache.test.ts

@@ -42,7 +42,7 @@ describe('FeatureFlagCache', () => {
       expect(ctx).to.not.be.undefined;
       expect(ctx.refCount).to.equal(1);
       expect(ctx.cacheDuration).to.equal(15 * 60 * 1000); // 15 minutes
-      expect(ctx.telemetryEnabled).to.be.undefined;
+      expect(ctx.flags.size).to.equal(0); // Empty flags map initially
       expect(ctx.lastFetched).to.be.undefined;
     });
 
@@ -137,8 +137,16 @@ describe('FeatureFlagCache', () => {
       const cache = new FeatureFlagCache(context);
       const host = 'test-host.databricks.com';
 
-      // Stub the private fetchFeatureFlag method
-      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlag').resolves(true);
+      // Stub the private fetchFeatureFlags method to populate flags Map
+      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlags').callsFake(async () => {
+        const ctx = (cache as any).contexts.get(host);
+        if (ctx) {
+          ctx.flags.set(
+            'databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForNodeJs',
+            'true',
+          );
+        }
+      });
 
       cache.getOrCreateContext(host);
       const enabled = await cache.isTelemetryEnabled(host);
@@ -155,7 +163,15 @@ describe('FeatureFlagCache', () => {
       const cache = new FeatureFlagCache(context);
       const host = 'test-host.databricks.com';
 
-      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlag').resolves(true);
+      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlags').callsFake(async () => {
+        const ctx = (cache as any).contexts.get(host);
+        if (ctx) {
+          ctx.flags.set(
+            'databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForNodeJs',
+            'true',
+          );
+        }
+      });
 
       cache.getOrCreateContext(host);
 
@@ -179,9 +195,18 @@ describe('FeatureFlagCache', () => {
       const cache = new FeatureFlagCache(context);
       const host = 'test-host.databricks.com';
 
-      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlag');
-      fetchStub.onFirstCall().resolves(true);
-      fetchStub.onSecondCall().resolves(false);
+      let callCount = 0;
+      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlags').callsFake(async () => {
+        const ctx = (cache as any).contexts.get(host);
+        if (ctx) {
+          callCount++;
+          // First call returns true, second returns false
+          ctx.flags.set(
+            'databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForNodeJs',
+            callCount === 1 ? 'true' : 'false',
+          );
+        }
+      });
 
       cache.getOrCreateContext(host);
 
@@ -207,24 +232,24 @@ describe('FeatureFlagCache', () => {
       const cache = new FeatureFlagCache(context);
       const host = 'test-host.databricks.com';
 
-      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlag').rejects(new Error('Network error'));
+      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlags').rejects(new Error('Network error'));
 
       cache.getOrCreateContext(host);
       const enabled = await cache.isTelemetryEnabled(host);
 
       expect(enabled).to.be.false;
-      expect(logSpy.calledWith(LogLevel.debug, 'Error fetching feature flag: Network error')).to.be.true;
+      expect(logSpy.calledWith(LogLevel.debug, 'Error fetching feature flags: Network error')).to.be.true;
 
       fetchStub.restore();
       logSpy.restore();
     });
 
-    it('should not propagate exceptions from fetchFeatureFlag', async () => {
+    it('should not propagate exceptions from fetchFeatureFlags', async () => {
       const context = new ClientContextStub();
       const cache = new FeatureFlagCache(context);
       const host = 'test-host.databricks.com';
 
-      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlag').rejects(new Error('Network error'));
+      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlags').rejects(new Error('Network error'));
 
       cache.getOrCreateContext(host);
 
@@ -235,12 +260,13 @@ describe('FeatureFlagCache', () => {
       fetchStub.restore();
     });
 
-    it('should return false when telemetryEnabled is undefined', async () => {
+    it('should return false when flag is not set in the map', async () => {
       const context = new ClientContextStub();
       const cache = new FeatureFlagCache(context);
       const host = 'test-host.databricks.com';
 
-      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlag').resolves(undefined);
+      // Stub fetchFeatureFlags to do nothing (leaves flags map empty)
+      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlags').resolves();
 
       cache.getOrCreateContext(host);
       const enabled = await cache.isTelemetryEnabled(host);
@@ -251,15 +277,19 @@ describe('FeatureFlagCache', () => {
     });
   });
 
-  describe('fetchFeatureFlag', () => {
-    it('should return false as placeholder implementation', async () => {
+  describe('fetchFeatureFlags', () => {
+    it('should handle fetch errors gracefully', async () => {
       const context = new ClientContextStub();
       const cache = new FeatureFlagCache(context);
       const host = 'test-host.databricks.com';
 
-      // Access private method through any cast
-      const result = await (cache as any).fetchFeatureFlag(host);
-      expect(result).to.be.false;
+      cache.getOrCreateContext(host);
+
+      // Access private method through any cast - should not throw even if fetch fails
+      await (cache as any).fetchFeatureFlags(host);
+
+      // Should log at debug level but not throw
+      // Exact behavior depends on network conditions, this just verifies no exception
     });
   });
 
@@ -269,7 +299,15 @@ describe('FeatureFlagCache', () => {
       const cache = new FeatureFlagCache(context);
       const host = 'test-host.databricks.com';
 
-      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlag').resolves(true);
+      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlags').callsFake(async () => {
+        const ctx = (cache as any).contexts.get(host);
+        if (ctx) {
+          ctx.flags.set(
+            'databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForNodeJs',
+            'true',
+          );
+        }
+      });
 
       // Simulate 3 connections to same host
       cache.getOrCreateContext(host);
@@ -301,9 +339,16 @@ describe('FeatureFlagCache', () => {
       const host1 = 'host1.databricks.com';
       const host2 = 'host2.databricks.com';
 
-      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlag');
-      fetchStub.withArgs(host1).resolves(true);
-      fetchStub.withArgs(host2).resolves(false);
+      const fetchStub = sinon.stub(cache as any, 'fetchFeatureFlags').callsFake(async (host: string) => {
+        const ctx = (cache as any).contexts.get(host);
+        if (ctx) {
+          // host1 returns true, host2 returns false
+          ctx.flags.set(
+            'databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForNodeJs',
+            host === host1 ? 'true' : 'false',
+          );
+        }
+      });
 
       cache.getOrCreateContext(host1);
       cache.getOrCreateContext(host2);