From 780aefa6fee9c0d551074542c082cc95bde7f399 Mon Sep 17 00:00:00 2001
From: Evgeny Smyshlyaev <amper43@corpglory.com>
Date: Mon, 12 Nov 2018 20:30:25 +0300
Subject: [PATCH] Tests for #231 (#230)

---
 analytics/tests/__init__.py     |  5 ++++-
 analytics/tests/test_dataset.py | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)
 create mode 100644 analytics/tests/test_dataset.py

diff --git a/analytics/tests/__init__.py b/analytics/tests/__init__.py
index 2ae2839..bcc8e88 100644
--- a/analytics/tests/__init__.py
+++ b/analytics/tests/__init__.py
@@ -1 +1,4 @@
-pass
+import sys
+import os
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'analytics'))
diff --git a/analytics/tests/test_dataset.py b/analytics/tests/test_dataset.py
new file mode 100644
index 0000000..e6f84fe
--- /dev/null
+++ b/analytics/tests/test_dataset.py
@@ -0,0 +1,33 @@
+import unittest
+import pandas as pd
+from analytic_unit_manager import prepare_data
+import models
+
+class TestDataset(unittest.TestCase):
+
+    def test_models_with_corrupted_dataframe(self):
+        data = [[1523889000000 + i, float('nan')] for i in range(10)]
+        dataframe = pd.DataFrame(data, columns=['timestamp', 'value'])
+        segments = []
+
+        model_instances = [
+            models.JumpModel(),
+            models.DropModel(),
+            models.GeneralModel(),
+            models.PeakModel(),
+            models.TroughModel()
+        ]
+        try:
+            for model in model_instances:
+                model_name = model.__class__.__name__
+                model.fit(dataframe, segments, dict())
+        except ValueError:
+            self.fail('Model {} raised unexpectedly'.format(model_name))
+
+    def test_data_preparation(self):
+        data = [[1523889000000 + i, float('nan')] for i in range(10)]
+
+        self.assertTrue(prepare_data(data).empty) # TODO: raise exception
+
+if __name__ == '__main__':
+    unittest.main()