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

Example 6 with MockContext

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

the class TestOperatorImplGraph method setup.

@Before
public void setup() {
    this.context = new MockContext();
    // individual tests can override this config if necessary
    when(this.context.getJobContext().getConfig()).thenReturn(mock(Config.class));
    TaskModel taskModel = mock(TaskModel.class);
    when(taskModel.getTaskName()).thenReturn(new TaskName("task 0"));
    when(this.context.getTaskContext().getTaskModel()).thenReturn(taskModel);
    when(this.context.getTaskContext().getTaskMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    when(this.context.getContainerContext().getContainerMetricsRegistry()).thenReturn(new MetricsRegistryMap());
}

Example 7 with MockContext

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

the class TestOperatorImpl method setup.

@Before
public void setup() {
    this.context = new MockContext();
    this.internalTaskContext = mock(InternalTaskContext.class);
    when(this.internalTaskContext.getContext()).thenReturn(this.context);
    // might be necessary in the future
    when(this.internalTaskContext.fetchObject(EndOfStreamStates.class.getName())).thenReturn(mock(EndOfStreamStates.class));
    when(this.internalTaskContext.fetchObject(WatermarkStates.class.getName())).thenReturn(mock(WatermarkStates.class));
    when(this.context.getTaskContext().getTaskMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    when(this.context.getTaskContext().getTaskModel()).thenReturn(mock(TaskModel.class));
    when(this.context.getContainerContext().getContainerMetricsRegistry()).thenReturn(new MetricsRegistryMap());
}

Example 8 with MockContext

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

the class TestStreamTableJoinOperatorImpl method testJoinFunctionIsInvokedOnlyOnce.

/**
 * Ensure join function is not invoked more than once when join function returns null on the first invocation
 */
@Test
public void testJoinFunctionIsInvokedOnlyOnce() {
    final String tableId = "testTable";
    final CountDownLatch joinInvokedLatch = new CountDownLatch(1);
    StreamTableJoinOperatorSpec mockJoinOpSpec = mock(StreamTableJoinOperatorSpec.class);
    when(mockJoinOpSpec.getTableId()).thenReturn(tableId);
    when(mockJoinOpSpec.getArgs()).thenReturn(new Object[0]);
    when(mockJoinOpSpec.getJoinFn()).thenReturn(new StreamTableJoinFunction<String, KV<String, String>, KV<String, String>, String>() {

        @Override
        public String apply(KV<String, String> message, KV<String, String> record) {
            joinInvokedLatch.countDown();
            return null;
        }

        @Override
        public String getMessageKey(KV<String, String> message) {
            return message.getKey();
        }

        @Override
        public String getRecordKey(KV<String, String> record) {
            return record.getKey();
        }
    });
    ReadWriteUpdateTable table = mock(ReadWriteUpdateTable.class);
    when(table.getAsync("1")).thenReturn(CompletableFuture.completedFuture("r1"));
    Context context = new MockContext();
    when(context.getTaskContext().getUpdatableTable(tableId)).thenReturn(table);
    MessageCollector mockMessageCollector = mock(MessageCollector.class);
    TaskCoordinator mockTaskCoordinator = mock(TaskCoordinator.class);
    StreamTableJoinOperatorImpl streamTableJoinOperator = new StreamTableJoinOperatorImpl(mockJoinOpSpec, context);
    // Table has the key
    streamTableJoinOperator.handleMessage(KV.of("1", "m1"), mockMessageCollector, mockTaskCoordinator);
    assertEquals("Join function should only be invoked once", 0, joinInvokedLatch.getCount());
}

Example 9 with MockContext

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

the class TestRemoteTableEndToEnd method createMockContext.

private Context createMockContext() {
    MetricsRegistry metricsRegistry = mock(MetricsRegistry.class);
    doReturn(new Counter("")).when(metricsRegistry).newCounter(anyString(), anyString());
    doReturn(new Timer("")).when(metricsRegistry).newTimer(anyString(), anyString());
    Context context = new MockContext();
    doReturn(new MapConfig()).when(context.getJobContext()).getConfig();
    doReturn(metricsRegistry).when(context.getContainerContext()).getContainerMetricsRegistry();
    return context;
}

Example 10 with MockContext

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

the class TestCachingTable method initTables.

private void initTables(boolean isTimerMetricsDisabled, ReadWriteUpdateTable... tables) {
    Map<String, String> config = new HashMap<>();
    if (isTimerMetricsDisabled) {
        config.put(MetricsConfig.METRICS_TIMER_ENABLED, "false");
    }
    Context context = new MockContext();
    doReturn(new MapConfig(config)).when(context.getJobContext()).getConfig();
    metricsRegistry = mock(MetricsRegistry.class);
    doReturn(mock(Timer.class)).when(metricsRegistry).newTimer(anyString(), anyString());
    doReturn(mock(Counter.class)).when(metricsRegistry).newCounter(anyString(), anyString());
    doReturn(mock(Gauge.class)).when(metricsRegistry).newGauge(anyString(), any());
    doReturn(metricsRegistry).when(context.getContainerContext()).getContainerMetricsRegistry();
    Arrays.asList(tables).forEach(t -> t.init(context));
}