#23 Add tests and improve filesystem connector

Author: OlivieFranklova

similarityRunner/connectors/filesystem_connector.py

@@ -5,24 +5,21 @@
 
 from functionsRunner import load_files_from_list
 from interfaces.ConnectorInterface import ConnectorInterface
-from models.connector_models import ConnectorSettings, Output, ConnectorOutput, FSConnectorSettings
+from models.connector_models import Output, ConnectorOutput, FSConnectorSettings
 
 
 class FilesystemConnector(ConnectorInterface):
-    def __init__(self, config):
-        self.config = config
 
     def _connect_and_load_data_source(self, settings: FSConnectorSettings) -> ConnectorOutput:
         file_list = settings.files_paths
         for folder in settings.directory_paths:
-            file_list = file_list + [folder + "/" +  s for s in os.listdir(folder)]
+            file_list = file_list + [os.path.join(folder, s) for s in os.listdir(folder)]
 
-        names, tables = load_files_from_list(os.listdir(file_list), settings.file_type)
+        tables, names = load_files_from_list(file_list, settings.file_type)
         return ConnectorOutput(names=names, tables=tables)
 
     def _format_data(self, data: ConnectorOutput) -> Output:
-        pass
+        return data.tables, data.names
 
     def close(self):
         pass
-

similarityRunner/functionsRunner.py

@@ -3,12 +3,9 @@
 from models.connector_models import FileType
 
 
-def load_files_from_list(folder: list[str], file_type: tuple[FileType] = FileType.CSV) -> tuple[list[pd.DataFrame], list[str]]:
+def load_files_from_list(folder: list[str], file_type: tuple = (FileType.CSV, )) -> tuple[list[pd.DataFrame], list[str]]:
     """
-    it loads cvs files from folder and returns list of loaded dataframe and list of names
-    :param folder: from which we load the files
-    :param file_type: type of file, csv, parquet, etc.
-    :return: two lists
+
     """
     data = []
     names = []
@@ -20,3 +17,8 @@ def load_files_from_list(folder: list[str], file_type: tuple[FileType] = FileTyp
             data.append(pd.read_parquet(file))
             names.append(file.replace(".parquet", ""))
     return data, names
+
+def csv_to_parquet(file: str):
+    df = pd.read_csv(file)
+    df.to_parquet(file.replace(".csv", ".parquet"))
+    return file.replace(".csv", ".parquet")

similarityRunner/models/connector_models.py

@@ -8,7 +8,7 @@
 import pandas as pd
 from pydantic import BaseModel
 
-Output = pd.DataFrame
+Output = tuple[list[pd.DataFrame], list[str]]
 
 class FileType(Enum):
     CSV = "csv"
@@ -20,17 +20,23 @@ class ConnectorSettings(BaseModel):
     ConnectorSettings class is a base class for connector settings.
     """
 
-    # here will be common fields for all connectors
     file_type: tuple[FileType] # csv, parquet, etc., tuple for immutability
+    class Config:
+        # arbitrary_types_allowed is set to True to allow tuple FileType
+        arbitrary_types_allowed = True
 
 
 class ConnectorOutput(BaseModel):
     """
     ConnectorOutput class is a base class for connector output.
     """
     names: list[str]
-    tables: tuple[list[pd.DataFrame]]
-    # here will be common fields for all connectors
+    tables: list[pd.DataFrame]
+
+    class Config:
+        # arbitrary_types_allowed is set to True to allow list of pandas DataFrames
+        arbitrary_types_allowed = True
+
 
 class FSConnectorSettings(ConnectorSettings):
     """

test/test_connectors.py

@@ -0,0 +1,25 @@
+import unittest
+
+from connectors.filesystem_connector import FilesystemConnector
+from models.connector_models import FSConnectorSettings
+
+
+class TestFileSystemConnector(unittest.TestCase):
+    def test_get_data_files(self):
+        connector = FilesystemConnector()
+        settings = FSConnectorSettings(files_paths=["../data/netflix_titles.csv",
+                                                  "../data/disney_movies.csv"],
+                                       directory_paths=[],
+                                     file_type=("csv",))
+        data, names = connector.get_data(settings)
+        self.assertEqual(len(data), 2)
+        self.assertEqual(names[0], "../data/netflix_titles")
+        self.assertEqual(names[1], "../data/disney_movies")
+
+    def test_get_data_folder(self):
+        connector = FilesystemConnector()
+        settings = FSConnectorSettings(files_paths=[],
+                                       directory_paths=["../data"],
+                                       file_type=("csv",))
+        data, _ = connector.get_data(settings)
+        self.assertEqual(len(data), 11)

test/test_functions_runner.py

@@ -0,0 +1,40 @@
+import unittest
+
+from functionsRunner import load_files_from_list, csv_to_parquet
+from models.connector_models import FileType
+
+
+class TestLoadFilesFromList(unittest.TestCase):
+    def test_load_csv_file(self):
+        data, names = load_files_from_list(["../data/netflix_titles.csv"], (FileType.CSV, ))
+        self.assertEqual(len(data), 1)
+        self.assertEqual(names[0], "../data/netflix_titles")
+
+    def test_load_csv_files(self):
+        data, names = load_files_from_list(["../data/netflix_titles.csv", "../data/disney_movies.csv"], (FileType.CSV, ))
+        self.assertEqual(len(data), 2)
+        self.assertEqual(names[0], "../data/netflix_titles")
+        self.assertEqual(names[1], "../data/disney_movies")
+
+
+    def test_load_parquet_file(self):
+        csv_to_parquet("../data/netflix_titles.csv")
+        data, names = load_files_from_list(["../data/netflix_titles.parquet"], (FileType.PARQUET, ))
+        self.assertEqual(len(data), 1)
+        self.assertEqual(names[0], "../data/netflix_titles")
+
+    def test_load_parquet_files(self):
+        csv_to_parquet("../data/netflix_titles.csv")
+        csv_to_parquet("../data/disney_movies.csv")
+        data, names = load_files_from_list(["../data/netflix_titles.parquet", "../data/disney_movies.parquet"], (FileType.PARQUET, ))
+        self.assertEqual(len(data), 2)
+        self.assertEqual(names[0], "../data/netflix_titles")
+        self.assertEqual(names[1], "../data/disney_movies")
+
+
+    def test_load_csv_and_parquet_files(self):
+        csv_to_parquet("../data/netflix_titles.csv")
+        data, names = load_files_from_list(["../data/netflix_titles.parquet", "../data/disney_movies.csv"], (FileType.PARQUET, FileType.CSV))
+        self.assertEqual(len(data), 2)
+        self.assertEqual(names[0], "../data/netflix_titles")
+        self.assertEqual(names[1], "../data/disney_movies")