org.apache.aurora.scheduler.events.WebhookTest.java Source code

Introduction

Here is the source code for org.apache.aurora.scheduler.events.WebhookTest.java

Source

/**
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.aurora.scheduler.events;

import java.io.IOException;
import java.net.URI;
import java.util.Map;

import com.google.common.collect.ImmutableMap;

import org.apache.aurora.common.testing.easymock.EasyMockTest;

import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.scheduler.base.TaskTestUtil;
import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.http.Header;
import org.apache.http.HttpEntity;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.util.EntityUtils;
import org.easymock.Capture;
import org.junit.Before;
import org.junit.Test;

import static org.easymock.EasyMock.capture;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.expectLastCall;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

public class WebhookTest extends EasyMockTest {

    private static final IScheduledTask TASK = TaskTestUtil.makeTask("id", TaskTestUtil.JOB);
    private final TaskStateChange change = TaskStateChange.initialized(TASK);
    private final TaskStateChange changeWithOldState = TaskStateChange.transition(TASK, ScheduleStatus.FAILED);
    private final String changeJson = changeWithOldState.toJson();

    private CloseableHttpClient httpClient;
    private Webhook webhook;

    @Before
    public void setUp() {
        WebhookInfo webhookInfo = WebhookModule.parseWebhookConfig(WebhookModule.readWebhookFile());
        httpClient = createMock(CloseableHttpClient.class);
        webhook = new Webhook(httpClient, webhookInfo);
    }

    @Test
    public void testTaskChangedStateNoOldState() throws Exception {
        // Should be a noop as oldState is MIA so this test would have throw an exception.
        // If it does not, then we are good.
        control.replay();
        webhook.taskChangedState(change);
    }

    @Test
    public void testTaskChangedWithOldState() throws Exception {
        CloseableHttpResponse httpResponse = createMock(CloseableHttpResponse.class);
        HttpEntity entity = createMock(HttpEntity.class);

        Capture<HttpPost> httpPostCapture = createCapture();
        expect(entity.isStreaming()).andReturn(false);
        expect(httpResponse.getEntity()).andReturn(entity);
        httpResponse.close();
        expectLastCall().once();
        expect(httpClient.execute(capture(httpPostCapture))).andReturn(httpResponse);

        control.replay();

        webhook.taskChangedState(changeWithOldState);

        assertTrue(httpPostCapture.hasCaptured());
        assertEquals(httpPostCapture.getValue().getURI(), new URI("http://localhost:5000/"));
        assertEquals(EntityUtils.toString(httpPostCapture.getValue().getEntity()), changeJson);
        Header[] producerTypeHeader = httpPostCapture.getValue().getHeaders("Producer-Type");
        assertEquals(producerTypeHeader.length, 1);
        assertEquals(producerTypeHeader[0].getName(), "Producer-Type");
        assertEquals(producerTypeHeader[0].getValue(), "reliable");
        Header[] contentTypeHeader = httpPostCapture.getValue().getHeaders("Content-Type");
        assertEquals(contentTypeHeader.length, 1);
        assertEquals(contentTypeHeader[0].getName(), "Content-Type");
        assertEquals(contentTypeHeader[0].getValue(), "application/vnd.kafka.json.v1+json");
        assertNotNull(httpPostCapture.getValue().getHeaders("Timestamp"));
    }

    @Test
    public void testCatchHttpClientException() throws Exception {
        // IOException should be silenced.
        Capture<HttpPost> httpPostCapture = createCapture();
        expect(httpClient.execute(capture(httpPostCapture))).andThrow(new IOException());
        control.replay();

        webhook.taskChangedState(changeWithOldState);
    }

    @Test
    public void testWebhookInfo() throws Exception {
        WebhookInfo parsedWebhookInfo = WebhookModule.parseWebhookConfig(WebhookModule.readWebhookFile());
        assertEquals(parsedWebhookInfo.toString(),
                "WebhookInfo{headers={" + "Content-Type=application/vnd.kafka.json.v1+json, "
                        + "Producer-Type=reliable" + "}, " + "targetURI=http://localhost:5000/, "
                        + "connectTimeoutMsec=50" + "}");
        // Verifying all attributes were parsed correctly.
        Map<String, String> headers = ImmutableMap.of("Content-Type", "application/vnd.kafka.json.v1+json",
                "Producer-Type", "reliable");
        assertEquals(parsedWebhookInfo.getHeaders(), headers);
        URI targetURI = new URI("http://localhost:5000/");
        assertEquals(parsedWebhookInfo.getTargetURI(), targetURI);
        Integer timeoutMsec = 50;
        assertEquals(parsedWebhookInfo.getConnectonTimeoutMsec(), timeoutMsec);

        control.replay();
    }
}