@Override public com.google.api.services.bigquery.model.Table call() { return bigQueryRpc.create(tablePb, optionsMap); } },
@Test public void testDeleteDataset() { EasyMock.expect(bigqueryRpcMock.deleteDataset(PROJECT, DATASET, EMPTY_RPC_OPTIONS)) .andReturn(true); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); assertTrue(bigquery.delete(DATASET)); }
@Test public void testListEmptyDatasets() { ImmutableList<com.google.api.services.bigquery.model.Dataset> datasets = ImmutableList.of(); Tuple<String, Iterable<com.google.api.services.bigquery.model.Dataset>> result = Tuple.<String, Iterable<com.google.api.services.bigquery.model.Dataset>>of(null, datasets); EasyMock.expect(bigqueryRpcMock.listDatasets(PROJECT, EMPTY_RPC_OPTIONS)).andReturn(result); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Page<Dataset> page = bigquery.listDatasets(); assertNull(page.getNextPageToken()); assertArrayEquals( ImmutableList.of().toArray(), Iterables.toArray(page.getValues(), Dataset.class)); }
@Test public void testGetQueryResults() { JobId queryJob = JobId.of(JOB); GetQueryResultsResponse responsePb = new GetQueryResultsResponse() .setEtag("etag") .setJobReference(queryJob.toPb()) .setRows(ImmutableList.of(TABLE_ROW)) .setJobComplete(true) .setCacheHit(false) .setPageToken(CURSOR) .setTotalBytesProcessed(42L) .setTotalRows(BigInteger.valueOf(1L)); EasyMock.expect(bigqueryRpcMock.getQueryResults(PROJECT, JOB, null, EMPTY_RPC_OPTIONS)) .andReturn(responsePb); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); QueryResponse response = bigquery.getQueryResults(queryJob); assertEquals(true, response.getCompleted()); assertEquals(null, response.getSchema()); }
Map<String, Object> row1 = ImmutableMap.<String, Object>of("field", "value1"); Map<String, Object> row2 = ImmutableMap.<String, Object>of("field", "value2"); List<RowToInsert> rows = ImmutableList.of(RowToInsert.of(row1), RowToInsert.of(row2)); InsertAllRequest request = InsertAllRequest.newBuilder(TABLE_ID) .setIgnoreUnknownValues(true) .setTemplateSuffix("suffix"); EasyMock.expect(bigqueryRpcMock.insertAll(PROJECT, DATASET, TABLE, requestPb)) .andThrow(new BigQueryException(500, "InternalError")); EasyMock.replay(bigqueryRpcMock); bigquery = options .toBuilder() .setRetrySettings(ServiceOptions.getDefaultRetrySettings()) .build() .getService(); thrown.expect(BigQueryException.class); bigquery.insertAll(request);
new GetQueryResultsResponse() .setJobReference(queryJob.toPb()) .setRows(ImmutableList.of(TABLE_ROW)) .setJobComplete(true) .setCacheHit(false) .setSchema(TABLE_SCHEMA.toPb()); EasyMock.expect( bigqueryRpcMock.create( JOB_INFO.toPb(), Collections.<BigQueryRpc.Option, Object>emptyMap())) .andReturn(jobResponsePb); EasyMock.expect( bigqueryRpcMock.getQueryResults( PROJECT, JOB, null, BigQueryImpl.optionMap(Job.DEFAULT_QUERY_WAIT_OPTIONS))) .andReturn(responsePb); EasyMock.expect( bigqueryRpcMock.listTableData( PROJECT, DATASET, TABLE, Collections.<BigQueryRpc.Option, Object>emptyMap())) .andReturn( new TableDataList() .setPageToken("") .setRows(ImmutableList.of(TABLE_ROW)) .setTotalRows(1L));
Map<String, Object> row2 = ImmutableMap.<String, Object>of("field", "value2"); List<RowToInsert> rows = ImmutableList.of(new RowToInsert("row1", row1), new RowToInsert("row2", row2)); TableId tableId = TableId.of(OTHER_PROJECT, DATASET, TABLE); InsertAllRequest request = new TableDataInsertAllResponse() .setInsertErrors( ImmutableList.of( new TableDataInsertAllResponse.InsertErrors() .setIndex(0L) .setErrors(ImmutableList.of(new ErrorProto().setMessage("ErrorMessage"))))); EasyMock.expect(bigqueryRpcMock.insertAll(OTHER_PROJECT, DATASET, TABLE, requestPb)) .andReturn(responsePb); EasyMock.replay(bigqueryRpcMock); BigQueryOptions bigQueryOptions = createBigQueryOptionsForProject(OTHER_PROJECT, rpcFactoryMock); bigquery = bigQueryOptions.getService(); InsertAllResponse response = bigquery.insertAll(request); assertNotNull(response.getErrorsFor(0L)); assertNull(response.getErrorsFor(1L));
.setSchema(TABLE_SCHEMA.toPb()); EasyMock.expect( bigqueryRpcMock.create( JOB_INFO.toPb(), Collections.<BigQueryRpc.Option, Object>emptyMap())) .andReturn(jobResponsePb1); EasyMock.expect(bigqueryRpcMock.getJob(eq(PROJECT), eq(JOB), anyString(), anyObject(Map.class))) .andReturn(jobResponsePb1); bigqueryRpcMock.getQueryResults( PROJECT, JOB, null, BigQueryImpl.optionMap(Job.DEFAULT_QUERY_WAIT_OPTIONS))) .andReturn(responsePb1); EasyMock.expect( bigqueryRpcMock.getQueryResults( bigqueryRpcMock.listTableData( .setPageToken("") .setRows(ImmutableList.of(TABLE_ROW)) .setTotalRows(1L)); bigquery = options.getService(); TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY, queryJob); assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA); assertThat(result.getTotalRows()).isEqualTo(1);
@Test public void testCreateJobSuccess() { String id = "testCreateJobSuccess-id"; JobId jobId = JobId.of(id); String query = "SELECT * in FOO"; Capture<com.google.api.services.bigquery.model.Job> jobCapture = EasyMock.newCapture(); EasyMock.expect( bigqueryRpcMock.create(EasyMock.capture(jobCapture), EasyMock.eq(EMPTY_RPC_OPTIONS))) .andReturn(newJobPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); assertThat(bigquery.create(JobInfo.of(jobId, QueryJobConfiguration.of(query)))).isNotNull(); assertThat(jobCapture.getValue().getJobReference().getJobId()).isEqualTo(id); }
@Test public void testGetJob() { EasyMock.expect(bigqueryRpcMock.getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS)) .andReturn(COMPLETE_COPY_JOB.toPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Job job = bigquery.getJob(JOB); assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_COPY_JOB)), job); }
@Test public void testGetTable() { EasyMock.expect(bigqueryRpcMock.getTable(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS)) .andReturn(TABLE_INFO_WITH_PROJECT.toPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Table table = bigquery.getTable(DATASET, TABLE); assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table); }
@Test public void testDeleteTable() { EasyMock.expect(bigqueryRpcMock.deleteTable(PROJECT, DATASET, TABLE)).andReturn(true); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); assertTrue(bigquery.delete(DATASET, TABLE)); }
@Test public void testRuntimeException() { String exceptionMessage = "Artificial runtime exception"; EasyMock.expect(bigqueryRpcMock.getDataset(PROJECT, DATASET, EMPTY_RPC_OPTIONS)) .andThrow(new RuntimeException(exceptionMessage)); EasyMock.replay(bigqueryRpcMock); bigquery = options .toBuilder() .setRetrySettings(ServiceOptions.getDefaultRetrySettings()) .build() .getService(); thrown.expect(BigQueryException.class); thrown.expectMessage(exceptionMessage); bigquery.getDataset(DATASET); }
@Test public void testCreateJobTryGet() { final String id = "testCreateJobTryGet-id"; String query = "SELECT * in FOO"; Supplier<JobId> idProvider = new Supplier<JobId>() { @Override public JobId get() { return JobId.of(id); } }; Capture<com.google.api.services.bigquery.model.Job> jobCapture = EasyMock.newCapture(); EasyMock.expect( bigqueryRpcMock.create(EasyMock.capture(jobCapture), EasyMock.eq(EMPTY_RPC_OPTIONS))) .andThrow(new BigQueryException(409, "already exists, for some reason")); EasyMock.expect( bigqueryRpcMock.getJob( anyString(), EasyMock.eq(id), EasyMock.eq((String) null), EasyMock.eq(EMPTY_RPC_OPTIONS))) .andReturn(newJobPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); ((BigQueryImpl) bigquery).create(JobInfo.of(QueryJobConfiguration.of(query)), idProvider); assertThat(jobCapture.getValue().getJobReference().getJobId()).isEqualTo(id); }
@Test public void testCreateDataset() { DatasetInfo datasetInfo = DATASET_INFO.setProjectId(OTHER_PROJECT); EasyMock.expect(bigqueryRpcMock.create(datasetInfo.toPb(), EMPTY_RPC_OPTIONS)) .andReturn(datasetInfo.toPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Dataset dataset = bigquery.create(datasetInfo); assertEquals(new Dataset(bigquery, new DatasetInfo.BuilderImpl(datasetInfo)), dataset); }
@Test public void testGetDataset() { EasyMock.expect(bigqueryRpcMock.getDataset(PROJECT, DATASET, EMPTY_RPC_OPTIONS)) .andReturn(DATASET_INFO_WITH_PROJECT.toPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Dataset dataset = bigquery.getDataset(DATASET); assertEquals( new Dataset(bigquery, new DatasetInfo.BuilderImpl(DATASET_INFO_WITH_PROJECT)), dataset); }
@Test public void testCreateJobWithSelectedFields() { Capture<Map<BigQueryRpc.Option, Object>> capturedOptions = Capture.newInstance(); EasyMock.expect( bigqueryRpcMock.create( EasyMock.anyObject(com.google.api.services.bigquery.model.Job.class), EasyMock.capture(capturedOptions))) .andReturn(newJobPb()); EasyMock.replay(bigqueryRpcMock); BigQuery.JobOption jobOptions = BigQuery.JobOption.fields(BigQuery.JobField.USER_EMAIL); bigquery = options.getService(); bigquery.create(JobInfo.of(QueryJobConfiguration.of("SOME QUERY")), jobOptions); String selector = (String) capturedOptions.getValue().get(jobOptions.getRpcOption()); // jobReference and configuration are always sent; the RPC call won't succeed otherwise. assertThat(selector.split(",")) .asList() .containsExactly("jobReference", "configuration", "user_email"); }
@Test public void testGetDatasetWithSelectedFields() { Capture<Map<BigQueryRpc.Option, Object>> capturedOptions = Capture.newInstance(); EasyMock.expect(bigqueryRpcMock.getDataset(eq(PROJECT), eq(DATASET), capture(capturedOptions))) .andReturn(DATASET_INFO_WITH_PROJECT.toPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Dataset dataset = bigquery.getDataset(DATASET, DATASET_OPTION_FIELDS); String selector = (String) capturedOptions.getValue().get(DATASET_OPTION_FIELDS.getRpcOption()); assertTrue(selector.contains("datasetReference")); assertTrue(selector.contains("access")); assertTrue(selector.contains("etag")); assertEquals(28, selector.length()); assertEquals( new Dataset(bigquery, new DatasetInfo.BuilderImpl(DATASET_INFO_WITH_PROJECT)), dataset); }
@Test public void testGetTableWithSelectedFields() { Capture<Map<BigQueryRpc.Option, Object>> capturedOptions = Capture.newInstance(); EasyMock.expect( bigqueryRpcMock.getTable(eq(PROJECT), eq(DATASET), eq(TABLE), capture(capturedOptions))) .andReturn(TABLE_INFO_WITH_PROJECT.toPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Table table = bigquery.getTable(TABLE_ID, TABLE_OPTION_FIELDS); String selector = (String) capturedOptions.getValue().get(TABLE_OPTION_FIELDS.getRpcOption()); assertTrue(selector.contains("tableReference")); assertTrue(selector.contains("schema")); assertTrue(selector.contains("etag")); assertEquals(31, selector.length()); assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table); }