Skip to content

Commit 0eef7ba

Browse files
tanmay-dbtanmay-db
committed
[Internal]
1 parent c7c1772 commit 0eef7ba

File tree

2 files changed

+140
-33
lines changed

2 files changed

+140
-33
lines changed

‎databricks-sdk-java/src/test/java/com/databricks/sdk/integration/DemoTest.java‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package com.databricks.sdk.integration;
2+
3+
import static org.junit.jupiter.api.Assertions.*;
4+
5+
import com.databricks.sdk.AccountClient;
6+
import com.databricks.sdk.WorkspaceClient;
7+
import com.databricks.sdk.core.DatabricksConfig;
8+
import com.databricks.sdk.mixin.NodeTypeSelector;
9+
import com.databricks.sdk.service.iam.Group;
10+
import com.databricks.sdk.service.iam.ListAccountGroupsRequest;
11+
import java.util.*;
12+
import org.junit.jupiter.api.Test;
13+
14+
// rm -r ~/.config/databricks-sdk-java/oauth/
15+
public class DemoTest {
16+
17+
18+
@Test
19+
void U2MAccounts() {
20+
com.databricks.sdk.core.DatabricksConfig cfg =
21+
new com.databricks.sdk.core.DatabricksConfig()
22+
.setHost("https://jumpy-arced-yeti.staging.databricks.com/")
23+
.setExperimentalIsUnifiedHost(true)
24+
.setAccountId("cdd2f1c0-e8c2-46a5-85ac-74d9bfd2d600")
25+
.setAuthType("external-browser")
26+
.setOAuthBrowserAuthTimeout(java.time.Duration.ofMinutes(2))
27+
.setOAuthRedirectUrl("http://localhost:8020");
28+
29+
AccountClient a = new AccountClient(cfg);
30+
31+
// Check access token generated at ~/.config/databricks-sdk-java/oauth/
32+
Iterable<Group> groups =
33+
a.groups()
34+
.list(new ListAccountGroupsRequest().setAttributes("displayName").setCount(10000L));
35+
36+
assertNotNull(groups);
37+
38+
// Display first group details
39+
Iterator<Group> iterator = groups.iterator();
40+
if (iterator.hasNext()) {
41+
Group firstGroup = iterator.next();
42+
System.out.println("First Group Details:");
43+
System.out.println(" Display Name: " + firstGroup.getDisplayName());
44+
System.out.println(" ID: " + firstGroup.getId());
45+
System.out.println(" Group: " + firstGroup);
46+
}
47+
}
48+
49+
@Test
50+
void U2mWorkspace() {
51+
com.databricks.sdk.core.DatabricksConfig cfg =
52+
new com.databricks.sdk.core.DatabricksConfig()
53+
.setHost("https://jumpy-arced-yeti.staging.databricks.com/")
54+
.setAccountId("cdd2f1c0-e8c2-46a5-85ac-74d9bfd2d600")
55+
.setWorkspaceId("4168072012817122")
56+
.setAuthType("external-browser")
57+
.setExperimentalIsUnifiedHost(true)
58+
.setOAuthBrowserAuthTimeout(java.time.Duration.ofMinutes(2))
59+
.setOAuthRedirectUrl("http://localhost:8020");
60+
61+
WorkspaceClient w = new WorkspaceClient(cfg);
62+
63+
// Oauth happens at account. Using the token above.
64+
String nodeTypeId = w.clusters().selectNodeType(new NodeTypeSelector().withLocalDisk());
65+
66+
assertNotNull(nodeTypeId);
67+
}
68+
69+
@Test
70+
void PAT() {
71+
DatabricksConfig config = new DatabricksConfig().setProfile("spog-pat").resolve();
72+
config.authenticate(); // Initialize auth before creating client
73+
74+
WorkspaceClient w = new WorkspaceClient(config);
75+
String nodeTypeId = w.clusters().selectNodeType(new NodeTypeSelector().withLocalDisk());
76+
77+
assertNotNull(nodeTypeId);
78+
}
79+
80+
@Test
81+
void M2M() {
82+
DatabricksConfig config = new DatabricksConfig().setProfile("spog-m2m").resolve();
83+
config.authenticate(); // Initialize auth before creating client
84+
85+
86+
AccountClient a = new AccountClient(config);
87+
88+
// Check access token generated at ~/.config/databricks-sdk-java/oauth/
89+
Iterable<Group> groups =
90+
a.groups()
91+
.list(new ListAccountGroupsRequest().setAttributes("displayName").setCount(10000L));
92+
93+
assertNotNull(groups);
94+
// Display first group details
95+
Iterator<Group> iterator = groups.iterator();
96+
if (iterator.hasNext()) {
97+
Group firstGroup = iterator.next();
98+
System.out.println("First Group Details:");
99+
System.out.println(" Display Name: " + firstGroup.getDisplayName());
100+
System.out.println(" ID: " + firstGroup.getId());
101+
System.out.println(" Group: " + firstGroup);
102+
}
103+
}
104+
105+
@Test
106+
void CLI() {
107+
DatabricksConfig config = new DatabricksConfig().setProfile("spog-demo-1").resolve();
108+
config.authenticate(); // Initialize auth before creating client
109+
110+
WorkspaceClient w = new WorkspaceClient(config);
111+
String nodeTypeId = w.clusters().selectNodeType(new NodeTypeSelector().withLocalDisk());
112+
113+
assertNotNull(nodeTypeId);
114+
}
115+
}

