10
10
import static org .mockito .ArgumentMatchers .eq ;
11
11
import static org .mockito .Mockito .doAnswer ;
12
12
import static org .mockito .Mockito .doThrow ;
13
- import static org .mockito .Mockito .mock ;
14
13
import static org .mockito .Mockito .never ;
15
- import static org .mockito .Mockito .spy ;
16
14
import static org .mockito .Mockito .times ;
17
15
import static org .mockito .Mockito .verify ;
18
16
import static org .mockito .Mockito .when ;
19
- import static org .opensearch .ml .common .CommonValue .CREATE_TIME_FIELD ;
20
- import static org .opensearch .ml .common .CommonValue .MASTER_KEY ;
21
17
import static org .opensearch .ml .common .CommonValue .ML_MODEL_INDEX ;
22
18
import static org .opensearch .ml .utils .TestHelper .ML_ROLE ;
23
19
import static org .opensearch .ml .utils .TestHelper .setupTestClusterState ;
34
30
import java .util .concurrent .atomic .AtomicInteger ;
35
31
36
32
import org .apache .lucene .search .TotalHits ;
37
- import org .junit .Assert ;
38
33
import org .junit .Before ;
39
34
import org .mockito .ArgumentCaptor ;
40
35
import org .mockito .Mock ;
41
36
import org .mockito .Mockito ;
42
37
import org .mockito .MockitoAnnotations ;
43
38
import org .opensearch .Version ;
44
39
import org .opensearch .action .bulk .BulkRequest ;
45
- import org .opensearch .action .get .GetResponse ;
46
- import org .opensearch .action .index .IndexResponse ;
47
40
import org .opensearch .action .search .SearchResponse ;
48
41
import org .opensearch .action .search .SearchResponseSections ;
49
42
import org .opensearch .action .search .ShardSearchFailure ;
71
64
import org .opensearch .ml .common .transport .sync .MLSyncUpAction ;
72
65
import org .opensearch .ml .common .transport .sync .MLSyncUpNodeResponse ;
73
66
import org .opensearch .ml .common .transport .sync .MLSyncUpNodesResponse ;
74
- import org .opensearch .ml .engine .encryptor .Encryptor ;
75
- import org .opensearch .ml .engine .encryptor .EncryptorImpl ;
76
67
import org .opensearch .ml .engine .indices .MLIndicesHandler ;
77
68
import org .opensearch .ml .utils .TestHelper ;
78
69
import org .opensearch .remote .metadata .client .SdkClient ;
87
78
import org .opensearch .threadpool .ThreadPool ;
88
79
import org .opensearch .transport .client .Client ;
89
80
90
- import com .google .common .collect .ImmutableMap ;
91
81
import com .google .common .collect .ImmutableSet ;
92
82
93
83
public class MLSyncUpCronTests extends OpenSearchTestCase {
@@ -112,7 +102,6 @@ public class MLSyncUpCronTests extends OpenSearchTestCase {
112
102
private final String mlNode2Id = "mlNode2" ;
113
103
114
104
private ClusterState testState ;
115
- private Encryptor encryptor ;
116
105
117
106
@ Mock
118
107
ThreadPool threadPool ;
@@ -124,64 +113,17 @@ public void setup() throws IOException {
124
113
MockitoAnnotations .openMocks (this );
125
114
mlNode1 = new DiscoveryNode (mlNode1Id , buildNewFakeTransportAddress (), emptyMap (), ImmutableSet .of (ML_ROLE ), Version .CURRENT );
126
115
mlNode2 = new DiscoveryNode (mlNode2Id , buildNewFakeTransportAddress (), emptyMap (), ImmutableSet .of (ML_ROLE ), Version .CURRENT );
127
- encryptor = spy (new EncryptorImpl (null , "m+dWmfmnNRiNlOdej/QelEkvMTyH//frS2TBeS2BP4w=" ));
128
116
129
117
testState = setupTestClusterState ("node" );
130
118
when (clusterService .state ()).thenReturn (testState );
131
119
132
- doAnswer (invocation -> {
133
- ActionListener <Boolean > actionListener = invocation .getArgument (0 );
134
- actionListener .onResponse (true );
135
- return null ;
136
- }).when (mlIndicesHandler ).initMLConfigIndex (any ());
137
-
138
120
Settings settings = Settings .builder ().build ();
139
121
sdkClient = Mockito .spy (SdkClientFactory .createSdkClient (client , NamedXContentRegistry .EMPTY , Collections .emptyMap ()));
140
122
threadContext = new ThreadContext (settings );
141
123
threadContext .putTransient (ConfigConstants .OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT , USER_STRING );
142
124
when (client .threadPool ()).thenReturn (threadPool );
143
125
when (threadPool .getThreadContext ()).thenReturn (threadContext );
144
- syncUpCron = new MLSyncUpCron (client , sdkClient , clusterService , nodeHelper , mlIndicesHandler , encryptor , mlFeatureEnabledSetting );
145
- }
146
-
147
- public void testInitMlConfig_MasterKeyNotExist () {
148
- doAnswer (invocation -> {
149
- ActionListener <GetResponse > listener = invocation .getArgument (1 );
150
- GetResponse response = mock (GetResponse .class );
151
- when (response .isExists ()).thenReturn (false );
152
- listener .onResponse (response );
153
- return null ;
154
- }).when (client ).get (any (), any ());
155
-
156
- doAnswer (invocation -> {
157
- ActionListener <IndexResponse > listener = invocation .getArgument (1 );
158
- IndexResponse indexResponse = mock (IndexResponse .class );
159
- listener .onResponse (indexResponse );
160
- return null ;
161
- }).when (client ).index (any (), any ());
162
-
163
- syncUpCron .initMLConfig ();
164
- Assert .assertNotNull (encryptor .encrypt ("test" , null ));
165
- syncUpCron .initMLConfig ();
166
- verify (encryptor , times (1 )).setMasterKey (any (), any ());
167
- }
168
-
169
- public void testInitMlConfig_MasterKeyExists () {
170
- doAnswer (invocation -> {
171
- ActionListener <GetResponse > listener = invocation .getArgument (1 );
172
- GetResponse response = mock (GetResponse .class );
173
- when (response .isExists ()).thenReturn (true );
174
- String masterKey = encryptor .generateMasterKey ();
175
- when (response .getSourceAsMap ())
176
- .thenReturn (ImmutableMap .of (MASTER_KEY , masterKey , CREATE_TIME_FIELD , Instant .now ().toEpochMilli ()));
177
- listener .onResponse (response );
178
- return null ;
179
- }).when (client ).get (any (), any ());
180
-
181
- syncUpCron .initMLConfig ();
182
- Assert .assertNotNull (encryptor .encrypt ("test" , null ));
183
- syncUpCron .initMLConfig ();
184
- verify (encryptor , times (1 )).setMasterKey (any (), any ());
126
+ syncUpCron = new MLSyncUpCron (client , sdkClient , clusterService , nodeHelper , mlIndicesHandler , mlFeatureEnabledSetting );
185
127
}
186
128
187
129
public void testRun_NoMLModelIndex () {
0 commit comments