Network Modelling Tool in Python: Part 2

Photo by Taylor Vick on Unsplash

Network Modelling Tool in Python: Part 2

In part 1 of this series, we managed to load the load CSV file of the network into a suitable data structure, an adjacent matrix. In this second part, we'd like to create tests for what we have worked on. Test-driven development is important because it ensures the code is thoroughly tested and reliable, reducing the likelihood of bugs and facilitating easier maintenance.

To get started writing tests, let's create a new folder called tests. Inside the tests folder, create a file, test_graph.py and add the following code:

import unittest
import sys

sys.path.append(".")

from graph.graph import Graph

We are importing the modules and the Graph class we will need to write tests. The sys.path.append(".") appends the current directory to the list of paths that Python searches for modules and packages. This is important because it enables our script to import the Graph class.

Add the following code to the test_graph.py:

...

class TestGraph(unittest.TestCase):
    def setUp(self) -> None:
        self.graph = Graph()

    def test_add_vertex(self) -> None:
        self.graph.add_vertex("A")
        actual = self.graph.get_vertices()
        expected = ["A"]
        self.assertEqual(actual, expected)


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

We have created a test class that inherits from unittes.TestCase. The class will contain all the methods that test the Graph class.

The setUp method runs before each test method. It initializes a new instance of the Graph class, ensuring each test runs in isolation using a fresh instance.

Next, we create a test method for the add_vertex method. This adds a vertex/node, A, to the graph. You'll notice that I'm calling the get_vertices method from the graph class. We didn't implement this method previously. It gets a list of all the vertices in the graph and stores it in the variable, actual. We are expecting to have a list containing a single vertex, A. The assertion method compares the actual list of vertices with the expected list of vertices. If they are not equal, the test will fail.

The method tests whether adding a vertex, A, to the graph results in the expected list of vertices.

Let's go back to our graph.py file and add the get_vertices method:

    def get_vertices(self):
        """Get the nodes."""
        return self.vertices

Reaching this far, we have implemented the test method for the add_vertex method in the Graph class. We can run the test using the following command:

python3 -m unittest tests.test_graph

Let's create more test methods to test the graph class:

...

from unittest.mock import patch


class TestGraph(unittest.TestCase):

    ...

    @patch("builtins.print")
    def test_add_vertex_already_exists(self, mock_print) -> None:
        self.graph.add_vertex("A")
        self.graph.add_vertex("A")
        mock_print.assert_called_with("Vertex ", "A", " already exists")

    def test_add_edge(self) -> None:
        self.graph.add_vertex("A")
        self.graph.add_vertex("B")
        self.graph.add_edge("A", "B", 5)
        actual = self.graph.get_weight(0, 1)
        expected = (5, 5)
        self.assertEqual(actual, expected)

    @patch("builtins.print")
    def test_add_edge_node_does_not_exists(self, mock_print) -> None:
        self.graph.add_vertex("A")
        self.graph.add_vertex("B")
        self.graph.add_edge("A", "B", 5)
        self.graph.add_edge("A", "C", 2)
        mock_print.assert_called_with("Vertex ", "C", " does not exist.")

We are importing the patch function which we will use for mocking objects. The @patch("builtins.print") decorator is used to mock the built-in print function with a mock object, mock_print, during the execution of the test method.

In the test_add_vertex_already_exists method, we pass in the mock_print parameter, a mock object that replaces the print function. We call the add_vertex method twice on the graph object to add a vertex, A. According to the behavior we defined in the Graph class, the second method call should trigger a print statement indicating that the vertex already exists. The assert_called_with method checks if the correct message is printed when adding a vertex to the graph.

Moving on, the test_add_edge method verifies that the Graph class correctly adds an edge between two vertices. It does this by adding two vertices A and B and adds an edge between them with a weight of 5. Next, it retrieves the weight of the edges between the two vertices and asserts that the retrieved weight matches the expected weight.

The test_add_edge_node_does_not_exists method verifies that the Graph class correctly handles the case of adding an edge where one of the vertices does not exist. It adds two vertices A and B and adds a valid edge between the two vertices. Then, it attempts to add an edge between A and a non-existent vertex C. We use unittest.mock.patch to mock the print function and check if the appropriate message Vertex C does not exist is printed when the non-existent vertex C is referenced.

In our tests, we have used the get_weight method from the Graph class. We did not implement this method previously, let's go back to the Graph class in the graph.py file and add the method:

    def get_weight(self, row: int, col: int) -> None:
        """Gets the weight between two nodes.

        Args:
        row: int. The row number in adjacent matrix.
        col: int. The col umber in adjacent matrix.
        """
        return (self.graph[row][col], self.graph[col][row])

Reaching this far, we have successfully implemented tests for the methods in the Graph class. In the next article, we'll determine the path traffic will take.

All the code for this second part can be found on GitHub.

Please Like, Comment, or Share if you found this article helpful.