From ae8ebaa75791d70d43df67f77e467d856845168f Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Fri, 21 Feb 2025 16:58:45 +0100
Subject: [PATCH] add acc test for oauth'

---
 .../upload-succeeds-oauth/out.requests.txt    | 45 +++++++++++++++++++
 .../out.upload_process.txt                    |  5 +++
 .../upload-succeeds-oauth/output.txt          |  4 ++
 .../script                                    |  7 ++-
 .../telemetry/upload-succeeds-oauth/test.toml | 19 ++++++++
 .../telemetry/upload-succeeds/test.toml       | 11 -----
 6 files changed, 79 insertions(+), 12 deletions(-)
 create mode 100644 acceptance/telemetry/upload-succeeds-oauth/out.requests.txt
 create mode 100644 acceptance/telemetry/upload-succeeds-oauth/out.upload_process.txt
 create mode 100644 acceptance/telemetry/upload-succeeds-oauth/output.txt
 rename acceptance/telemetry/{upload-succeeds => upload-succeeds-oauth}/script (65%)
 create mode 100644 acceptance/telemetry/upload-succeeds-oauth/test.toml
 delete mode 100644 acceptance/telemetry/upload-succeeds/test.toml

diff --git a/acceptance/telemetry/upload-succeeds-oauth/out.requests.txt b/acceptance/telemetry/upload-succeeds-oauth/out.requests.txt
new file mode 100644
index 000000000..1cbde659a
--- /dev/null
+++ b/acceptance/telemetry/upload-succeeds-oauth/out.requests.txt
@@ -0,0 +1,45 @@
+{
+  "method": "GET",
+  "path": "/oidc/.well-known/oauth-authorization-server"
+}
+{
+  "headers": {
+    "Authorization": [
+      "Basic [ENCODED_AUTH]"
+    ]
+  },
+  "method": "POST",
+  "path": "/oidc/v1/token",
+  "raw_body": "grant_type=client_credentials\u0026scope=all-apis"
+}
+{
+  "method": "GET",
+  "path": "/oidc/.well-known/oauth-authorization-server"
+}
+{
+  "headers": {
+    "Authorization": [
+      "Basic [ENCODED_AUTH]"
+    ]
+  },
+  "method": "POST",
+  "path": "/oidc/v1/token",
+  "raw_body": "grant_type=client_credentials\u0026scope=all-apis"
+}
+{
+  "headers": {
+    "Authorization": [
+      "Bearer oauth-token"
+    ]
+  },
+  "method": "POST",
+  "path": "/telemetry-ext",
+  "body": {
+    "uploadTime": "UNIX_TIME_MILLIS",
+    "items": [],
+    "protoLogs": [
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"OS\",\"execution_time_ms\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE1\"}}}}",
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"OS\",\"execution_time_ms\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
+    ]
+  }
+}
diff --git a/acceptance/telemetry/upload-succeeds-oauth/out.upload_process.txt b/acceptance/telemetry/upload-succeeds-oauth/out.upload_process.txt
new file mode 100644
index 000000000..0b26be5b9
--- /dev/null
+++ b/acceptance/telemetry/upload-succeeds-oauth/out.upload_process.txt
@@ -0,0 +1,5 @@
+Telemetry logs uploaded successfully
+Response:
+{"errors":null,"numProtoSuccess":2}
+
+2025/02/21 16:56:09 INFO Ignoring basic auth, because oauth-m2m is preferred sdk=true
diff --git a/acceptance/telemetry/upload-succeeds-oauth/output.txt b/acceptance/telemetry/upload-succeeds-oauth/output.txt
new file mode 100644
index 000000000..1a1f61588
--- /dev/null
+++ b/acceptance/telemetry/upload-succeeds-oauth/output.txt
@@ -0,0 +1,4 @@
+
+>>> [CLI] selftest send-telemetry
+waiting for telemetry process to finish
+[wait_pid] process has ended
diff --git a/acceptance/telemetry/upload-succeeds/script b/acceptance/telemetry/upload-succeeds-oauth/script
similarity index 65%
rename from acceptance/telemetry/upload-succeeds/script
rename to acceptance/telemetry/upload-succeeds-oauth/script
index 6aadde9f1..53bc02dba 100644
--- a/acceptance/telemetry/upload-succeeds/script
+++ b/acceptance/telemetry/upload-succeeds-oauth/script
@@ -1,7 +1,12 @@
 export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
 export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt
 
-# TODO: Move the environment variable calculation to the daemon library.
+# Unset the token which is configured by default
+# in acceptance tests
+export DATABRICKS_TOKEN=""
+
+export DATABRICKS_CLIENT_ID=client_id
+export DATABRICKS_CLIENT_SECRET=client_secret
 
 trace $CLI selftest send-telemetry
 
diff --git a/acceptance/telemetry/upload-succeeds-oauth/test.toml b/acceptance/telemetry/upload-succeeds-oauth/test.toml
new file mode 100644
index 000000000..c94972b18
--- /dev/null
+++ b/acceptance/telemetry/upload-succeeds-oauth/test.toml
@@ -0,0 +1,19 @@
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "numProtoSuccess": 2
+}
+'''
+
+# TODO: Remove the INFO log line.
+
+[[Repls]]
+Old = 'execution_time_ms\\\":\d{1,5},'
+New = 'execution_time_ms\":\"SMALL_INT\",'
+
+# "client_id:client_secret" in base64 is Y2xpZW50X2lkOmNsaWVudF9zZWNyZXQ=, expect to
+# see this in Authorization header
+[[Repls]]
+Old = "Y2xpZW50X2lkOmNsaWVudF9zZWNyZXQ="
+New = "[ENCODED_AUTH]"
diff --git a/acceptance/telemetry/upload-succeeds/test.toml b/acceptance/telemetry/upload-succeeds/test.toml
deleted file mode 100644
index 4826f42b0..000000000
--- a/acceptance/telemetry/upload-succeeds/test.toml
+++ /dev/null
@@ -1,11 +0,0 @@
-[[Server]]
-Pattern = "POST /telemetry-ext"
-Response.Body = '''
-{
-    "numProtoSuccess": 2
-}
-'''
-
-[[Repls]]
-Old = 'execution_time_ms\\\":\d{1,5},'
-New = 'execution_time_ms\":\"SMALL_INT\",'