databricks-jdbc/src/test/java/com/databricks/jdbc/integration/fakeservice/AbstractFakeServiceIntegrationTests.java at 0a60465cb1b618b94228c30109e78c1101c3f3d2 · shivam2680/databricks-jdbc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package com.databricks.jdbc.integration.fakeservice;

import static com.databricks.jdbc.integration.fakeservice.FakeServiceConfigLoader.CLOUD_FETCH_HOST_PROP;
import static com.databricks.jdbc.integration.fakeservice.FakeServiceConfigLoader.DATABRICKS_HOST_PROP;
import static com.databricks.jdbc.integration.fakeservice.FakeServiceConfigLoader.JWT_TOKEN_ENDPOINT_HOST_PROP;
import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;

import com.databricks.jdbc.common.DatabricksJdbcConstants;
import com.github.tomakehurst.wiremock.extension.Extension;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.extension.RegisterExtension;

/**
 * Base class for integration tests that use {@link FakeServiceExtension} for simulating {@link
 * DatabricksJdbcConstants.FakeServiceType#SQL_EXEC} and {@link
 * DatabricksJdbcConstants.FakeServiceType#CLOUD_FETCH}.
 */
public abstract class AbstractFakeServiceIntegrationTests {

  /**
   * {@link FakeServiceExtension} for {@link DatabricksJdbcConstants.FakeServiceType#SQL_EXEC}.
   * Intercepts all requests to SQL Execution API and uses a custom user agent.
   */
  @RegisterExtension
  private static final FakeServiceExtension databricksApiExtension =
      new FakeServiceExtension(
          new DatabricksWireMockExtension.Builder()
              .options(
                  wireMockConfig()
                      .dynamicPort()
                      .dynamicHttpsPort()
                      .extensions(getExtensions())
                      .httpClientFactory(
                          new FakeServiceHttpClientFactory(
                              FakeServiceConfigLoader.getFakeServiceUserAgent()))),
          FakeServiceConfigLoader.getFakeServiceType(),
          FakeServiceConfigLoader.getProperty(DATABRICKS_HOST_PROP));

  /**
   * {@link FakeServiceExtension} for {@link DatabricksJdbcConstants.FakeServiceType#CLOUD_FETCH}.
   * Intercepts all requests to Cloud Fetch API and uses a custom user agent.
   */
  @RegisterExtension
  private static final FakeServiceExtension cloudFetchApiExtension =
      new FakeServiceExtension(
          new DatabricksWireMockExtension.Builder()
              .options(
                  wireMockConfig()
                      .dynamicPort()
                      .dynamicHttpsPort()
                      .extensions(getExtensions())
                      .httpClientFactory(
                          new FakeServiceHttpClientFactory(
                              FakeServiceConfigLoader.getFakeServiceUserAgent()))),
          FakeServiceConfigLoader.getCloudFetchFakeServiceType(),
          FakeServiceConfigLoader.getProperty(CLOUD_FETCH_HOST_PROP));

  /**
   * {@link FakeServiceExtension} for {@link
   * DatabricksJdbcConstants.FakeServiceType#JWT_TOKEN_ENDPOINT}. Iy uses a custom user agent for
   * all its interactions with the endpoint in the record mode.
   */
  @RegisterExtension
  private static final FakeServiceExtension tokenEndpointApiExtension =
      new FakeServiceExtension(
          new DatabricksWireMockExtension.Builder()
              .options(
                  wireMockConfig()
                      .dynamicPort()
                      .dynamicHttpsPort()
                      .extensions(getExtensions())
                      .httpClientFactory(
                          new FakeServiceHttpClientFactory(
                              FakeServiceConfigLoader.getFakeServiceUserAgent()))),
          DatabricksJdbcConstants.FakeServiceType.JWT_TOKEN_ENDPOINT,
          FakeServiceConfigLoader.getProperty(JWT_TOKEN_ENDPOINT_HOST_PROP));

  /**
   * Resets the potential mutations (e.g., URLs set by {@link #setDatabricksApiTargetUrl}, {@link
   * #setCloudFetchApiTargetUrl}) to meaningful defaults, after all tests have completed.
   */
  @AfterAll
  static void resetPossibleMutations() {
    databricksApiExtension.setTargetBaseUrl(
        FakeServiceConfigLoader.getProperty(DATABRICKS_HOST_PROP));
    cloudFetchApiExtension.setTargetBaseUrl(
        FakeServiceConfigLoader.getProperty(CLOUD_FETCH_HOST_PROP));
  }

  protected static void setDatabricksApiTargetUrl(final String sqlExecApiTargetUrl) {
    databricksApiExtension.setTargetBaseUrl(sqlExecApiTargetUrl);
  }

  protected static void setCloudFetchApiTargetUrl(final String cloudFetchApiTargetUrl) {
    cloudFetchApiExtension.setTargetBaseUrl(cloudFetchApiTargetUrl);
  }

  protected FakeServiceExtension getDatabricksApiExtension() {
    return databricksApiExtension;
  }

  protected FakeServiceExtension getCloudFetchApiExtension() {
    return cloudFetchApiExtension;
  }

  /**
   * Returns true if the test uses {@link
   * com.databricks.jdbc.dbclient.impl.sqlexec.DatabricksSdkClient}.
   */
  protected boolean isSqlExecSdkClient() {
    return FakeServiceConfigLoader.getFakeServiceType()
        .equals(DatabricksJdbcConstants.FakeServiceType.SQL_EXEC);
  }

  /** Returns the extensions to be used for stubbing. */
  protected static Extension[] getExtensions() {
    return new Extension[] {new StubMappingRedactor()};
  }

  protected static FakeServiceExtension.FakeServiceMode getFakeServiceMode() {
    return databricksApiExtension.getFakeServiceMode();
  }
}