fix(test): add test cases for train in api

williamfzc · Sep 23, 2019 · c009dcc · c009dcc
1 parent af05c72
commit c009dcc
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/test/test_api.py b/test/test_api.py
@@ -1,6 +1,6 @@
 import os
 
-from stagesepx.api import cut, classify, one_step
+from stagesepx.api import cut, classify, one_step, train
 from stagesepx.reporter import Reporter
 
 PROJECT_PATH = os.path.dirname(os.path.dirname(__file__))
@@ -12,10 +12,13 @@ def test_one_step():
 
 
 def test_cut_and_classify():
-    # --- cutter ---
+    # test cut
     res, data_home = cut(VIDEO_PATH)
 
-    # --- classify ---
+    # test train
+    train(data_home, "hello.pkl")
+
+    # test classify
     classify_result = classify(VIDEO_PATH, data_home)
 
     # --- draw ---

diff --git a/test/test_cli.py b/test/test_cli.py
@@ -17,6 +17,9 @@ def test_cli():
     subprocess.check_call(['stagesepx', 'cut', VIDEO_PATH])
     subprocess.check_call(['stagesepx', 'cut', VIDEO_PATH, 'output'])
 
+    logger.info('checking train ...')
+    subprocess.check_call(['stagesepx', 'train', 'output', 'model.pkl'])
+
     logger.info('checking classifier ...')
     subprocess.check_call(['stagesepx', 'classify', VIDEO_PATH, 'output'])
     shutil.rmtree('output')