application/src/test/java/com/yahoo/application/container/docprocs/MockDispatchDocproc.java


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.application.container.docprocs;

import com.yahoo.docproc.DocumentProcessor;
import com.yahoo.docproc.Processing;
import com.yahoo.document.DocumentOperation;
import com.yahoo.document.DocumentPut;
import com.yahoo.documentapi.messagebus.protocol.DocumentMessage;
import com.yahoo.documentapi.messagebus.protocol.PutDocumentMessage;
import com.yahoo.jdisc.Request;
import com.yahoo.jdisc.Response;
import com.yahoo.jdisc.handler.RequestDispatch;
import com.yahoo.jdisc.service.CurrentContainer;
import com.yahoo.messagebus.jdisc.MbusRequest;
import com.yahoo.messagebus.routing.Route;

import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutionException;

/**
 * @author Christian Andersen
 */
public class MockDispatchDocproc extends DocumentProcessor {
    private final Route route;
    private final URI uri;
    private final CurrentContainer currentContainer;
    private final List<Response> responses = new ArrayList<>();

    public MockDispatchDocproc(CurrentContainer currentContainer) {
        this.route = Route.parse("default");
        this.uri = URI.create("mbus://remotehost/source");
        this.currentContainer = currentContainer;
    }

    @Override
    public Progress process(Processing processing) {
        for (DocumentOperation op : processing.getDocumentOperations()) {
            PutDocumentMessage message = new PutDocumentMessage((DocumentPut)op);
            var future = createRequest(message).dispatch();
            try {
                responses.add(future.get());
            } catch (ExecutionException | InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        return Progress.DONE;
    }

    private RequestDispatch createRequest(final DocumentMessage message) {
        return new RequestDispatch() {
            @Override
            protected Request newRequest() {
                return new MbusRequest(currentContainer, uri, message.setRoute(route), false);
            }
        };
    }

    public List<Response> getResponses() {
        return responses;
    }
}