From e8b5ce9a10130e4297c73eb416a7e758937d2bb5 Mon Sep 17 00:00:00 2001
From: Ning Shang <nishang@microsoft.com>
Date: Mon, 21 Sep 2020 09:34:47 +0800
Subject: [PATCH] add /experiment/imported-data & add verify condition

---
 src/nni_manager/common/manager.ts                    |  1 +
 src/nni_manager/core/nnimanager.ts                   |  4 ++++
 src/nni_manager/rest_server/restHandler.ts           | 11 +++++++++++
 src/nni_manager/rest_server/test/mockedNNIManager.ts |  4 ++++
 test/nni_test/nnitest/utils.py                       |  1 +
 test/nni_test/nnitest/validators.py                  |  6 +++++-
 6 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/nni_manager/common/manager.ts b/src/nni_manager/common/manager.ts
index c003598abc..1f3972ae43 100644
--- a/src/nni_manager/common/manager.ts
+++ b/src/nni_manager/common/manager.ts
@@ -87,6 +87,7 @@ abstract class Manager {
     public abstract getExperimentProfile(): Promise<ExperimentProfile>;
     public abstract updateExperimentProfile(experimentProfile: ExperimentProfile, updateType: ProfileUpdateType): Promise<void>;
     public abstract importData(data: string): Promise<void>;
+    public abstract getImportedData(): Promise<string[]>;
     public abstract exportData(): Promise<string>;
 
     public abstract addCustomizedTrialJob(hyperParams: string): Promise<number>;
diff --git a/src/nni_manager/core/nnimanager.ts b/src/nni_manager/core/nnimanager.ts
index ad243f4835..6ec4d0e21d 100644
--- a/src/nni_manager/core/nnimanager.ts
+++ b/src/nni_manager/core/nnimanager.ts
@@ -108,6 +108,10 @@ class NNIManager implements Manager {
         return this.dataStore.storeTrialJobEvent('IMPORT_DATA', '', data);
     }
 
+    public getImportedData(): Promise<string[]> {
+        return this.dataStore.getImportedData();
+    }
+
     public async exportData(): Promise<string> {
         return this.dataStore.exportTrialHpConfigs();
     }
diff --git a/src/nni_manager/rest_server/restHandler.ts b/src/nni_manager/rest_server/restHandler.ts
index af44d71a01..1a4d87ec02 100644
--- a/src/nni_manager/rest_server/restHandler.ts
+++ b/src/nni_manager/rest_server/restHandler.ts
@@ -47,6 +47,7 @@ class NNIRestHandler {
         this.getExperimentProfile(router);
         this.updateExperimentProfile(router);
         this.importData(router);
+        this.getImportedData(router);
         this.startExperiment(router);
         this.getTrialJobStatistics(router);
         this.setClusterMetaData(router);
@@ -143,6 +144,16 @@ class NNIRestHandler {
         });
     }
 
+    private getImportedData(router: Router): void {
+        router.get('/experiment/imported-data', (req: Request, res: Response) => {
+            this.nniManager.getImportedData().then((importedData: string[]) => {
+                res.send(JSON.stringify(importedData));
+            }).catch((err: Error) => {
+                this.handleError(err, res);
+            });
+        });
+    }
+
     private startExperiment(router: Router): void {
         router.post('/experiment', expressJoi(ValidationSchemas.STARTEXPERIMENT), (req: Request, res: Response) => {
             if (isNewExperiment()) {
diff --git a/src/nni_manager/rest_server/test/mockedNNIManager.ts b/src/nni_manager/rest_server/test/mockedNNIManager.ts
index e45819d6cb..b3bd549c03 100644
--- a/src/nni_manager/rest_server/test/mockedNNIManager.ts
+++ b/src/nni_manager/rest_server/test/mockedNNIManager.ts
@@ -33,6 +33,10 @@ export class MockedNNIManager extends Manager {
     public importData(data: string): Promise<void> {
         return Promise.resolve();
     }
+    public getImportedData(): Promise<string[]> {
+        const ret: string[] = ["1", "2"];
+        return Promise.resolve(ret);
+    }
     public async exportData(): Promise<string> {
         const ret: string = '';
         return Promise.resolve(ret);
diff --git a/test/nni_test/nnitest/utils.py b/test/nni_test/nnitest/utils.py
index 795bd8aefe..bdd3fd7fea 100644
--- a/test/nni_test/nnitest/utils.py
+++ b/test/nni_test/nnitest/utils.py
@@ -24,6 +24,7 @@
 STATUS_URL = API_ROOT_URL + '/check-status'
 TRIAL_JOBS_URL = API_ROOT_URL + '/trial-jobs'
 METRICS_URL = API_ROOT_URL + '/metric-data'
+GET_IMPORTED_DATA_URL = API_ROOT_URL + '/experiment/imported-data'
 
 def read_last_line(file_name):
     '''read last line of a file and return None if file not found'''
diff --git a/test/nni_test/nnitest/validators.py b/test/nni_test/nnitest/validators.py
index af554ad781..ff349ebc2b 100644
--- a/test/nni_test/nnitest/validators.py
+++ b/test/nni_test/nnitest/validators.py
@@ -7,7 +7,8 @@
 import json
 import requests
 from nnicli import Experiment
-from utils import METRICS_URL
+from nni_cmd.updater import load_search_space
+from utils import METRICS_URL, GET_IMPORTED_DATA_URL
 
 
 class ITValidator:
@@ -40,6 +41,9 @@ def __call__(self, rest_endpoint, experiment_dir, nni_source_dir, **kwargs):
         proc = subprocess.run(['nnictl', 'experiment', 'import', exp_id, '-f', import_data_file_path])
         assert proc.returncode == 0, \
             '`nnictl experiment import {0} -f {1}` failed with code {2}'.format(exp_id, import_data_file_path, proc.returncode)
+        imported_data = requests.get(GET_IMPORTED_DATA_URL).json()
+        origin_data = load_search_space(import_data_file_path).replace(' ', '')
+        assert origin_data in imported_data
 
 class MetricsValidator(ITValidator):
     def __call__(self, rest_endpoint, experiment_dir, nni_source_dir, **kwargs):