/** * Default constructor used for deserialization. */ public WorkUnitState() { this.workUnit = WorkUnit.createEmpty(); this.jobState = new State(); // Not available on deserialization this.taskBroker = null; }
@Override public void readFields(DataInput in) throws IOException { int numWorkUnits = in.readInt(); for (int i = 0; i < numWorkUnits; i++) { WorkUnit workUnit = WorkUnit.createEmpty(); workUnit.readFields(in); this.workUnits.add(workUnit); } super.readFields(in); }
@Override public void map(LongWritable key, Text value, Context context) throws IOException, InterruptedException { WorkUnit workUnit = (value.toString().endsWith(MULTI_WORK_UNIT_FILE_EXTENSION) ? MultiWorkUnit.createEmpty() : WorkUnit.createEmpty()); SerializationUtils.deserializeState(this.fs, new Path(value.toString()), workUnit); if (workUnit instanceof MultiWorkUnit) { List<WorkUnit> flattenedWorkUnits = JobLauncherUtils.flattenWorkUnits(((MultiWorkUnit) workUnit).getWorkUnits()); this.workUnits.addAll(flattenedWorkUnits); } else { this.workUnits.add(workUnit); } } }
@Test public void testFlattenWorkUnits() { List<WorkUnit> workUnitsOnly = Arrays.asList(WorkUnit.createEmpty(), WorkUnit.createEmpty(), WorkUnit.createEmpty()); Assert.assertEquals(JobLauncherUtils.flattenWorkUnits(workUnitsOnly).size(), 3); MultiWorkUnit multiWorkUnit1 = MultiWorkUnit.createEmpty(); multiWorkUnit1.addWorkUnits(Arrays.asList(WorkUnit.createEmpty(), WorkUnit.createEmpty(), WorkUnit.createEmpty())); MultiWorkUnit multiWorkUnit2 = MultiWorkUnit.createEmpty(); multiWorkUnit1.addWorkUnits(Arrays.asList(WorkUnit.createEmpty(), WorkUnit.createEmpty(), WorkUnit.createEmpty())); List<WorkUnit> workUnitsAndMultiWorkUnits = Arrays.asList(WorkUnit.createEmpty(), WorkUnit.createEmpty(), WorkUnit.createEmpty(), multiWorkUnit1, multiWorkUnit2); Assert.assertEquals(JobLauncherUtils.flattenWorkUnits(workUnitsAndMultiWorkUnits).size(), 9); }
@Override public List<WorkUnit> getWorkunits(SourceState state) { WorkUnit workUnit = WorkUnit.createEmpty(); workUnit.addAll(state); return Collections.singletonList(workUnit); }
@Override public List<WorkUnit> getWorkunits(SourceState state) { WorkUnit workUnit = WorkUnit.createEmpty(); workUnit.addAll(state); return Collections.singletonList(workUnit); }
@Test(dependsOnMethods = "testSerializeState") public void testDeserializeState() throws IOException { WorkUnit workUnit1 = WorkUnit.createEmpty(); WorkUnit workUnit2 = WorkUnit.createEmpty(); SerializationUtils.deserializeState(this.fs, new Path(this.outputPath, "wu1"), workUnit1); SerializationUtils.deserializeState(this.fs, new Path(this.outputPath, "wu2"), workUnit2); Assert.assertEquals(workUnit1.getPropertyNames().size(), 2); Assert.assertEquals(workUnit1.getProp("foo"), "bar"); Assert.assertEquals(workUnit1.getPropAsInt("a"), 10); Assert.assertEquals(workUnit2.getPropertyNames().size(), 2); Assert.assertEquals(workUnit2.getProp("foo"), "baz"); Assert.assertEquals(workUnit2.getPropAsInt("b"), 20); }
@Test public void testSerializeState() throws IOException { WorkUnit workUnit1 = WorkUnit.createEmpty(); workUnit1.setProp("foo", "bar"); workUnit1.setProp("a", 10); SerializationUtils.serializeState(this.fs, new Path(this.outputPath, "wu1"), workUnit1); WorkUnit workUnit2 = WorkUnit.createEmpty(); workUnit2.setProp("foo", "baz"); workUnit2.setProp("b", 20); SerializationUtils.serializeState(this.fs, new Path(this.outputPath, "wu2"), workUnit2); }
@Test(dependsOnMethods = "testSerializeToFile") public void testDeserializeFromFile() throws IOException { WorkUnit workUnit1 = WorkUnit.createEmpty(); WorkUnit workUnit2 = WorkUnit.createEmpty(); try (ParallelRunner parallelRunner = new ParallelRunner(2, this.fs)) { parallelRunner.deserializeFromFile(workUnit1, new Path(this.outputPath, "wu1")); parallelRunner.deserializeFromFile(workUnit2, new Path(this.outputPath, "wu2")); } Assert.assertEquals(workUnit1.getPropertyNames().size(), 2); Assert.assertEquals(workUnit1.getProp("foo"), "bar"); Assert.assertEquals(workUnit1.getPropAsInt("a"), 10); Assert.assertEquals(workUnit2.getPropertyNames().size(), 2); Assert.assertEquals(workUnit2.getProp("foo"), "baz"); Assert.assertEquals(workUnit2.getPropAsInt("b"), 20); }
public List<WorkUnit> getWorkunits(SourceState state) { List<WorkUnit> workUnits = new ArrayList<>(); if (state.contains(TEST_WORKUNIT_PERSISTENCE)) { testSkipWorkUnitPersistence(state); return workUnits; } for (int i = 0; i < NUMBER_OF_WORKUNITS; i++) { WorkUnit workUnit = WorkUnit.createEmpty(); if (i < NUMBER_OF_SKIP_WORKUNITS) { workUnit.skip(); } workUnits.add(workUnit); } return workUnits; }
@BeforeMethod private void setup() throws SQLException { this.state = new State(); this.state.setProp(ConfigurationKeys.WRITER_DESTINATION_TYPE_KEY, DestinationType.MYSQL.name()); this.state.setProp(JdbcPublisher.JDBC_PUBLISHER_DATABASE_NAME, DB); this.state.setProp(JdbcPublisher.JDBC_PUBLISHER_FINAL_TABLE_NAME, DEST_TABLE); this.workUnit = WorkUnit.createEmpty(); this.workUnits = Lists.newArrayList(); this.workUnits.add(this.workUnit); this.factory = mock(JdbcWriterCommandsFactory.class); this.commands = mock(JdbcWriterCommands.class); this.conn = mock(Connection.class); doReturn(this.commands).when(this.factory).newInstance(any(Destination.class), eq(this.conn)); this.initializer = new JdbcWriterInitializer(this.state, this.workUnits, this.factory, 1, 0); this.initializer = spy(this.initializer); doReturn(this.conn).when(this.initializer).createConnection(); }
@Test public void testSerializeToFile() throws IOException { try (ParallelRunner parallelRunner = new ParallelRunner(2, this.fs)) { WorkUnit workUnit1 = WorkUnit.createEmpty(); workUnit1.setProp("foo", "bar"); workUnit1.setProp("a", 10); parallelRunner.serializeToFile(workUnit1, new Path(this.outputPath, "wu1")); WorkUnit workUnit2 = WorkUnit.createEmpty(); workUnit2.setProp("foo", "baz"); workUnit2.setProp("b", 20); parallelRunner.serializeToFile(workUnit2, new Path(this.outputPath, "wu2")); } }
/** * Test for {@link MultiWorkUnitWeightedQueue#MultiWorkUnitWeightedQueue()}. It adds a series of WorkUnits to an * instance of MultiWorkUnitWeightedQueue of checks the size of all the WorkUnits returned by * {@link MultiWorkUnitWeightedQueue#getQueueAsList()}. */ @Test public void testDefaultConstructor() { int numWorkUnits = 10; int weight = 1; MultiWorkUnitWeightedQueue multiWorkUnitWeightedQueue = new MultiWorkUnitWeightedQueue(); WorkUnit workUnit = WorkUnit.createEmpty(); for (int i = 0; i < numWorkUnits; i++) { multiWorkUnitWeightedQueue.addWorkUnit(workUnit, weight); } List<WorkUnit> multiWorkUnitWeightedQueueList = multiWorkUnitWeightedQueue.getQueueAsList(); Assert.assertEquals(multiWorkUnitWeightedQueueList.size(), numWorkUnits); MultiWorkUnit multiWorkUnit; for (WorkUnit workUnitElement : multiWorkUnitWeightedQueueList) { multiWorkUnit = (MultiWorkUnit) workUnitElement; Assert.assertEquals(multiWorkUnit.getWorkUnits().size(), 1); } }
/** * Test for {@link MultiWorkUnitWeightedQueue#MultiWorkUnitWeightedQueue(int)}. It sets a limit on the maximum number * of MultiWorkUnits that can be created, adds a series of WorkUnits to the list, and checks the results of * {@link MultiWorkUnitWeightedQueue#getQueueAsList()} to ensure each MultiWorkUnit created is of proper length. */ @Test public void testWithQueueSizeLimit() { int maxMultiWorkUnits = 10; int numWorkUnits = 100; int weight = 1; MultiWorkUnitWeightedQueue multiWorkUnitWeightedQueue = new MultiWorkUnitWeightedQueue(maxMultiWorkUnits); WorkUnit workUnit = WorkUnit.createEmpty(); for (int i = 0; i < numWorkUnits; i++) { multiWorkUnitWeightedQueue.addWorkUnit(workUnit, weight); } MultiWorkUnit multiWorkUnit; for (WorkUnit workUnitElement : multiWorkUnitWeightedQueue.getQueueAsList()) { multiWorkUnit = (MultiWorkUnit) workUnitElement; Assert.assertEquals(multiWorkUnit.getWorkUnits().size(), numWorkUnits / maxMultiWorkUnits); } } }
@BeforeClass public void setUp() throws Exception { WorkUnit workUnit = WorkUnit.createEmpty(); Properties properties = new Properties(); properties.load(new StringReader(TEST_JOB_CONFIG)); workUnit.addAll(properties); workUnit.setProp(ConfigurationKeys.JOB_ID_KEY, JobLauncherUtils.newJobId("GobblinTest1")); workUnit.setProp(ConfigurationKeys.TASK_ID_KEY, JobLauncherUtils.newTaskId(workUnit.getProp(ConfigurationKeys.JOB_ID_KEY), 0)); this.taskContext = new TaskContext(new WorkUnitState(workUnit)); }
@BeforeClass public void setUp() { this.multiWorkUnit = new MultiWorkUnit(); WorkUnit workUnit1 = WorkUnit.createEmpty(); workUnit1.setHighWaterMark(1000); workUnit1.setLowWaterMark(0); workUnit1.setProp("k1", "v1"); this.multiWorkUnit.addWorkUnit(workUnit1); WorkUnit workUnit2 = WorkUnit.createEmpty(); workUnit2.setHighWaterMark(2000); workUnit2.setLowWaterMark(1001); workUnit2.setProp("k2", "v2"); this.multiWorkUnit.addWorkUnit(workUnit2); }
@Test public void testBuilder() throws IOException, InterruptedException, TimeoutException { final String eventBusId = "/GobblinTestEventBusWriterTest/testBuilder"; GobblinTestEventBusWriter.Builder writerBuilder = new GobblinTestEventBusWriter.Builder(); WorkUnit wu = WorkUnit.createEmpty(); wu.setProp(GobblinTestEventBusWriter.FULL_EVENTBUSID_KEY, eventBusId); writerBuilder.writeTo(Destination.of(DestinationType.HDFS, wu)); Assert.assertEquals(writerBuilder.getEventBusId(), eventBusId); try(TestingEventBusAsserter asserter = new TestingEventBusAsserter(eventBusId)) { GobblinTestEventBusWriter writer = writerBuilder.build(); writer.write("event1"); writer.write("event2"); asserter.assertNextValueEq("event1"); asserter.assertNextValueEq("event2"); Assert.assertEquals(writer.recordsWritten(), 2); } }
@Test public void testDataPullUpperBoundsRemovedInLastWorkUnit() { int totalCount = 5; ArrayList<DataRecord> records = this.generateRecords(totalCount); WorkUnit workUnit = WorkUnit.createEmpty(); workUnit.setProp(Partition.IS_LAST_PARTIITON, true); workUnit.setProp(ConfigurationKeys.SOURCE_QUERYBASED_EXTRACT_TYPE, "SNAPSHOT"); WorkUnitState workUnitState = new WorkUnitState(workUnit, new State()); workUnitState.setId("testDataPullUpperBoundsRemovedInLastWorkUnit"); TestQueryBasedExtractor testExtractor = new TestQueryBasedExtractor(workUnitState, records); testExtractor.setRangePredicates(1, 3); this.verify(testExtractor, totalCount); }
protected WorkUnit createNewWorkUnit(HivePartitionDataset dataset, int executionAttempts) { WorkUnit workUnit = WorkUnit.createEmpty(); workUnit.setProp(ComplianceConfigurationKeys.PARTITION_NAME, dataset.datasetURN()); workUnit.setProp(ComplianceConfigurationKeys.EXECUTION_ATTEMPTS, executionAttempts); workUnit.setProp(ComplianceConfigurationKeys.TIMESTAMP, this.timeStamp); workUnit.setProp(ComplianceConfigurationKeys.GOBBLIN_COMPLIANCE_SHOULD_PROXY, this.shouldProxy); workUnit.setProp(ComplianceConfigurationKeys.EXECUTION_COUNT, this.executionCount); workUnit.setProp(ComplianceConfigurationKeys.NUM_ROWS, DatasetUtils .getProperty(dataset, ComplianceConfigurationKeys.NUM_ROWS, ComplianceConfigurationKeys.DEFAULT_NUM_ROWS)); workUnit.setProp(ComplianceConfigurationKeys.RAW_DATA_SIZE, DatasetUtils .getProperty(dataset, ComplianceConfigurationKeys.RAW_DATA_SIZE, ComplianceConfigurationKeys.DEFAULT_RAW_DATA_SIZE)); workUnit.setProp(ComplianceConfigurationKeys.TOTAL_SIZE, DatasetUtils .getProperty(dataset, ComplianceConfigurationKeys.TOTAL_SIZE, ComplianceConfigurationKeys.DEFAULT_TOTAL_SIZE)); submitWorkUnitGeneratedEvent(dataset.datasetURN(), executionAttempts); return workUnit; }
private WorkUnitState getMockWorkUnitState(Long lowWaterMark, Long highWaterMark) { WorkUnit mockWorkUnit = WorkUnit.createEmpty(); mockWorkUnit.setWatermarkInterval(new WatermarkInterval(new MultiLongWatermark(new ArrayList<Long>(){{add(lowWaterMark);}}), new MultiLongWatermark(new ArrayList<Long>(){{add(highWaterMark);}}))); WorkUnitState mockWorkUnitState = new WorkUnitState(mockWorkUnit, new State()); mockWorkUnitState.setProp(KafkaSource.TOPIC_NAME, TEST_TOPIC_NAME); mockWorkUnitState.setProp(KafkaSource.PARTITION_ID, "1"); mockWorkUnitState.setProp(ConfigurationKeys.KAFKA_BROKERS, "localhost:8080"); mockWorkUnitState.setProp(KafkaSchemaRegistry.KAFKA_SCHEMA_REGISTRY_URL, TEST_URL); return mockWorkUnitState; }