org.apache.beam.sdk.options.ValueProvider.StaticValueProvider Java Examples

@Test
public void testBuildWrite() {
  SimpleSink<Void> sink = makeSimpleSink();
  WriteFiles<String, ?, String> write = WriteFiles.to(sink).withNumShards(3);
  assertThat((SimpleSink<Void>) write.getSink(), is(sink));
  PTransform<PCollection<String>, PCollectionView<Integer>> originalSharding =
      write.getComputeNumShards();

  assertThat(write.getComputeNumShards(), is(nullValue()));
  assertThat(write.getNumShardsProvider(), instanceOf(StaticValueProvider.class));
  assertThat(write.getNumShardsProvider().get(), equalTo(3));
  assertThat(write.getComputeNumShards(), equalTo(originalSharding));

  WriteFiles<String, ?, ?> write2 = write.withSharding(SHARDING_TRANSFORM);
  assertThat((SimpleSink<Void>) write2.getSink(), is(sink));
  assertThat(write2.getComputeNumShards(), equalTo(SHARDING_TRANSFORM));
  // original unchanged

  WriteFiles<String, ?, ?> writeUnsharded = write2.withRunnerDeterminedSharding();
  assertThat(writeUnsharded.getComputeNumShards(), nullValue());
  assertThat(write.getComputeNumShards(), equalTo(originalSharding));
}
 

/**
 * The windowed filename method will construct filenames per window according to the baseFile,
 * suffix, and shardTemplate supplied. Directories with date templates in them will automatically
 * have their values resolved. For example the outputDirectory of /YYYY/MM/DD would resolve to
 * /2017/01/08 on January 8th, 2017.
 */
@Override
public ResourceId windowedFilename(
        int shardNumber,
        int numShards,
        BoundedWindow window,
        PaneInfo paneInfo,
        OutputFileHints outputFileHints) {

    ResourceId outputFile =
            resolveWithDateTemplates(outputDirectory, window)
                    .resolve(outputFilenamePrefix.get(), StandardResolveOptions.RESOLVE_FILE);

    DefaultFilenamePolicy policy =
            DefaultFilenamePolicy.fromStandardParameters(
                    StaticValueProvider.of(outputFile), shardTemplate.get(), suffix.get(), true);
    ResourceId result =
            policy.windowedFilename(shardNumber, numShards, window, paneInfo, outputFileHints);
    LOG.debug("Windowed file name policy created: {}", result.toString());
    return result;
}
 

@Test(timeout = 30000)
public void canSendPubsubOutput() throws Exception {
  final List<String> inputLines = Lines.resources("testdata/pubsub-integration/input.ndjson");

  pipeline.getOptions().as(DirectOptions.class).setBlockOnRun(false);

  SinkOptions.Parsed sinkOptions = pipeline.getOptions().as(SinkOptions.Parsed.class);
  sinkOptions.setOutput(pipeline.newProvider(topicName.toString()));
  // We would normally use pipeline.newProvider instead of StaticValueProvider in tests,
  // but something about this configuration causes the pipeline to stall when CompressPayload
  // accesses a method on the underlying enum value when defined via pipeline.newProvider.
  sinkOptions.setOutputPubsubCompression(StaticValueProvider.of(Compression.UNCOMPRESSED));

  pipeline.apply(Create.of(inputLines)).apply(InputFileFormat.json.decode())
      .apply(OutputType.pubsub.write(sinkOptions));

  final PipelineResult result = pipeline.run();

  System.err.println("Waiting for subscriber to receive messages published in the pipeline...");
  List<String> expectedLines = Lines.resources("testdata/pubsub-integration/truncated.ndjson");
  List<String> received = receiveLines(expectedLines.size());
  assertThat(received, matchesInAnyOrder(expectedLines));
  result.cancel();
}
 

