summaryrefslogtreecommitdiffstats
path: root/container-search/src/test/java/com/yahoo/search/dispatch/MockSearchCluster.java
blob: ee903fd3fa09a7cb349ce02c8940ab3ba772b1f8 (plain) (blame)
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.search.dispatch;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMultimap;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

/**
 * @author ollivir
 */
public class MockSearchCluster extends SearchCluster {
    private final int numGroups;
    private final int numNodesPerGroup;
    private final ImmutableMap<Integer, Group> groups;
    private final ImmutableMultimap<String, Node> nodesByHost;

    public MockSearchCluster(int groups, int nodesPerGroup) {
        super(100, 100, 0, Collections.emptyList(), null, 1, null);

        ImmutableMap.Builder<Integer, Group> groupBuilder = ImmutableMap.builder();
        ImmutableMultimap.Builder<String, Node> hostBuilder = ImmutableMultimap.builder();
        int dk = 1;
        for (int group = 0; group < groups; group++) {
            List<Node> nodes = new ArrayList<>();
            for (int node = 0; node < nodesPerGroup; node++) {
                Node n = new Node(dk, "host" + dk, -1, group);
                n.setWorking(true);
                nodes.add(n);
                hostBuilder.put(n.hostname(), n);
                dk++;
            }
            groupBuilder.put(group, new Group(group, nodes));
        }
        this.groups = groupBuilder.build();
        this.nodesByHost = hostBuilder.build();
        this.numGroups = groups;
        this.numNodesPerGroup = nodesPerGroup;
    }

    @Override
    public int size() {
        return numGroups * numNodesPerGroup;
    }

    public ImmutableMap<Integer, Group> groups() {
        return groups;
    }

    public int groupSize() {
        return numNodesPerGroup;
    }

    public Optional<Group> group(int n) {
        if(n < numGroups) {
            return Optional.of(groups.get(n));
        } else {
            return Optional.empty();
        }
    }

    public ImmutableMultimap<String, Node> nodesByHost() {
        return nodesByHost;
    }

    public Optional<Node> directDispatchTarget() {
        return Optional.empty();
    }

    public void working(Node node) {
        node.setWorking(true);
    }

    public void failed(Node node) {
        node.setWorking(false);
    }
}