Examples with TaskContextImpl org.apache.samza.context.TaskContextImpl used on opensource projects

Example 1 with TaskContextImpl

use of org.apache.samza.context.TaskContextImpl in project samza by apache.

the class TestRemoteTableDescriptor method createMockContext.

private Context createMockContext(TableDescriptor tableDescriptor) {
    Context context = mock(Context.class);
    ContainerContext containerContext = mock(ContainerContext.class);
    when(context.getContainerContext()).thenReturn(containerContext);
    MetricsRegistry metricsRegistry = mock(MetricsRegistry.class);
    when(metricsRegistry.newTimer(anyString(), anyString())).thenReturn(mock(Timer.class));
    when(metricsRegistry.newCounter(anyString(), anyString())).thenReturn(mock(Counter.class));
    when(containerContext.getContainerMetricsRegistry()).thenReturn(metricsRegistry);
    TaskContextImpl taskContext = mock(TaskContextImpl.class);
    when(context.getTaskContext()).thenReturn(taskContext);
    TaskName taskName = new TaskName("MyTask");
    TaskModel taskModel = mock(TaskModel.class);
    when(taskModel.getTaskName()).thenReturn(taskName);
    when(context.getTaskContext().getTaskModel()).thenReturn(taskModel);
    ContainerModel containerModel = mock(ContainerModel.class);
    when(containerModel.getTasks()).thenReturn(ImmutableMap.of(taskName, taskModel));
    when(containerContext.getContainerModel()).thenReturn(containerModel);
    String containerId = "container-1";
    JobModel jobModel = mock(JobModel.class);
    when(taskContext.getJobModel()).thenReturn(jobModel);
    when(jobModel.getContainers()).thenReturn(ImmutableMap.of(containerId, containerModel));
    JobContext jobContext = mock(JobContext.class);
    Config jobConfig = new MapConfig(tableDescriptor.toConfig(new MapConfig()));
    when(jobContext.getConfig()).thenReturn(jobConfig);
    when(context.getJobContext()).thenReturn(jobContext);
    return context;
}

Example 2 with TaskContextImpl

use of org.apache.samza.context.TaskContextImpl in project samza by apache.

the class TestWindowOperator method setup.

