Send environment type (tests, prod) in telemetry events (#944)

## Changes
Set the type based on a presence of TEST_DEFAULT_CLUSTER_ID env var,
which is only set in e2e tests.

Also don't sent telemetry to a prod endpoint when we run unit tests.p

## Tests
Manually and with new unit tests

Author: ilia-db

packages/databricks-vscode/src/extension.ts

@@ -41,7 +41,7 @@ import {FeatureId, FeatureManager} from "./feature-manager/FeatureManager";
 import {DbConnectAccessVerifier} from "./language/DbConnectAccessVerifier";
 import {MsPythonExtensionWrapper} from "./language/MsPythonExtensionWrapper";
 import {DatabricksEnvFileManager} from "./file-managers/DatabricksEnvFileManager";
-import {Telemetry, toUserMetadata} from "./telemetry";
+import {getContextMetadata, Telemetry, toUserMetadata} from "./telemetry";
 import "./telemetry/commandExtensions";
 import {Events, Metadata} from "./telemetry/constants";
 import {DbConnectInstallPrompt} from "./language/DbConnectInstallPrompt";
@@ -100,6 +100,7 @@ export async function activate(
     }
 
     const telemetry = Telemetry.createDefault();
+    telemetry.setMetadata(Metadata.CONTEXT, getContextMetadata());
 
     const packageMetadata = await PackageJsonUtils.getMetadata(context);
     logging.NamedLogger.getOrCreate(Loggers.Extension).debug("Metadata", {

packages/databricks-vscode/src/telemetry/constants.ts

@@ -110,12 +110,15 @@ export type EventProperties = {
         : never;
 };
 
+export type EnvironmentType = "tests" | "prod";
+
 /**
  * Additional metadata collected from the extension, independent of the event itself.
  */
 /* eslint-disable @typescript-eslint/naming-convention */
 export enum Metadata {
     USER = "user",
+    CONTEXT = "context",
 }
 /* eslint-enable @typescript-eslint/naming-convention */
 
@@ -141,6 +144,12 @@ export class MetadataTypes {
             comment: "The kind of authentication used by the user",
         },
     };
+    [Metadata.CONTEXT]: EventType<{environmentType: EnvironmentType}> = {
+        environmentType: {
+            comment:
+                "A type of the environment this extension is running with (test, staging, prod)",
+        },
+    };
 }
 
 /** The type of all extra metadata collected by the extension. */

packages/databricks-vscode/src/telemetry/index.test.ts

@@ -2,7 +2,7 @@
 import TelemetryReporter from "@vscode/extension-telemetry";
 import assert from "assert";
 import {mock, instance, capture, when} from "ts-mockito";
-import {Telemetry, toUserMetadata} from ".";
+import {Telemetry, getContextMetadata, toUserMetadata} from ".";
 import {Events, Metadata} from "./constants";
 import {DatabricksWorkspace} from "../configuration/DatabricksWorkspace";
 import {Uri} from "vscode";
@@ -12,10 +12,14 @@ import {ApiClient, Config} from "@databricks/databricks-sdk";
 describe(__filename, () => {
     let reporter: TelemetryReporter;
     let telemetry: Telemetry;
+    const defaultClusterId = process.env["TEST_DEFAULT_CLUSTER_ID"];
     beforeEach(async () => {
         reporter = mock(TelemetryReporter);
         telemetry = new Telemetry(instance(reporter));
     });
+    afterEach(() => {
+        process.env["TEST_DEFAULT_CLUSTER_ID"] = defaultClusterId;
+    });
     it("should record expected properties and metrics", async () => {
         telemetry.recordEvent(Events.COMMAND_EXECUTION, {
             command: "testCommand",
@@ -36,6 +40,42 @@ describe(__filename, () => {
         });
     });
 
+    it("sets context metadata with prod env type", async () => {
+        delete process.env["TEST_DEFAULT_CLUSTER_ID"];
+        telemetry.setMetadata(Metadata.CONTEXT, getContextMetadata());
+        telemetry.recordEvent(Events.COMMAND_EXECUTION, {
+            command: "testCommand",
+            success: true,
+            duration: 100,
+        });
+        const [eventName, props] = capture(reporter.sendTelemetryEvent).last();
+        assert.equal(eventName, "commandExecution");
+        assert.deepEqual(props, {
+            "version": "1.0",
+            "event.command": "testCommand",
+            "event.success": "true",
+            "context.environmentType": "prod",
+        });
+    });
+
+    it("sets context metadata with tests env type", async () => {
+        process.env["TEST_DEFAULT_CLUSTER_ID"] = "123";
+        telemetry.setMetadata(Metadata.CONTEXT, getContextMetadata());
+        telemetry.recordEvent(Events.COMMAND_EXECUTION, {
+            command: "testCommand",
+            success: true,
+            duration: 100,
+        });
+        const [eventName, props] = capture(reporter.sendTelemetryEvent).last();
+        assert.equal(eventName, "commandExecution");
+        assert.deepEqual(props, {
+            "version": "1.0",
+            "event.command": "testCommand",
+            "event.success": "true",
+            "context.environmentType": "tests",
+        });
+    });
+
     it("sets user metadata correctly after logged in", async () => {
         const ws = mock(DatabricksWorkspace);
         when(ws.userName).thenReturn("miles@databricks.com");

packages/databricks-vscode/src/telemetry/index.ts

@@ -79,6 +79,14 @@ export async function toUserMetadata(
     return {...dbWorkspaceMetadata, ...authType};
 }
 
+export function getContextMetadata(): ExtraMetadata[Metadata.CONTEXT] {
+    return {
+        environmentType: process.env["TEST_DEFAULT_CLUSTER_ID"]
+            ? "tests"
+            : "prod",
+    };
+}
+
 function getTelemetryKey(): string {
     if (isDevExtension()) {
         return DEV_APP_INSIGHTS_CONFIGURATION_STRING;

packages/databricks-vscode/src/test/runTest.ts

@@ -2,6 +2,7 @@ import path from "path";
 import os from "os";
 import fs from "fs/promises";
 
+import {EXTENSION_DEVELOPMENT} from "../utils/developmentUtils";
 import {downloadAndUnzipVSCode, runTests} from "@vscode/test-electron";
 
 async function main() {
@@ -28,6 +29,9 @@ async function main() {
             extensionDevelopmentPath,
             extensionTestsPath,
             launchArgs: ["--user-data-dir", `${os.tmpdir()}`],
+            extensionTestsEnv: {
+                [EXTENSION_DEVELOPMENT]: "true",
+            },
         });
     } catch (err) {
         // eslint-disable-next-line no-console