Fixed staging operation tests

databricks · Sep 28, 2023 · a0f7fa4 · a0f7fa4
1 parent 06381db
commit a0f7fa4
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 2 deletions.
diff --git a/connection_test.go b/connection_test.go
@@ -121,6 +121,11 @@ func TestConn_executeStatement(t *testing.T) {
  },
  }
 
+ getResultSetMetadata := func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (_r *cli_service.TGetResultSetMetadataResp, _err error) {
+ var b = false
+ return &cli_service.TGetResultSetMetadataResp{IsStagingOperation: &b}, nil
+ }
+
  testClient := &client.TestClient{
  FnExecuteStatement: func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) {
  executeStatementCount++
@@ -130,6 +135,7 @@ func TestConn_executeStatement(t *testing.T) {
  closeOperationCount++
  return &cli_service.TCloseOperationResp{}, nil
  },
+ FnGetResultSetMetadata: getResultSetMetadata,
  }
  testConn := &conn{
  session: getTestSession(),
@@ -1103,6 +1109,10 @@ func TestConn_ExecContext(t *testing.T) {
  }
  return getOperationStatusResp, nil
  }
+ getResultSetMetadata := func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (_r *cli_service.TGetResultSetMetadataResp, _err error) {
+ var b = false
+ return &cli_service.TGetResultSetMetadataResp{IsStagingOperation: &b}, nil
+ }
 
  testClient := &client.TestClient{
  FnExecuteStatement: executeStatement,
@@ -1112,6 +1122,7 @@ func TestConn_ExecContext(t *testing.T) {
  assert.NoError(t, ctxErr)
  return &cli_service.TCloseOperationResp{}, nil
  },
+ FnGetResultSetMetadata: getResultSetMetadata,
  }
  testConn := &conn{
  session: getTestSession(),
@@ -1155,6 +1166,11 @@ func TestConn_ExecContext(t *testing.T) {
  return getOperationStatusResp, nil
  }
 
+ getResultSetMetadata := func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (_r *cli_service.TGetResultSetMetadataResp, _err error) {
+ var b = false
+ return &cli_service.TGetResultSetMetadataResp{IsStagingOperation: &b}, nil
+ }
+
  testClient := &client.TestClient{
  FnExecuteStatement: executeStatement,
  FnGetOperationStatus: getOperationStatus,
@@ -1173,7 +1189,9 @@ func TestConn_ExecContext(t *testing.T) {
  }
  return cancelOperationResp, nil
  },
+ FnGetResultSetMetadata: getResultSetMetadata,
  }
+
  testConn := &conn{
  session: getTestSession(),
  client: testClient,

diff --git a/driver_e2e_test.go b/driver_e2e_test.go
@@ -31,6 +31,7 @@ func TestWorkflowExample(t *testing.T) {
  )
  state := &callState{}
  // load basic responses
+ loadTestData(t, "GetResultSetMetadataNotStaging.json", &state.getResultSetMetadataResp)
  loadTestData(t, "OpenSessionSuccess.json", &state.openSessionResp)
  loadTestData(t, "CloseSessionSuccess.json", &state.closeSessionResp)
  loadTestData(t, "CloseOperationSuccess.json", &state.closeOperationResp)
@@ -259,6 +260,7 @@ func TestContextTimeoutExample(t *testing.T) {
  _ = logger.SetLogLevel("debug")
  state := &callState{}
  // load basic responses
+ loadTestData(t, "GetResultSetMetadataNotStaging.json", &state.getResultSetMetadataResp)
  loadTestData(t, "OpenSessionSuccess.json", &state.openSessionResp)
  loadTestData(t, "CloseSessionSuccess.json", &state.closeSessionResp)
  loadTestData(t, "CloseOperationSuccess.json", &state.closeOperationResp)
@@ -308,6 +310,7 @@ func TestRetries(t *testing.T) {
  _ = logger.SetLogLevel("debug")
  state := &callState{}
  // load basic responses
+ loadTestData(t, "GetResultSetMetadataNotStaging.json", &state.getResultSetMetadataResp)
  loadTestData(t, "OpenSessionSuccess.json", &state.openSessionResp)
  loadTestData(t, "CloseSessionSuccess.json", &state.closeSessionResp)
  loadTestData(t, "CloseOperationSuccess.json", &state.closeOperationResp)
@@ -333,6 +336,7 @@ func TestRetries(t *testing.T) {
  _ = logger.SetLogLevel("debug")
  state := &callState{}
  // load basic responses
+ loadTestData(t, "GetResultSetMetadataNotStaging.json", &state.getResultSetMetadataResp)
  loadTestData(t, "OpenSessionSuccess.json", &state.openSessionResp)
  loadTestData(t, "CloseSessionSuccess.json", &state.closeSessionResp)
  loadTestData(t, "CloseOperationSuccess.json", &state.closeOperationResp)
@@ -358,6 +362,7 @@ func TestRetries(t *testing.T) {
  _ = logger.SetLogLevel("debug")
  state := &callState{}
  // load basic responses
+ loadTestData(t, "GetResultSetMetadataNotStaging.json", &state.getResultSetMetadataResp)
  loadTestData(t, "OpenSessionSuccess.json", &state.openSessionResp)
  loadTestData(t, "CloseSessionSuccess.json", &state.closeSessionResp)
  loadTestData(t, "CloseOperationSuccess.json", &state.closeOperationResp)
@@ -392,6 +397,7 @@ func TestRetries(t *testing.T) {
  _ = logger.SetLogLevel("debug")
  state := &callState{}
  // load basic responses
+ loadTestData(t, "GetResultSetMetadataNotStaging.json", &state.getResultSetMetadataResp)
  loadTestData(t, "OpenSessionSuccess.json", &state.openSessionResp)
  loadTestData(t, "CloseSessionSuccess.json", &state.closeSessionResp)
  loadTestData(t, "CloseOperationSuccess.json", &state.closeOperationResp)
@@ -426,6 +432,7 @@ func TestRetries(t *testing.T) {
  _ = logger.SetLogLevel("debug")
  state := &callState{}
  // load basic responses
+ loadTestData(t, "GetResultSetMetadataNotStaging.json", &state.getResultSetMetadataResp)
  loadTestData(t, "OpenSessionSuccess.json", &state.openSessionResp)
  loadTestData(t, "CloseSessionSuccess.json", &state.closeSessionResp)
  loadTestData(t, "CloseOperationSuccess.json", &state.closeOperationResp)

diff --git a/statement_test.go b/statement_test.go
@@ -87,15 +87,16 @@ func TestStmt_ExecContext(t *testing.T) {
  }
  return getOperationStatusResp, nil
  }
- fetchResultSetMetadata := func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (_r *cli_service.TGetResultSetMetadataResp, _err error) {
+
+ getResultSetMetadata := func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (_r *cli_service.TGetResultSetMetadataResp, _err error) {
  var b = false
  return &cli_service.TGetResultSetMetadataResp{IsStagingOperation: &b}, nil
  }
 
  testClient := &client.TestClient{
  FnExecuteStatement: executeStatement,
  FnGetOperationStatus: getOperationStatus,
- FnGetResultSetMetadata: fetchResultSetMetadata,
+ FnGetResultSetMetadata: getResultSetMetadata,
  }
  testConn := &conn{
  session: getTestSession(),

diff --git a/testdata/GetResultSetMetadataNotStaging.json b/testdata/GetResultSetMetadataNotStaging.json
@@ -0,0 +1,6 @@
+{
+ "status": {
+ "statusCode": "SUCCESS_STATUS"
+ },
+ "isStagingOperation": false
+}