‎databricks-sdk-java/src/test/java/com/databricks/sdk/integration/framework/EnvTest.java‎

Lines changed: 25 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -45,42 +45,16 @@ public class EnvTest
4545
TEST_EXECUTION_RULES_BY_ENV = new HashMap<>();
4646
TEST_EXECUTION_RULES_BY_ENV.put(
4747
"workspace",
48-
(env) -> {
49-
if (env.containsKey("DATABRICKS_ACCOUNT_ID")) {
50-
return disabled("Skipping workspace-level test in account environment");
51-
}
52-
return enabled("okay to run");
53-
});
48+
(env) -> skipIfNotEnvironmentType(env, "WORKSPACE"));
5449
TEST_EXECUTION_RULES_BY_ENV.put(
5550
"ucws",
56-
(env) -> {
57-
if (env.containsKey("DATABRICKS_ACCOUNT_ID")) {
58-
return disabled("Skipping workspace-level test in account environment");
59-
} else if (!env.containsKey("TEST_METASTORE_ID")) {
60-
return disabled(
61-
"Skipping Unity Catalog workspace-level test in non-Unity Catalog environment");
62-
}
63-
return enabled("okay to run");
64-
});
51+
(env) -> skipIfNotEnvironmentType(env, "UC_WORKSPACE"));
6552
TEST_EXECUTION_RULES_BY_ENV.put(
6653
"account",
67-
(env) -> {
68-
if (!env.containsKey("DATABRICKS_ACCOUNT_ID")) {
69-
return disabled("Skipping account-level test in workspace environment");
70-
}
71-
return enabled("okay to run");
72-
});
54+
(env) -> skipIfNotEnvironmentType(env, "ACCOUNT"));
7355
TEST_EXECUTION_RULES_BY_ENV.put(
7456
"ucacct",
75-
(env) -> {
76-
if (!env.containsKey("DATABRICKS_ACCOUNT_ID")) {
77-
return disabled("Skipping account-level test in workspace environment");
78-
} else if (!env.containsKey("TEST_METASTORE_ID")) {
79-
return disabled(
80-
"Skipping Unity Catalog account-level test in non-Unity Catalog environment");
81-
}
82-
return enabled("okay to run");
83-
});
57+
(env) -> skipIfNotEnvironmentType(env, "UC_ACCOUNT"));
8458
MAPPER = makeObjectMapper();
8559
}
8660

@@ -127,10 +101,11 @@ public ConditionEvaluationResult evaluateExecutionCondition(ExtensionContext con
127101
if (methodParams.isPresent()) {
128102
for (Parameter parameter : methodParams.get()) {
129103
Class<?> type = parameter.getType();
130-
boolean hasAccount = env.containsKey("DATABRICKS_ACCOUNT_ID");
131-
if (type == WorkspaceClient.class && hasAccount) {
104+
String envType = env.getOrDefault("TEST_ENVIRONMENT_TYPE", "");
105+
boolean isAccountEnv = envType.equals("ACCOUNT") || envType.equals("UC_ACCOUNT");
106+
if (type == WorkspaceClient.class && isAccountEnv) {
132107
return ConditionEvaluationResult.disabled("Can't use workspace client in account env");
133-
} else if (type == AccountClient.class && !hasAccount) {
108+
} else if (type == AccountClient.class && !isAccountEnv) {
134109
return ConditionEvaluationResult.disabled("Can't use account client in workspace env");
135110
} else if (type == String.class) {
136111
EnvOrSkip envOrSkip = parameter.getAnnotation(EnvOrSkip.class);
@@ -309,6 +284,23 @@ private ConditionEvaluationResult evaluate(
309284
name, actual, regex));
310285
}
311286

287+
/**
288+
* Skips the test if TEST_ENVIRONMENT_TYPE doesn't match the expected type. TEST_ENVIRONMENT_TYPE
289+
* values: "ACCOUNT", "WORKSPACE", "UC_ACCOUNT", "UC_WORKSPACE".
290+
*/
291+
private static ConditionEvaluationResult skipIfNotEnvironmentType(
292+
Map<String, String> env, String expectedType) {
293+
String envType = env.get("TEST_ENVIRONMENT_TYPE");
294+
if (envType == null || envType.isEmpty()) {
295+
return disabled("Skipping test because TEST_ENVIRONMENT_TYPE is not set");
296+
}
297+
if (!envType.equals(expectedType)) {
298+
return disabled(
299+
format("Skipping %s test in %s environment", expectedType, envType));
300+
}
301+
return enabled("okay to run");
302+
}
303+
312304
private static ObjectMapper makeObjectMapper() {
313305
ObjectMapper mapper = new ObjectMapper();
314306
mapper

0 commit comments

Comments
 (0)