@Test
public void testTableSourceCreateReader() throws Exception {
  BigQueryStorageTableSource<TableRow> tableSource =
      BigQueryStorageTableSource.create(
          ValueProvider.StaticValueProvider.of(
              BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table")),
          null,
          null,
          null,
          new TableRowParser(),
          TableRowJsonCoder.of(),
          new FakeBigQueryServices().withDatasetService(fakeDatasetService));

  thrown.expect(UnsupportedOperationException.class);
  thrown.expectMessage("BigQuery storage source must be split before reading");
  tableSource.createReader(options);
}
 

@Test
public void testTableSourceEstimatedSize_WithDefaultProject() throws Exception {
  fakeDatasetService.createDataset("project-id", "dataset", "", "", null);
  TableReference tableRef = BigQueryHelpers.parseTableSpec("project-id:dataset.table");
  Table table = new Table().setTableReference(tableRef).setNumBytes(100L);
  fakeDatasetService.createTable(table);

  BigQueryStorageTableSource<TableRow> tableSource =
      BigQueryStorageTableSource.create(
          ValueProvider.StaticValueProvider.of(BigQueryHelpers.parseTableSpec("dataset.table")),
          null,
          null,
          null,
          new TableRowParser(),
          TableRowJsonCoder.of(),
          new FakeBigQueryServices().withDatasetService(fakeDatasetService));

  assertEquals(100, tableSource.getEstimatedSizeBytes(options));
}
 

@Test
public void testDestinationFunction_generatesProperFileParams() {
  SerializableFunction<BillingEvent, Params> destinationFunction =
      InvoicingUtils.makeDestinationFunction("my/directory", StaticValueProvider.of("2017-10"));

  BillingEvent billingEvent = mock(BillingEvent.class);
  // We mock BillingEvent to make the test independent of the implementation of toFilename()
  when(billingEvent.toFilename(any())).thenReturn("invoice_details_2017-10_registrar_tld");

  assertThat(destinationFunction.apply(billingEvent))
      .isEqualTo(
          new Params()
              .withShardTemplate("")
              .withSuffix(".csv")
              .withBaseFilename(
                  FileBasedSink.convertToFileResourceIfPossible(
                      "my/directory/2017-10/invoice_details_2017-10_registrar_tld")));
}
 

/** Tests the {@link BulkCompressor.Compressor} performs compression properly. */
@Test
public void testCompressFile() throws Exception {
  // Setup test
  final Compression compression = Compression.GZIP;

  final ValueProvider<String> outputDirectoryProvider =
      pipeline.newProvider(tempFolderCompressedPath.toString());

  final ValueProvider<Compression> compressionProvider = StaticValueProvider.of(compression);

  final Metadata metadata = FileSystems.matchSingleFileSpec(textFile.toString());

  // Execute the compressor
  PCollection<String> lines = pipeline
      .apply("Create File Input", Create.of(metadata))
      .apply("Compress", ParDo.of(new Compressor(outputDirectoryProvider, compressionProvider)))
      .apply("Read the Files", TextIO.readAll().withCompression(Compression.AUTO));

  // Test the result
  PAssert.that(lines).containsInAnyOrder(FILE_CONTENT);
  pipeline.run();
}
 

/**
 * Test {@link TransformTextViaJavascript} allows for script to disregard some data by returning
 * null / undefined.
 */
@Test
@Category(NeedsRunner.class)
public void testDoFnDisregards() {
  List<String> inJson = Arrays.asList("{\"answerToLife\":42}", "{\"answerToLife\":43}");
  List<String> expectedJson = Arrays.asList("{\"answerToLife\":42}");

  PCollection<String> transformedJson =
      pipeline
          .apply("Create", Create.of(inJson))
          .apply(
              TransformTextViaJavascript.newBuilder()
                  .setFileSystemPath(StaticValueProvider.of(TRANSFORM_FILE_PATH))
                  .setFunctionName(StaticValueProvider.of("transformWithFilter"))
                  .build());

  PAssert.that(transformedJson).containsInAnyOrder(expectedJson);

  pipeline.run();
}
 

/**
 * Test building {@link SplunkEventWriter} with custom batchcount and certificate validation .
 */
@Test
public void eventWriterCustomBatchCountAndValidation() {

  Integer batchCount = 30;
  Boolean certificateValidation = false;
  SplunkEventWriter writer =
      SplunkEventWriter.newBuilder()
          .withUrl("test-url")
          .withToken("test-token")
          .withInputBatchCount(StaticValueProvider.of(batchCount))
          .withDisableCertificateValidation(StaticValueProvider.of(certificateValidation))
          .build();

  assertThat(writer.inputBatchCount().get()).isEqualTo(batchCount);
  assertThat(writer.disableCertificateValidation().get()).isEqualTo(certificateValidation);
}
 

@Test
public void testNestedValueProviderStatic() throws Exception {
  ValueProvider<String> xvp = StaticValueProvider.of("foo");
  ValueProvider<Integer> yvp = StaticValueProvider.of(1);
  ValueProvider<String> zvp =
      DualInputNestedValueProvider.of(
          xvp,
          yvp,
          new SerializableFunction<TranslatorInput<String, Integer>, String>() {
            @Override
            public String apply(TranslatorInput<String, Integer> input) {
              return input.getX() + (input.getY() + 1);
            }
          });
  assertTrue(zvp.isAccessible());
  assertEquals("foo2", zvp.get());
}
 

@Test
public void testCreateOnlyOnce() {
  SpannerConfig config1 =
      SpannerConfig.create()
          .toBuilder()
          .setServiceFactory(serviceFactory)
          .setProjectId(StaticValueProvider.of("project"))
          .setInstanceId(StaticValueProvider.of("test1"))
          .setDatabaseId(StaticValueProvider.of("test1"))
          .build();

  SpannerAccessor acc1 = SpannerAccessor.getOrCreate(config1);
  SpannerAccessor acc2 = SpannerAccessor.getOrCreate(config1);
  SpannerAccessor acc3 = SpannerAccessor.getOrCreate(config1);

  acc1.close();
  acc2.close();
  acc3.close();

  // getDatabaseClient and close() only called once.
  verify(serviceFactory.mockSpanner(), times(1))
      .getDatabaseClient(DatabaseId.of("project", "test1", "test1"));
  verify(serviceFactory.mockSpanner(), times(1)).close();
}
 

@Override
public PDone expand(PCollection<PubsubMessage> input) {
  List<Destination> destinations = baseOptions.getPerNamespaceDestinations().entrySet().stream()
      .map(entry -> new Destination(entry.getKey(), entry.getValue()))
      .collect(Collectors.toList());
  int numDestinations = destinations.size();
  int numPartitions = numDestinations + 1;
  PCollectionList<PubsubMessage> partitioned = input.apply("PartitionByNamespace",
      Partition.of(numPartitions, new PartitionFn(destinations)));

  for (int i = 0; i < numDestinations; i++) {
    Destination destination = destinations.get(i);
    RepublisherOptions.Parsed opts = baseOptions.as(RepublisherOptions.Parsed.class);
    opts.setOutput(StaticValueProvider.of(destination.dest));
    String name = String.join("_", "republish", destination.namespace);
    partitioned.get(i).apply(name, opts.getOutputType().write(opts));
  }

  return PDone.in(input.getPipeline());
}
 

private void doTableSourceEstimatedSizeTest(boolean useStreamingBuffer) throws Exception {
  fakeDatasetService.createDataset("foo.com:project", "dataset", "", "", null);
  TableReference tableRef = BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table");
  Table table = new Table().setTableReference(tableRef).setNumBytes(100L);
  if (useStreamingBuffer) {
    table.setStreamingBuffer(new Streamingbuffer().setEstimatedBytes(BigInteger.TEN));
  }

  fakeDatasetService.createTable(table);

  BigQueryStorageTableSource<TableRow> tableSource =
      BigQueryStorageTableSource.create(
          ValueProvider.StaticValueProvider.of(tableRef),
          null,
          null,
          null,
          new TableRowParser(),
          TableRowJsonCoder.of(),
          new FakeBigQueryServices().withDatasetService(fakeDatasetService));

  assertEquals(100, tableSource.getEstimatedSizeBytes(options));
}
 

/**
 * Returns the table to write, or {@code null} if writing with {@code tableFunction}.
 *
 * <p>If the table's project is not specified, use the executing project.
 */
@Nullable
ValueProvider<TableReference> getTableWithDefaultProject(BigQueryOptions bqOptions) {
  ValueProvider<TableReference> table = getTable();
  if (table == null) {
    return table;
  }

  if (!table.isAccessible()) {
    LOG.info(
        "Using a dynamic value for table input. This must contain a project"
            + " in the table reference: {}",
        table);
    return table;
  }
  if (Strings.isNullOrEmpty(table.get().getProjectId())) {
    // If user does not specify a project we assume the table to be located in
    // the default project.
    TableReference tableRef = table.get();
    tableRef.setProjectId(bqOptions.getProject());
    return NestedValueProvider.of(
        StaticValueProvider.of(BigQueryHelpers.toJsonString(tableRef)),
        new JsonTableRefToTableRef());
  }
  return table;
}
 

private void setupOneMessage(Iterable<IncomingMessage> incoming) {
  now = new AtomicLong(REQ_TIME);
  clock = () -> now.get();
  factory = PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S, incoming);
  PubsubUnboundedSource source =
      new PubsubUnboundedSource(
          clock,
          factory,
          null,
          null,
          StaticValueProvider.of(SUBSCRIPTION),
          TIMESTAMP_ATTRIBUTE,
          ID_ATTRIBUTE,
          true /* needsAttributes */);
  primSource = new PubsubSource(source);
}
 

