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

Example 1 with Context

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

the class TestLocalTableProvider method testInit.

@Test
public void testInit() {
    Context context = mock(Context.class);
    JobContext jobContext = mock(JobContext.class);
    when(context.getJobContext()).thenReturn(jobContext);
    when(jobContext.getConfig()).thenReturn(new MapConfig());
    ContainerContext containerContext = mock(ContainerContext.class);
    when(context.getContainerContext()).thenReturn(containerContext);
    when(containerContext.getContainerMetricsRegistry()).thenReturn(new NoOpMetricsRegistry());
    TaskContext taskContext = mock(TaskContext.class);
    when(context.getTaskContext()).thenReturn(taskContext);
    when(taskContext.getStore(any())).thenReturn(mock(KeyValueStore.class));
    TableProvider tableProvider = createTableProvider("t1");
    tableProvider.init(context);
    Assert.assertNotNull(tableProvider.getTable());
}

Example 2 with Context

use of org.apache.samza.context.Context 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 3 with Context

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

the class TestCachingTable method doTestCacheOps.

private void doTestCacheOps(boolean isWriteAround) {
    CachingTableDescriptor desc = new CachingTableDescriptor("1", createDummyTableDescriptor("realTable"), createDummyTableDescriptor("cacheTable"));
    if (isWriteAround) {
        desc.withWriteAround();
    }
    Context context = new MockContext();
    final ReadWriteUpdateTable cacheTable = getMockCache().getLeft();
    final ReadWriteUpdateTable realTable = mock(ReadWriteUpdateTable.class);
    doAnswer(invocation -> {
        String key = invocation.getArgumentAt(0, String.class);
        return CompletableFuture.completedFuture("test-data-" + key);
    }).when(realTable).getAsync(any());
    doReturn(CompletableFuture.completedFuture(null)).when(realTable).putAsync(any(), any());
    doAnswer(invocation -> {
        String tableId = invocation.getArgumentAt(0, String.class);
        if (tableId.equals("realTable")) {
            // cache
            return realTable;
        } else if (tableId.equals("cacheTable")) {
            return cacheTable;
        }
        Assert.fail();
        return null;
    }).when(context.getTaskContext()).getUpdatableTable(anyString());
    when(context.getContainerContext().getContainerMetricsRegistry()).thenReturn(new NoOpMetricsRegistry());
    Map<String, String> tableConfig = desc.toConfig(new MapConfig());
    when(context.getJobContext().getConfig()).thenReturn(new MapConfig(tableConfig));
    CachingTableProvider tableProvider = new CachingTableProvider(desc.getTableId());
    tableProvider.init(context);
    CachingTable cachingTable = (CachingTable) tableProvider.getTable();
    Assert.assertEquals("test-data-1", cachingTable.get("1"));
    verify(realTable, times(1)).getAsync(any());
    // cache miss
    verify(cacheTable, times(1)).get(any());
    verify(cacheTable, times(1)).put(any(), any());
    // 0 hit, 1 request
    Assert.assertEquals(cachingTable.hitRate(), 0.0, 0.0);
    Assert.assertEquals(cachingTable.missRate(), 1.0, 0.0);
    Assert.assertEquals("test-data-1", cachingTable.get("1"));
    // no change
    verify(realTable, times(1)).getAsync(any());
    verify(cacheTable, times(2)).get(any());
    // no change
    verify(cacheTable, times(1)).put(any(), any());
    // 1 hit, 2 requests
    Assert.assertEquals(0.5, cachingTable.hitRate(), 0.0);
    Assert.assertEquals(0.5, cachingTable.missRate(), 0.0);
    cachingTable.put("2", "test-data-XXXX");
    verify(cacheTable, times(isWriteAround ? 1 : 2)).put(any(), any());
    verify(realTable, times(1)).putAsync(any(), any());
    if (isWriteAround) {
        // expects value from table
        Assert.assertEquals("test-data-2", cachingTable.get("2"));
        // should have one more fetch
        verify(realTable, times(2)).getAsync(any());
        // 1 hit, 3 requests
        Assert.assertEquals(cachingTable.hitRate(), 0.33, 0.1);
    } else {
        // expect value from cache
        Assert.assertEquals("test-data-XXXX", cachingTable.get("2"));
        // no change
        verify(realTable, times(1)).getAsync(any());
        // 2 hits, 3 requests
        Assert.assertEquals(cachingTable.hitRate(), 0.66, 0.1);
    }
}

Example 4 with Context

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

the class TestRemoteTable method getMockContext.

public static Context getMockContext() {
    Context context = new MockContext();
    MetricsRegistry metricsRegistry = mock(MetricsRegistry.class);
    doAnswer(args -> new Timer((String) args.getArguments()[0])).when(metricsRegistry).newTimer(anyString(), anyString());
    doAnswer(args -> new Counter((String) args.getArguments()[0])).when(metricsRegistry).newCounter(anyString(), anyString());
    doAnswer(args -> new Gauge((String) args.getArguments()[0], 0)).when(metricsRegistry).newGauge(anyString(), any());
    doReturn(metricsRegistry).when(context.getContainerContext()).getContainerMetricsRegistry();
    return context;
}

Example 5 with Context

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

the class TestStreamTableJoinOperatorImpl method testHandleMessage.

@Test
public void testHandleMessage() {
    String tableId = "t1";
    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) {
            if ("1".equals(message.getKey())) {
                Assert.assertEquals("m1", message.getValue());
                Assert.assertEquals("r1", record.getValue());
                return "m1r1";
            } else if ("2".equals(message.getKey())) {
                Assert.assertEquals("m2", message.getValue());
                Assert.assertNull(record);
                return null;
            }
            throw new SamzaException("Should never reach here!");
        }

        @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"));
    when(table.getAsync("2")).thenReturn(CompletableFuture.completedFuture(null));
    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
    Collection<TestMessageEnvelope> result;
    result = streamTableJoinOperator.handleMessage(KV.of("1", "m1"), mockMessageCollector, mockTaskCoordinator);
    Assert.assertEquals(1, result.size());
    Assert.assertEquals("m1r1", result.iterator().next());
    // Table doesn't have the key
    result = streamTableJoinOperator.handleMessage(KV.of("2", "m2"), mockMessageCollector, mockTaskCoordinator);
    Assert.assertEquals(0, result.size());
}