Search in sources :

Example 1 with DataflowPackage

use of com.google.api.services.dataflow.model.DataflowPackage in project beam by apache.

the class PackageUtilTest method testStagingPreservesClasspath.

@Test
public void testStagingPreservesClasspath() throws Exception {
    File smallFile = makeFileWithContents("small.txt", "small");
    File largeFile = makeFileWithContents("large.txt", "large contents");
    when(mockGcsUtil.getObjects(anyListOf(GcsPath.class))).thenReturn(ImmutableList.of(StorageObjectOrIOException.create(new FileNotFoundException("some/path"))));
    when(mockGcsUtil.create(any(GcsPath.class), anyString())).thenAnswer(new Answer<SinkChannel>() {

        @Override
        public SinkChannel answer(InvocationOnMock invocation) throws Throwable {
            return Pipe.open().sink();
        }
    });
    List<DataflowPackage> targets = PackageUtil.stageClasspathElements(ImmutableList.of(smallFile.getAbsolutePath(), largeFile.getAbsolutePath()), STAGING_PATH, createOptions);
    // Verify that the packages are returned small, then large, matching input order even though
    // the large file would be uploaded first.
    assertThat(targets.get(0).getName(), startsWith("small"));
    assertThat(targets.get(1).getName(), startsWith("large"));
}
Also used : InvocationOnMock(org.mockito.invocation.InvocationOnMock) GcsPath(org.apache.beam.sdk.util.gcsfs.GcsPath) FileNotFoundException(java.io.FileNotFoundException) File(java.io.File) SinkChannel(java.nio.channels.Pipe.SinkChannel) DataflowPackage(com.google.api.services.dataflow.model.DataflowPackage) Test(org.junit.Test)

Example 2 with DataflowPackage

use of com.google.api.services.dataflow.model.DataflowPackage in project beam by apache.

the class PackageUtilTest method testPackageNamingWithFilesHavingSameContentsButDifferentNames.

@Test
public void testPackageNamingWithFilesHavingSameContentsButDifferentNames() throws Exception {
    File tmpDirectory1 = tmpFolder.newFolder("folder1", "folderA");
    makeFileWithContents("folder1/folderA/uniqueName1", "This is a test!");
    DataflowPackage target1 = makePackageAttributes(tmpDirectory1, null).getDataflowPackage();
    File tmpDirectory2 = tmpFolder.newFolder("folder2", "folderA");
    makeFileWithContents("folder2/folderA/uniqueName2", "This is a test!");
    DataflowPackage target2 = makePackageAttributes(tmpDirectory2, null).getDataflowPackage();
    assertNotEquals(target1.getName(), target2.getName());
    assertNotEquals(target1.getLocation(), target2.getLocation());
}
Also used : File(java.io.File) DataflowPackage(com.google.api.services.dataflow.model.DataflowPackage) Test(org.junit.Test)

Example 3 with DataflowPackage

use of com.google.api.services.dataflow.model.DataflowPackage in project beam by apache.

the class PackageUtilTest method testPackageUploadWithEmptyDirectorySucceeds.

@Test
public void testPackageUploadWithEmptyDirectorySucceeds() throws Exception {
    Pipe pipe = Pipe.open();
    File tmpDirectory = tmpFolder.newFolder("folder");
    when(mockGcsUtil.getObjects(anyListOf(GcsPath.class))).thenReturn(ImmutableList.of(StorageObjectOrIOException.create(new FileNotFoundException("some/path"))));
    when(mockGcsUtil.create(any(GcsPath.class), anyString())).thenReturn(pipe.sink());
    List<DataflowPackage> targets = PackageUtil.stageClasspathElements(ImmutableList.of(tmpDirectory.getAbsolutePath()), STAGING_PATH, createOptions);
    DataflowPackage target = Iterables.getOnlyElement(targets);
    verify(mockGcsUtil).getObjects(anyListOf(GcsPath.class));
    verify(mockGcsUtil).create(any(GcsPath.class), anyString());
    verifyNoMoreInteractions(mockGcsUtil);
    assertThat(target.getName(), RegexMatcher.matches("folder-" + HASH_PATTERN + ".jar"));
    assertThat(target.getLocation(), equalTo(STAGING_PATH + target.getName()));
    assertNull(new ZipInputStream(Channels.newInputStream(pipe.source())).getNextEntry());
}
Also used : ZipInputStream(java.util.zip.ZipInputStream) GcsPath(org.apache.beam.sdk.util.gcsfs.GcsPath) FileNotFoundException(java.io.FileNotFoundException) Pipe(java.nio.channels.Pipe) File(java.io.File) DataflowPackage(com.google.api.services.dataflow.model.DataflowPackage) Test(org.junit.Test)

Example 4 with DataflowPackage

use of com.google.api.services.dataflow.model.DataflowPackage in project beam by apache.