/**
 * This test ensures that protobuf creation and interactions with {@link BigtableDataClient} work
 * as expected.
 *
 * @throws IOException
 * @throws InterruptedException
 */
@Test
public void testRead() throws IOException {
  ByteKey start = ByteKey.copyFrom("a".getBytes(StandardCharsets.UTF_8));
  ByteKey end = ByteKey.copyFrom("b".getBytes(StandardCharsets.UTF_8));
  when(mockBigtableSource.getRanges()).thenReturn(Arrays.asList(ByteKeyRange.of(start, end)));
  when(mockBigtableSource.getTableId()).thenReturn(StaticValueProvider.of("table_name"));
  @SuppressWarnings("unchecked")
  ResultScanner<Row> mockResultScanner = Mockito.mock(ResultScanner.class);
  Row expectedRow = Row.newBuilder().setKey(ByteString.copyFromUtf8("a")).build();
  when(mockResultScanner.next()).thenReturn(expectedRow).thenReturn(null);
  when(mockBigtableDataClient.readRows(any(ReadRowsRequest.class))).thenReturn(mockResultScanner);
  BigtableService.Reader underTest =
      new BigtableServiceImpl.BigtableReaderImpl(mockSession, mockBigtableSource);

  underTest.start();
  Assert.assertEquals(expectedRow, underTest.getCurrentRow());
  Assert.assertFalse(underTest.advance());
  underTest.close();

  verify(mockResultScanner, times(1)).close();
}
 

