test_line_finder.py

import unittest
import pandas as pd
from line_finder import LineFinder  # Adjust the import based on your actual module structure

class TestLineFinder(unittest.TestCase):

    def test_get_the_lines(self):
        # Create a sample DataFrame for testing
        sample_data = pd.DataFrame({
            'Line': ['720', '150'],
            'Stop name': ['Zurich Wollishofen', 'Allaman']
        })

        # Initialize an instance of LineFinder
        line_finder = LineFinder()

        # Override the data attribute of the instance for testing
        line_finder.data = sample_data

        # Call the function to be tested
        result1 = line_finder.get_the_lines('Zurich Wollishofen')
        result2 =line_finder.get_the_lines('Allaman')

        # Define the expected result
        expected_result1 = {'720'}
        expected_result2 ={'150'}

        # Check if the result matches the expected output
        self.assertEqual(result1, expected_result1)
        self.assertEqual(result2, expected_result2)

if __name__ == '__main__':
    unittest.main()