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

Example 1 with MockContext

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

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

use of org.apache.samza.context.MockContext 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());
}

Example 4 with MockContext

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

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

the class TestJoinOperator method createStreamOperatorTask.

private StreamOperatorTask createStreamOperatorTask(Clock clock, StreamApplicationDescriptorImpl graphSpec) throws Exception {
    Map<String, String> mapConfig = new HashMap<>();
    mapConfig.put("job.name", "jobName");
    mapConfig.put("job.id", "jobId");
    StreamTestUtils.addStreamConfigs(mapConfig, "inStream", "insystem", "instream");
    StreamTestUtils.addStreamConfigs(mapConfig, "inStream2", "insystem", "instream2");
    Context context = new MockContext(new MapConfig(mapConfig));
    TaskModel taskModel = mock(TaskModel.class);
    when(taskModel.getSystemStreamPartitions()).thenReturn(ImmutableSet.of(new SystemStreamPartition("insystem", "instream", new Partition(0)), new SystemStreamPartition("insystem", "instream2", new Partition(0))));
    when(context.getTaskContext().getTaskModel()).thenReturn(taskModel);
    when(context.getTaskContext().getTaskMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    when(context.getContainerContext().getContainerMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    // need to return different stores for left and right side
    IntegerSerde integerSerde = new IntegerSerde();
    TimestampedValueSerde timestampedValueSerde = new TimestampedValueSerde(new KVSerde(integerSerde, integerSerde));
    when(context.getTaskContext().getStore(eq("jobName-jobId-join-j1-L"))).thenReturn(new TestInMemoryStore(integerSerde, timestampedValueSerde));
    when(context.getTaskContext().getStore(eq("jobName-jobId-join-j1-R"))).thenReturn(new TestInMemoryStore(integerSerde, timestampedValueSerde));
    StreamOperatorTask sot = new StreamOperatorTask(graphSpec.getOperatorSpecGraph(), clock);
    sot.init(context);
    return sot;
}