public static PTransform<PCollection<String>, WriteFilesResult<Void>> createWrite(
    String filenamePrefix, String filenameSuffix, Schema schema, JdbcAvroArgs jdbcAvroArgs) {
  filenamePrefix = filenamePrefix.replaceAll("/+$", "") + "/part";
  ValueProvider<ResourceId> prefixProvider =
      StaticValueProvider.of(FileBasedSink.convertToFileResourceIfPossible(filenamePrefix));
  FileBasedSink.FilenamePolicy filenamePolicy =
      DefaultFilenamePolicy.fromStandardParameters(
          prefixProvider, DEFAULT_SHARD_TEMPLATE, filenameSuffix, false);

  final DynamicAvroDestinations<String, Void, String> destinations =
      AvroIO.constantDestinations(
          filenamePolicy,
          schema,
          ImmutableMap.of(),
          // since Beam does not support zstandard
          CodecFactory.nullCodec(),
          SerializableFunctions.identity());
  final FileBasedSink<String, Void, String> sink =
      new JdbcAvroSink<>(prefixProvider, destinations, jdbcAvroArgs);
  return WriteFiles.to(sink);
}
 

/**
 * Tests that windowedFilename() constructs the filename correctly according to the parameters
 * when the filename suffix is a null value.
 */
@Test
public void testWindowedFilenameNullSuffix() throws IOException {
  // Arrange
  //
  ResourceId outputDirectory = getBaseTempDirectory();
  WindowedContext context = mock(WindowedContext.class);
  BoundedWindow window = mock(BoundedWindow.class);
  PaneInfo paneInfo = PaneInfo.createPane(false, true, Timing.ON_TIME, 0, 0);
  WindowedFilenamePolicy policy =
      new WindowedFilenamePolicy(
          StaticValueProvider.of(outputDirectory.toString()),
          StaticValueProvider.of("output"),
          StaticValueProvider.of("-SSS-of-NNN"),
          StaticValueProvider.of(null));

  // Act
  //
  ResourceId filename =
      policy.windowedFilename(1, 1, window, paneInfo, new TestOutputFileHints());
  // Assert
  //
  assertThat(filename, is(notNullValue()));
  assertThat(filename.getFilename(), is(equalTo("output-001-of-001")));
}
 

/**
 * Tests that windowedFilename() constructs the filename correctly according to the parameters
 * when using ValueProviders.
 */
