Examples with WorkloadRecordStoreConnection com.yahoo.athenz.common.server.workload.WorkloadRecordStoreConnection used on opensource projects

Example 6 with WorkloadRecordStoreConnection

use of com.yahoo.athenz.common.server.workload.WorkloadRecordStoreConnection in project athenz by yahoo.

the class InstanceCertManagerTest method testGetWorkloadsByIp.

@Test
public void testGetWorkloadsByIp() {
    InstanceCertManager instance = new InstanceCertManager(null, null, null, true, null);
    WorkloadRecordStore store = Mockito.mock(WorkloadRecordStore.class);
    instance.setWorkloadStore(store);
    WorkloadRecordStoreConnection storeConn = Mockito.mock(WorkloadRecordStoreConnection.class);
    Mockito.when(store.getConnection()).thenReturn(storeConn);
    long currTime = System.currentTimeMillis();
    Date d = new Date(currTime);
    WorkloadRecord w1 = ZTSTestUtils.createWorkloadRecord(d, d, "aws", "i-123", "test-host1.corp.yahoo.com", "10.0.0.1", "athenz.api", d);
    WorkloadRecord w5 = ZTSTestUtils.createWorkloadRecord(d, d, "aws", "i-123", "test-host2.corp.yahoo.com", "10.0.0.1", "athenz.secondapi", d);
    List<WorkloadRecord> workloadRecordList = new ArrayList<>();
    workloadRecordList.add(w1);
    workloadRecordList.add(w5);
    Mockito.when(storeConn.getWorkloadRecordsByIp(any())).thenReturn(workloadRecordList);
    List<Workload> workloadList = instance.getWorkloadsByIp("10.0.0.1");
    assertNotNull(workloadList);
    assertEquals(workloadList.size(), 2);
    List<String> expectedServices = workloadList.stream().map(Workload::getServiceName).collect(Collectors.toList());
    assertTrue(expectedServices.contains("api"));
    assertTrue(expectedServices.contains("secondapi"));
    instance.shutdown();
}

Example 7 with WorkloadRecordStoreConnection

use of com.yahoo.athenz.common.server.workload.WorkloadRecordStoreConnection in project athenz by yahoo.

the class InstanceCertManager method getWorkloadsByService.

public List<Workload> getWorkloadsByService(String domain, String service) {
    if (workloadStore == null) {
        return Collections.emptyList();
    }
    try (WorkloadRecordStoreConnection storeConnection = workloadStore.getConnection()) {
        List<WorkloadRecord> workloadRecords = storeConnection.getWorkloadRecordsByService(domain, service);
        Map<String, List<String>> flattenedIpAddresses = new HashMap<>();
        String mapKey;
        for (WorkloadRecord workloadRecord : workloadRecords) {
            mapKey = workloadRecord.getInstanceId() + ":" + workloadRecord.getProvider() + ":" + workloadRecord.getUpdateTime().getTime() + ":" + workloadRecord.getCertExpiryTime().getTime() + ":" + workloadRecord.getHostname();
            if (flattenedIpAddresses.containsKey(mapKey)) {
                flattenedIpAddresses.get(mapKey).add(workloadRecord.getIp());
            } else {
                List<String> ipList = new ArrayList<>();
                ipList.add(workloadRecord.getIp());
                flattenedIpAddresses.put(mapKey, ipList);
            }
        }
        return flattenedIpAddresses.entrySet().stream().map(entry -> {
            Workload wl = new Workload();
            String[] tempArr = entry.getKey().split(":");
            wl.setUuid(tempArr[0]).setProvider(tempArr[1]).setUpdateTime(Timestamp.fromMillis(Long.parseLong(tempArr[2]))).setCertExpiryTime(Timestamp.fromMillis(Long.parseLong(tempArr[3]))).setHostname(tempArr[4]).setIpAddresses(entry.getValue());
            return wl;
        }).collect(Collectors.toList());
    }
}