@Before
public void setup() {
    Map<String, String> configMap = new HashMap<>();
    configMap.put("job.default.system", "kafka");
    configMap.put("job.name", "jobName");
    configMap.put("job.id", "jobId");
    this.config = new MapConfig(configMap);
    this.context = new MockContext();
    when(this.context.getJobContext().getConfig()).thenReturn(this.config);
    Serde storeKeySerde = new TimeSeriesKeySerde(new IntegerSerde());
    Serde storeValSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    SystemStreamPartition ssp = new SystemStreamPartition("kafka", "integers", new Partition(0));
    TaskModel taskModel = mock(TaskModel.class);
    when(taskModel.getSystemStreamPartitions()).thenReturn(ImmutableSet.of(ssp));
    when(taskModel.getTaskName()).thenReturn(new TaskName("task 1"));
    when(this.context.getTaskContext().getTaskModel()).thenReturn(taskModel);
    when(((TaskContextImpl) this.context.getTaskContext()).getSspsExcludingSideInputs()).thenReturn(ImmutableSet.of(ssp));
    when(this.context.getTaskContext().getTaskMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    when(this.context.getContainerContext().getContainerMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    when(this.context.getTaskContext().getStore("jobName-jobId-window-w1")).thenReturn(new TestInMemoryStore<>(storeKeySerde, storeValSerde));
}

Example 3 with TaskContextImpl

use of org.apache.samza.context.TaskContextImpl in project samza by apache.

the class TestEmbeddedTaggedRateLimiter method initRateLimiter.

static void initRateLimiter(RateLimiter rateLimiter) {
    Map<TaskName, TaskModel> tasks = IntStream.range(0, NUMBER_OF_TASKS).mapToObj(i -> new TaskName("task-" + i)).collect(Collectors.toMap(Function.identity(), x -> mock(TaskModel.class)));
    ContainerModel containerModel = mock(ContainerModel.class);
    when(containerModel.getTasks()).thenReturn(tasks);
    JobModel jobModel = mock(JobModel.class);
    Map<String, ContainerModel> containerModelMap = new HashMap<>();
    containerModelMap.put("container-1", containerModel);
    when(jobModel.getContainers()).thenReturn(containerModelMap);
    Context context = mock(Context.class);
    TaskContextImpl taskContext = mock(TaskContextImpl.class);
    when(context.getTaskContext()).thenReturn(taskContext);
    when(taskContext.getJobModel()).thenReturn(jobModel);
    when(context.getTaskContext().getTaskModel()).thenReturn(mock(TaskModel.class));
    rateLimiter.init(context);
}

Example 4 with TaskContextImpl

use of org.apache.samza.context.TaskContextImpl in project samza by apache.

the class TestOperatorImplGraph method testPartitionByChain.

@Test
public void testPartitionByChain() {
    String inputStreamId = "input";
    String inputSystem = "input-system";
    String inputPhysicalName = "input-stream";
    String outputStreamId = "output";
    String outputSystem = "output-system";
    String outputPhysicalName = "output-stream";
    String intermediateStreamId = "jobName-jobId-partition_by-p1";
    String intermediateSystem = "intermediate-system";
    HashMap<String, String> configs = new HashMap<>();
    configs.put(JobConfig.JOB_NAME, "jobName");
    configs.put(JobConfig.JOB_ID, "jobId");
    configs.put(JobConfig.JOB_DEFAULT_SYSTEM, intermediateSystem);
    StreamTestUtils.addStreamConfigs(configs, inputStreamId, inputSystem, inputPhysicalName);
    StreamTestUtils.addStreamConfigs(configs, outputStreamId, outputSystem, outputPhysicalName);
    Config config = new MapConfig(configs);
    when(this.context.getJobContext().getConfig()).thenReturn(config);
    StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> {
        GenericSystemDescriptor isd = new GenericSystemDescriptor(inputSystem, "mockFactoryClass");
        GenericSystemDescriptor osd = new GenericSystemDescriptor(outputSystem, "mockFactoryClass");
        GenericInputDescriptor inputDescriptor = isd.getInputDescriptor(inputStreamId, mock(Serde.class));
        GenericOutputDescriptor outputDescriptor = osd.getOutputDescriptor(outputStreamId, KVSerde.of(mock(IntegerSerde.class), mock(StringSerde.class)));
        MessageStream<Object> inputStream = appDesc.getInputStream(inputDescriptor);
        OutputStream<KV<Integer, String>> outputStream = appDesc.getOutputStream(outputDescriptor);
        inputStream.partitionBy(Object::hashCode, Object::toString, KVSerde.of(mock(IntegerSerde.class), mock(StringSerde.class)), "p1").sendTo(outputStream);
    }, config);
    JobModel jobModel = mock(JobModel.class);
    ContainerModel containerModel = mock(ContainerModel.class);
    TaskModel taskModel = mock(TaskModel.class);
    when(jobModel.getContainers()).thenReturn(Collections.singletonMap("0", containerModel));
    when(containerModel.getTasks()).thenReturn(Collections.singletonMap(new TaskName("task 0"), taskModel));
    when(taskModel.getSystemStreamPartitions()).thenReturn(Collections.emptySet());
    when(((TaskContextImpl) this.context.getTaskContext()).getJobModel()).thenReturn(jobModel);
    OperatorImplGraph opImplGraph = new OperatorImplGraph(graphSpec.getOperatorSpecGraph(), this.context, mock(Clock.class));
    InputOperatorImpl inputOpImpl = opImplGraph.getInputOperator(new SystemStream(inputSystem, inputPhysicalName));
    assertEquals(1, inputOpImpl.registeredOperators.size());
    OperatorImpl partitionByOpImpl = (PartitionByOperatorImpl) inputOpImpl.registeredOperators.iterator().next();
    // is terminal but paired with an input operator
    assertEquals(0, partitionByOpImpl.registeredOperators.size());
    assertEquals(OpCode.PARTITION_BY, partitionByOpImpl.getOperatorSpec().getOpCode());
    InputOperatorImpl repartitionedInputOpImpl = opImplGraph.getInputOperator(new SystemStream(intermediateSystem, intermediateStreamId));
    assertEquals(1, repartitionedInputOpImpl.registeredOperators.size());
    OperatorImpl sendToOpImpl = (OutputOperatorImpl) repartitionedInputOpImpl.registeredOperators.iterator().next();
    assertEquals(0, sendToOpImpl.registeredOperators.size());
    assertEquals(OpCode.SEND_TO, sendToOpImpl.getOperatorSpec().getOpCode());
}