@Test
public void testWindowedFilenameFormatValueProvider() throws IOException {
  // Arrange
  //
  ResourceId outputDirectory = getBaseTempDirectory();
  WindowedContext context = mock(WindowedContext.class);
  BoundedWindow window = mock(BoundedWindow.class);
  PaneInfo paneInfo = PaneInfo.createPane(false, true, Timing.ON_TIME, 0, 0);
  WindowedFilenamePolicy policy =
      new WindowedFilenamePolicy(
          StaticValueProvider.of(outputDirectory.toString()),
          StaticValueProvider.of("output"),
          StaticValueProvider.of("-SSS-of-NNN"),
          StaticValueProvider.of(".txt"));

  // Act
  //
  ResourceId filename =
      policy.windowedFilename(1, 1, window, paneInfo, new TestOutputFileHints());

  // Assert
  //
  assertThat(filename, is(notNullValue()));
  assertThat(filename.getFilename(), is(equalTo("output-001-of-001.txt")));
}
 

private void datastoreWriterFnTest(int numMutations) throws Exception {
  // Create the requested number of mutations.
  List<Mutation> mutations = new ArrayList<>(numMutations);
  for (int i = 0; i < numMutations; ++i) {
    mutations.add(
        makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
  }

  DatastoreWriterFn datastoreWriter =
      new DatastoreWriterFn(
          StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher());
  DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  doFnTester.processBundle(mutations);

  int start = 0;
  while (start < numMutations) {
    int end = Math.min(numMutations, start + DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START);
    CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
    commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
    commitRequest.addAllMutations(mutations.subList(start, end));
    // Verify all the batch requests were made with the expected mutations.
    verify(mockDatastore, times(1)).commit(commitRequest.build());
    start = end;
  }
}
 

/** Like {@link #run} but with the given potentially modified options. */
@Override
public PipelineResult run(PipelineOptions options) {
  checkState(
      enforcement.isPresent(),
      "Is your TestPipeline declaration missing a @Rule annotation? Usage: "
          + "@Rule public final transient TestPipeline pipeline = TestPipeline.create();");

  final PipelineResult pipelineResult;
  try {
    enforcement.get().beforePipelineExecution();
    PipelineOptions updatedOptions =
        MAPPER.convertValue(MAPPER.valueToTree(options), PipelineOptions.class);
    updatedOptions
        .as(TestValueProviderOptions.class)
        .setProviderRuntimeValues(StaticValueProvider.of(providerRuntimeValues));
    pipelineResult = super.run(updatedOptions);
    verifyPAssertsSucceeded(this, pipelineResult);
  } catch (RuntimeException exc) {
    Throwable cause = exc.getCause();
    if (cause instanceof AssertionError) {
      throw (AssertionError) cause;
    } else {
      throw exc;
    }
  }

  // If we reach this point, the pipeline has been run and no exceptions have been thrown during
  // its execution.
  enforcement.get().afterPipelineExecution();
  return pipelineResult;
}
 

/**
 * Instantiates a new Import.
 *
 * @param fhirStore the fhir store
 * @param tempGcsPath the temp gcs path
 * @param deadLetterGcsPath the dead letter gcs path
 * @param contentStructure the content structure
 */
Import(
    String fhirStore,
    String tempGcsPath,
    String deadLetterGcsPath,
    @Nullable ContentStructure contentStructure) {
  this.fhirStore = StaticValueProvider.of(fhirStore);
  this.tempGcsPath = StaticValueProvider.of(tempGcsPath);
  this.deadLetterGcsPath = StaticValueProvider.of(deadLetterGcsPath);
  if (contentStructure == null) {
    this.contentStructure = ContentStructure.CONTENT_STRUCTURE_UNSPECIFIED;
  } else {
    this.contentStructure = contentStructure;
  }
}
 

/**
 * Disables ssl certificate validation.
 *
 * @param disableCertificateValidation for disabling certificate validation
 */
public Write withDisableCertificateValidation(Boolean disableCertificateValidation) {
  checkArgument(
      disableCertificateValidation != null,
      "withDisableCertificateValidation(disableCertificateValidation) called with null input.");
  return toBuilder()
      .setDisableCertificateValidation(StaticValueProvider.of((disableCertificateValidation)))
      .build();
}
 

@Test
public void noSubscriptionNoSplitGeneratesSubscription() throws Exception {
  TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
  factory = PubsubTestClient.createFactoryForCreateSubscription();
  PubsubUnboundedSource source =
      new PubsubUnboundedSource(
          factory,
          StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")),
          StaticValueProvider.of(topicPath),
          null /* subscription */,
          null /* timestampLabel */,
          null /* idLabel */,
          false /* needsAttributes */);
  assertThat(source.getSubscription(), nullValue());

  assertThat(source.getSubscription(), nullValue());

  PipelineOptions options = PipelineOptionsFactory.create();
  PubsubSource actualSource = new PubsubSource(source);
  PubsubReader reader = actualSource.createReader(options, null);
  SubscriptionPath createdSubscription = reader.subscription;
  assertThat(createdSubscription, not(nullValue()));

  PubsubCheckpoint checkpoint = reader.getCheckpointMark();
  assertThat(checkpoint.subscriptionPath, equalTo(createdSubscription.getPath()));

  checkpoint.finalizeCheckpoint();
  PubsubCheckpoint deserCheckpoint =
      CoderUtils.clone(actualSource.getCheckpointMarkCoder(), checkpoint);
  assertThat(checkpoint.subscriptionPath, not(nullValue()));
  assertThat(checkpoint.subscriptionPath, equalTo(deserCheckpoint.subscriptionPath));

  PubsubReader readerFromOriginal = actualSource.createReader(options, checkpoint);
  PubsubReader readerFromDeser = actualSource.createReader(options, deserCheckpoint);

  assertThat(readerFromOriginal.subscription, equalTo(createdSubscription));
  assertThat(readerFromDeser.subscription, equalTo(createdSubscription));
}
 

@Test
public void testReadSubscriptionDisplayData() {
  String subscription = "projects/project/subscriptions/subscription";
  PubsubIO.Read<String> read =
      PubsubIO.readStrings()
          .fromSubscription(StaticValueProvider.of(subscription))
          .withTimestampAttribute("myTimestamp")
          .withIdAttribute("myId");

  DisplayData displayData = DisplayData.from(read);

  assertThat(displayData, hasDisplayItem("subscription", subscription));
  assertThat(displayData, hasDisplayItem("timestampAttribute", "myTimestamp"));
  assertThat(displayData, hasDisplayItem("idAttribute", "myId"));
}
 

@Test(timeout = 30000)
public void canSendPubsubErrorOutput() throws Exception {
  final List<String> inputLines = Lines
      .resources("testdata/pubsub-integration/error-input.ndjson");

  pipeline.getOptions().as(DirectOptions.class).setBlockOnRun(false);

  SinkOptions.Parsed sinkOptions = pipeline.getOptions().as(SinkOptions.Parsed.class);
  sinkOptions.setInput(pipeline.newProvider("test input"));
  sinkOptions.setJobName("test job name");
  sinkOptions.setErrorOutput(pipeline.newProvider(topicName.toString()));
  // We would normally use pipeline.newProvider instead of StaticValueProvider in tests,
  // but something about this configuration causes the pipeline to stall when CompressPayload
  // accesses a method on the underlying enum value when defined via pipeline.newProvider.
  sinkOptions.setErrorOutputPubsubCompression(StaticValueProvider.of(Compression.UNCOMPRESSED));

  pipeline.apply(Create.of(inputLines)).apply(InputFileFormat.json.decode())
      .apply(ErrorOutputType.pubsub.write(sinkOptions));

  final PipelineResult result = pipeline.run();

  System.err.println("Waiting for subscriber to receive messages published in the pipeline...");
  List<String> expectedLines = Lines.resources("testdata/pubsub-integration/error-output.ndjson");
  List<String> received = receiveLines(expectedLines.size());
  assertThat(received, matchesInAnyOrder(expectedLines));
  result.cancel();
}
 

@Override
public Result expand(PCollection<String> input) {
  PCollection<HealthcareIOError<String>> failedBundles;
  PCollection<HealthcareIOError<String>> failedImports;
  switch (this.getWriteMethod()) {
    case IMPORT:
      LOG.warn(
          "Make sure the Cloud Healthcare Service Agent has permissions when using import:"
              + " https://cloud.google.com/healthcare/docs/how-tos/permissions-healthcare-api-gcp-products#fhir_store_cloud_storage_permissions");
      ValueProvider<String> deadPath =
          getImportGcsDeadLetterPath().orElseThrow(IllegalArgumentException::new);
      FhirIO.Import.ContentStructure contentStructure =
          getContentStructure().orElseThrow(IllegalArgumentException::new);
      ValueProvider<String> tempPath =
          getImportGcsTempPath()
              .orElse(
                  StaticValueProvider.of(input.getPipeline().getOptions().getTempLocation()));

      return input.apply(new Import(getFhirStore(), tempPath, deadPath, contentStructure));
    case EXECUTE_BUNDLE:
    default:
      failedBundles =
          input
              .apply(
                  "Execute FHIR Bundles",
                  ParDo.of(new ExecuteBundles.ExecuteBundlesFn(this.getFhirStore())))
              .setCoder(HealthcareIOErrorCoder.of(StringUtf8Coder.of()));
  }
  return Result.in(input.getPipeline(), failedBundles);
}
 

@Test
public void sendOneMessage() throws IOException {
  List<OutgoingMessage> outgoing =
      ImmutableList.of(
          OutgoingMessage.of(
              com.google.pubsub.v1.PubsubMessage.newBuilder()
                  .setData(ByteString.copyFromUtf8(DATA))
                  .putAllAttributes(ATTRIBUTES)
                  .build(),
              TIMESTAMP,
              getRecordId(DATA)));
  int batchSize = 1;
  int batchBytes = 1;
  try (PubsubTestClientFactory factory =
      PubsubTestClient.createFactoryForPublish(TOPIC, outgoing, ImmutableList.of())) {
    PubsubUnboundedSink sink =
        new PubsubUnboundedSink(
            factory,
            StaticValueProvider.of(TOPIC),
            TIMESTAMP_ATTRIBUTE,
            ID_ATTRIBUTE,
            NUM_SHARDS,
            batchSize,
            batchBytes,
            Duration.standardSeconds(2),
            RecordIdMethod.DETERMINISTIC);
    p.apply(Create.of(ImmutableList.of(DATA))).apply(ParDo.of(new Stamp(ATTRIBUTES))).apply(sink);
    p.run();
  }
  // The PubsubTestClientFactory will assert fail on close if the actual published
  // message does not match the expected publish message.
}
 

@Test
public void testValueProviderTopic() {
  StaticValueProvider<String> provider = StaticValueProvider.of("projects/project/topics/topic");
  Read<String> pubsubRead = PubsubIO.readStrings().fromTopic(provider);
  Pipeline.create().apply(pubsubRead);
  assertThat(pubsubRead.getTopicProvider(), not(nullValue()));
  assertThat(pubsubRead.getTopicProvider().isAccessible(), is(true));
  assertThat(pubsubRead.getTopicProvider().get().asPath(), equalTo(provider.get()));
}
 

/**
 * Tests that windowedFilename() produces the correct directory when the directory has dynamic
 * date variables.
 */
@Test
public void testWithDynamicDirectory() throws IOException {
  // Arrange
  //
  ResourceId outputDirectory = getBaseTempDirectory()
      .resolve("YYYY/MM/DD/HH:mm", StandardResolveOptions.RESOLVE_DIRECTORY);
  WindowedContext context = mock(WindowedContext.class);
  IntervalWindow window = mock(IntervalWindow.class);
  PaneInfo paneInfo = PaneInfo.createPane(false, true, Timing.ON_TIME, 0, 0);
  Instant windowBegin = new DateTime(2017, 1, 8, 10, 55, 0).toInstant();
  Instant windowEnd = new DateTime(2017, 1, 8, 10, 56, 0).toInstant();
  when(window.maxTimestamp()).thenReturn(windowEnd);
  when(window.start()).thenReturn(windowBegin);
  when(window.end()).thenReturn(windowEnd);

  WindowedFilenamePolicy policy =
      new WindowedFilenamePolicy(
          StaticValueProvider.of(outputDirectory.toString()),
          StaticValueProvider.of("output"),
          StaticValueProvider.of("-SSS-of-NNN"),
          StaticValueProvider.of(null));

  // Act
  //
  ResourceId filename =
      policy.windowedFilename(1, 1, window, paneInfo, new TestOutputFileHints());

  // Assert
  //
  assertThat(filename, is(notNullValue()));
  assertThat(
      filename.getCurrentDirectory().toString().endsWith("2017/01/08/10:56/"), is(equalTo(true)));
  assertThat(filename.getFilename(), is(equalTo("output-001-of-001")));
}