1 package org.apache.helix.controller.stages;
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22 import java.util.Map;
23
24 import org.apache.helix.ZNRecord;
25 import org.apache.helix.PropertyKey.Builder;
26 import org.apache.helix.controller.stages.AttributeName;
27 import org.apache.helix.controller.stages.CurrentStateComputationStage;
28 import org.apache.helix.controller.stages.CurrentStateOutput;
29 import org.apache.helix.controller.stages.ReadClusterDataStage;
30 import org.apache.helix.model.CurrentState;
31 import org.apache.helix.model.Message;
32 import org.apache.helix.model.Partition;
33 import org.apache.helix.model.Resource;
34 import org.testng.AssertJUnit;
35 import org.testng.annotations.Test;
36
37
38 public class TestCurrentStateComputationStage extends BaseStageTest
39 {
40
41 @Test
42 public void testEmptyCS()
43 {
44 Map<String, Resource> resourceMap = getResourceMap();
45 event.addAttribute(AttributeName.RESOURCES.toString(), resourceMap);
46 CurrentStateComputationStage stage = new CurrentStateComputationStage();
47 runStage(event, new ReadClusterDataStage());
48 runStage(event, stage);
49 CurrentStateOutput output =
50 event.getAttribute(AttributeName.CURRENT_STATE.toString());
51 AssertJUnit.assertEquals(output.getCurrentStateMap("testResourceName",
52 new Partition("testResourceName_0"))
53 .size(),
54 0);
55 }
56
57 @Test
58 public void testSimpleCS()
59 {
60
61 Map<String, Resource> resourceMap = getResourceMap();
62
63 setupLiveInstances(5);
64
65 event.addAttribute(AttributeName.RESOURCES.toString(), resourceMap);
66 CurrentStateComputationStage stage = new CurrentStateComputationStage();
67 runStage(event, new ReadClusterDataStage());
68 runStage(event, stage);
69 CurrentStateOutput output1 =
70 event.getAttribute(AttributeName.CURRENT_STATE.toString());
71 AssertJUnit.assertEquals(output1.getCurrentStateMap("testResourceName",
72 new Partition("testResourceName_0"))
73 .size(),
74 0);
75
76
77 Message message = new Message(Message.MessageType.STATE_TRANSITION, "msg1");
78 message.setFromState("OFFLINE");
79 message.setToState("SLAVE");
80 message.setResourceName("testResourceName");
81 message.setPartitionName("testResourceName_1");
82 message.setTgtName("localhost_3");
83 message.setTgtSessionId("session_3");
84
85 Builder keyBuilder = accessor.keyBuilder();
86 accessor.setProperty(keyBuilder.message("localhost_" + 3, message.getId()), message);
87
88 runStage(event, new ReadClusterDataStage());
89 runStage(event, stage);
90 CurrentStateOutput output2 =
91 event.getAttribute(AttributeName.CURRENT_STATE.toString());
92 String pendingState =
93 output2.getPendingState("testResourceName",
94 new Partition("testResourceName_1"),
95 "localhost_3");
96 AssertJUnit.assertEquals(pendingState, "SLAVE");
97
98 ZNRecord record1 = new ZNRecord("testResourceName");
99
100 CurrentState stateWithLiveSession = new CurrentState(record1);
101 stateWithLiveSession.setSessionId("session_3");
102 stateWithLiveSession.setStateModelDefRef("MasterSlave");
103 stateWithLiveSession.setState("testResourceName_1", "OFFLINE");
104 ZNRecord record2 = new ZNRecord("testResourceName");
105 CurrentState stateWithDeadSession = new CurrentState(record2);
106 stateWithDeadSession.setSessionId("session_dead");
107 stateWithDeadSession.setStateModelDefRef("MasterSlave");
108 stateWithDeadSession.setState("testResourceName_1", "MASTER");
109
110 accessor.setProperty(keyBuilder.currentState("localhost_3",
111 "session_3",
112 "testResourceName"),
113 stateWithLiveSession);
114 accessor.setProperty(keyBuilder.currentState("localhost_3",
115 "session_dead",
116 "testResourceName"),
117 stateWithDeadSession);
118 runStage(event, new ReadClusterDataStage());
119 runStage(event, stage);
120 CurrentStateOutput output3 =
121 event.getAttribute(AttributeName.CURRENT_STATE.toString());
122 String currentState =
123 output3.getCurrentState("testResourceName",
124 new Partition("testResourceName_1"),
125 "localhost_3");
126 AssertJUnit.assertEquals(currentState, "OFFLINE");
127
128 }
129
130 }