17
17
18
18
package org .apache .uniffle .server .storage ;
19
19
20
- import java .io .BufferedReader ;
21
20
import java .io .File ;
22
21
import java .io .IOException ;
23
- import java .io .InputStreamReader ;
24
22
import java .nio .file .Files ;
25
23
import java .nio .file .Path ;
26
- import java .nio .file .Paths ;
27
24
import java .util .ArrayList ;
28
25
import java .util .Arrays ;
29
26
import java .util .Collections ;
30
27
import java .util .List ;
31
28
import java .util .Map ;
32
29
33
30
import org .apache .commons .io .FileUtils ;
34
- import org .apache .commons .lang3 .SystemUtils ;
35
31
import org .junit .jupiter .api .AfterAll ;
36
32
import org .junit .jupiter .api .BeforeAll ;
37
33
import org .junit .jupiter .api .Test ;
50
46
import org .apache .uniffle .server .ShuffleServerConf ;
51
47
import org .apache .uniffle .server .ShuffleServerMetrics ;
52
48
import org .apache .uniffle .server .ShuffleTaskInfo ;
49
+ import org .apache .uniffle .storage .common .DefaultStorageMediaProvider ;
53
50
import org .apache .uniffle .storage .common .LocalStorage ;
54
51
import org .apache .uniffle .storage .common .Storage ;
55
52
import org .apache .uniffle .storage .util .StorageType ;
@@ -285,14 +282,12 @@ public void testInitializeLocalStorage() throws IOException {
285
282
@ Test
286
283
public void testGetLocalStorageInfo () throws IOException {
287
284
Path testBaseDir = Files .createTempDirectory ("rss-test" );
288
- final Path storageBaseDir1 =
289
- Files .createDirectory (Paths .get (testBaseDir .toString (), "rss-data-1" ));
290
- final Path storageBaseDir2 =
291
- Files .createDirectory (Paths .get (testBaseDir .toString (), "rss-data-2" ));
292
- final Path storageBaseDir3 =
293
- Files .createDirectory (Paths .get (testBaseDir .toString (), "rss-data-3" ));
285
+ final Path storageBaseDir1 = Files .createDirectory (testBaseDir .resolve ("rss-data-1" ));
286
+ final Path storageBaseDir2 = Files .createDirectory (testBaseDir .resolve ("rss-data-2" ));
287
+ final Path storageBaseDir3 = Files .createDirectory (testBaseDir .resolve ("rss-data-3" ));
288
+
294
289
String [] storagePaths = {
295
- storageBaseDir1 .toString (), storageBaseDir2 .toString (), storageBaseDir3 .toString (),
290
+ storageBaseDir1 .toString (), storageBaseDir2 .toString (), storageBaseDir3 .toString ()
296
291
};
297
292
ShuffleServerConf conf = new ShuffleServerConf ();
298
293
conf .set (ShuffleServerConf .RSS_STORAGE_BASE_PATH , Arrays .asList (storagePaths ));
@@ -301,45 +296,31 @@ public void testGetLocalStorageInfo() throws IOException {
301
296
ShuffleServerConf .RSS_STORAGE_TYPE .key (),
302
297
org .apache .uniffle .storage .util .StorageType .LOCALFILE .name ());
303
298
conf .setDouble (ShuffleServerConf .HEALTH_STORAGE_MAX_USAGE_PERCENTAGE , 100.0D );
304
- LocalStorageManager localStorageManager = new LocalStorageManager (conf );
305
- // Create and write 3 files in each storage dir
306
- final Path file1 = Files .createFile (Paths .get (storageBaseDir1 .toString (), "partition1.data" ));
307
- final Path file2 = Files .createFile (Paths .get (storageBaseDir2 .toString (), "partition2.data" ));
308
- final Path file3 = Files .createFile (Paths .get (storageBaseDir3 .toString (), "partition3.data" ));
309
- FileUtils .writeByteArrayToFile (file1 .toFile (), new byte [] {0x1 });
310
- FileUtils .writeByteArrayToFile (file2 .toFile (), new byte [] {0x2 });
311
- FileUtils .writeByteArrayToFile (file3 .toFile (), new byte [] {0x3 });
312
-
313
- boolean healthy = localStorageManager .getStorageChecker ().checkIsHealthy ();
314
- assertTrue (healthy , "should be healthy" );
299
+
300
+ // Write one file to each storage dir
301
+ final LocalStorageManager localStorageManager = new LocalStorageManager (conf );
302
+ Files .write (storageBaseDir1 .resolve ("partition1.data" ), new byte [] {0x1 });
303
+ Files .write (storageBaseDir2 .resolve ("partition2.data" ), new byte [] {0x2 });
304
+ Files .write (storageBaseDir3 .resolve ("partition3.data" ), new byte [] {0x3 });
305
+
306
+ assertTrue (localStorageManager .getStorageChecker ().checkIsHealthy (), "should be healthy" );
307
+
315
308
Map <String , StorageInfo > storageInfo = localStorageManager .getStorageInfo ();
316
309
assertEquals (1 , storageInfo .size ());
317
- try {
318
- final String path = testBaseDir .toString ();
319
- final String mountPoint = Files .getFileStore (new File (path ).toPath ()).name ();
320
- assertNotNull (storageInfo .get (mountPoint ));
321
- // on Linux environment, it can detect SSD as local storage type
322
- if (SystemUtils .IS_OS_LINUX ) {
323
- final String cmd =
324
- String .format (
325
- "%s | %s | %s" ,
326
- "lsblk -a -o name,rota" ,
327
- "grep $(df --output=source " + path + " | tail -n 1 | sed -E 's_^.+/__')" ,
328
- "awk '{print $2}'" );
329
- Process process = Runtime .getRuntime ().exec (new String [] {"bash" , "-c" , cmd });
330
- BufferedReader br = new BufferedReader (new InputStreamReader (process .getInputStream ()));
331
- final String line = br .readLine ();
332
- br .close ();
333
- final StorageMedia expected = "0" .equals (line ) ? StorageMedia .SSD : StorageMedia .HDD ;
334
- assertEquals (expected , storageInfo .get (mountPoint ).getType ());
335
- } else {
336
- assertEquals (StorageMedia .HDD , storageInfo .get (mountPoint ).getType ());
337
- }
338
- assertEquals (StorageStatus .NORMAL , storageInfo .get (mountPoint ).getStatus ());
339
- assertEquals (3L , storageInfo .get (mountPoint ).getUsedBytes ());
340
- } catch (IOException e ) {
341
- throw new RuntimeException (e );
342
- }
310
+
311
+ String mountPoint = Files .getFileStore (testBaseDir ).name ();
312
+ assertNotNull (storageInfo .get (mountPoint ));
313
+
314
+ // Use production logic to get expected media type
315
+ StorageMedia expected =
316
+ new DefaultStorageMediaProvider ().getStorageMediaFor (testBaseDir .toString ());
317
+
318
+ // Assert media type from storage manager matches production logic
319
+ assertEquals (expected , storageInfo .get (mountPoint ).getType ());
320
+
321
+ // Assert status and used bytes
322
+ assertEquals (StorageStatus .NORMAL , storageInfo .get (mountPoint ).getStatus ());
323
+ assertEquals (3L , storageInfo .get (mountPoint ).getUsedBytes ());
343
324
}
344
325
345
326
@ Test
0 commit comments