the class PackageUtilTest method testPackageUploadWithExplicitPackageName.

@Test
public void testPackageUploadWithExplicitPackageName() throws Exception {
    Pipe pipe = Pipe.open();
    File tmpFile = makeFileWithContents("file.txt", "This is a test!");
    final String overriddenName = "alias.txt";
    when(mockGcsUtil.getObjects(anyListOf(GcsPath.class))).thenReturn(ImmutableList.of(StorageObjectOrIOException.create(new FileNotFoundException("some/path"))));
    when(mockGcsUtil.create(any(GcsPath.class), anyString())).thenReturn(pipe.sink());
    List<DataflowPackage> targets = PackageUtil.stageClasspathElements(ImmutableList.of(overriddenName + "=" + tmpFile.getAbsolutePath()), STAGING_PATH, createOptions);
    DataflowPackage target = Iterables.getOnlyElement(targets);
    verify(mockGcsUtil).getObjects(anyListOf(GcsPath.class));
    verify(mockGcsUtil).create(any(GcsPath.class), anyString());
    verifyNoMoreInteractions(mockGcsUtil);
    assertThat(target.getName(), equalTo(overriddenName));
    assertThat(target.getLocation(), RegexMatcher.matches(STAGING_PATH + "file-" + HASH_PATTERN + ".txt"));
}
Also used : GcsPath(org.apache.beam.sdk.util.gcsfs.GcsPath) FileNotFoundException(java.io.FileNotFoundException) Pipe(java.nio.channels.Pipe) Matchers.anyString(org.mockito.Matchers.anyString) File(java.io.File) DataflowPackage(com.google.api.services.dataflow.model.DataflowPackage) Test(org.junit.Test)

Example 5 with DataflowPackage

use of com.google.api.services.dataflow.model.DataflowPackage in project beam by apache.

the class PackageUtilTest method testPackageUploadWithFileSucceeds.

@Test
public void testPackageUploadWithFileSucceeds() throws Exception {
    Pipe pipe = Pipe.open();
    String contents = "This is a test!";
    File tmpFile = makeFileWithContents("file.txt", contents);
    when(mockGcsUtil.getObjects(anyListOf(GcsPath.class))).thenReturn(ImmutableList.of(StorageObjectOrIOException.create(new FileNotFoundException("some/path"))));
    when(mockGcsUtil.create(any(GcsPath.class), anyString())).thenReturn(pipe.sink());
    List<DataflowPackage> targets = PackageUtil.stageClasspathElements(ImmutableList.of(tmpFile.getAbsolutePath()), STAGING_PATH, createOptions);
    DataflowPackage target = Iterables.getOnlyElement(targets);
    verify(mockGcsUtil).getObjects(anyListOf(GcsPath.class));
    verify(mockGcsUtil).create(any(GcsPath.class), anyString());
    verifyNoMoreInteractions(mockGcsUtil);
    assertThat(target.getName(), RegexMatcher.matches("file-" + HASH_PATTERN + ".txt"));
    assertThat(target.getLocation(), equalTo(STAGING_PATH + target.getName()));
    assertThat(new LineReader(Channels.newReader(pipe.source(), "UTF-8")).readLine(), equalTo(contents));
}
Also used : LineReader(com.google.common.io.LineReader) GcsPath(org.apache.beam.sdk.util.gcsfs.GcsPath) FileNotFoundException(java.io.FileNotFoundException) Pipe(java.nio.channels.Pipe) Matchers.anyString(org.mockito.Matchers.anyString) File(java.io.File) DataflowPackage(com.google.api.services.dataflow.model.DataflowPackage) Test(org.junit.Test)

Aggregations

DataflowPackage (com.google.api.services.dataflow.model.DataflowPackage)34 Test (org.junit.Test)30 Job (com.google.api.services.dataflow.model.Job)22 DataflowPipelineOptions (org.apache.beam.runners.dataflow.options.DataflowPipelineOptions)22 Pipeline (org.apache.beam.sdk.Pipeline)21 Structs.getString (org.apache.beam.runners.dataflow.util.Structs.getString)13 File (java.io.File)11 Step (com.google.api.services.dataflow.model.Step)8 ImmutableMap (com.google.common.collect.ImmutableMap)6 Map (java.util.Map)6 Structs.addObject (org.apache.beam.runners.dataflow.util.Structs.addObject)6 FileNotFoundException (java.io.FileNotFoundException)5 GcsPath (org.apache.beam.sdk.util.gcsfs.GcsPath)4 Matchers.anyString (org.mockito.Matchers.anyString)4 Pipe (java.nio.channels.Pipe)3 LinkedList (java.util.LinkedList)3 WorkerPool (com.google.api.services.dataflow.model.WorkerPool)2 ImmutableList (com.google.common.collect.ImmutableList)2 IOException (java.io.IOException)2 ArrayList (java.util.ArrayList)2