update flightsql tests

apache · Apr 13, 2024 · 03c6461 · 03c6461
1 parent f57d8bd
commit 03c6461
Showing 1 changed file with 15 additions and 5 deletions.
diff --git a/go/arrow/flight/flightsql/client_test.go b/go/arrow/flight/flightsql/client_test.go
@@ -417,15 +417,15 @@ func (s *FlightSqlClientSuite) TestPreparedStatementExecuteParamBinding() {
  closeAct.Type = flightsql.ClosePreparedStatementActionType
 
  // results from createprepared statement
- result := &pb.ActionCreatePreparedStatementResult{
+ actionResult := &pb.ActionCreatePreparedStatementResult{
  PreparedStatementHandle: []byte(query),
  }
  schema := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true}}, nil)
- result.ParameterSchema = flight.SerializeSchema(schema, memory.DefaultAllocator)
+ actionResult.ParameterSchema = flight.SerializeSchema(schema, memory.DefaultAllocator)
 
  // mocked client stream
  var out anypb.Any
- out.MarshalFrom(result)
+ out.MarshalFrom(actionResult)
  data, _ := proto.Marshal(&out)
 
  createRsp := &mockDoActionClient{}
@@ -445,14 +445,19 @@ func (s *FlightSqlClientSuite) TestPreparedStatementExecuteParamBinding() {
 
  expectedDesc := getDesc(&pb.CommandPreparedStatementQuery{PreparedStatementHandle: []byte(query)})
 
+ // mocked DoPut result
+ doPutPreparedStatementResult := &pb.DoPutPreparedStatementResult{PreparedStatementHandle: []byte(query)}
+ resdata, _ := proto.Marshal(doPutPreparedStatementResult)
+ putResult := &pb.PutResult{ AppMetadata: resdata }
+
  // mocked client stream for DoPut
  mockedPut := &mockDoPutClient{}
  s.mockClient.On("DoPut", s.callOpts).Return(mockedPut, nil)
  mockedPut.On("Send", mock.MatchedBy(func(fd *flight.FlightData) bool {
  return proto.Equal(expectedDesc, fd.FlightDescriptor)
  })).Return(nil).Twice() // first sends schema message, second sends data
  mockedPut.On("CloseSend").Return(nil)
- mockedPut.On("Recv").Return((*pb.PutResult)(nil), nil)
+ mockedPut.On("Recv").Return(putResult, nil)
 
  infoCmd := &pb.CommandPreparedStatementQuery{PreparedStatementHandle: []byte(query)}
  desc := getDesc(infoCmd)
@@ -516,6 +521,11 @@ func (s *FlightSqlClientSuite) TestPreparedStatementExecuteReaderBinding() {
 
  expectedDesc := getDesc(&pb.CommandPreparedStatementQuery{PreparedStatementHandle: []byte(query)})
 
+ // mocked DoPut result
+ doPutPreparedStatementResult := &pb.DoPutPreparedStatementResult{PreparedStatementHandle: []byte(query)}
+ resdata, _ := proto.Marshal(doPutPreparedStatementResult)
+ putResult := &pb.PutResult{ AppMetadata: resdata }
+
  // mocked client stream for DoPut
  mockedPut := &mockDoPutClient{}
  s.mockClient.On("DoPut", s.callOpts).Return(mockedPut, nil)
@@ -528,7 +538,7 @@ func (s *FlightSqlClientSuite) TestPreparedStatementExecuteReaderBinding() {
  return fd.FlightDescriptor == nil
  })).Return(nil).Times(3)
  mockedPut.On("CloseSend").Return(nil)
- mockedPut.On("Recv").Return((*pb.PutResult)(nil), nil)
+ mockedPut.On("Recv").Return(putResult, nil)
 
  infoCmd := &pb.CommandPreparedStatementQuery{PreparedStatementHandle: []byte(query)}
  desc := getDesc(infoCmd)