py-test

test_main.py

+from scripts.common.config_parser import config
+from scripts.common.constants import ModelObjectConstants, ComponentExceptions
+import os
+from main import RandomForest
+
+
+def test1():
+    assert config['model_name'] == ModelObjectConstants.RANDOMFOREST
+
+
+def test2():
+    obj = RandomForest(config['model_name'], config['component_output_dir'], config)
+    if config['model_name'] == ModelObjectConstants.RANDOMFOREST:
+        val = obj.random_forest()
+        assert val == True
+
+
+def test3():
+    obj = RandomForest(config['model_name'], config['component_output_dir'], config)
+    if config['model_name'] == ModelObjectConstants.RANDOMFOREST:
+        val = obj.random_forest()
+        if val:
+            x = os.listdir(config['component_output_dir'])[0]
+            assert x == "random_forest.pkl"