Add tests for LSTM class

Author: sergey-misuk-valor

tests/contrib/name_parser/test_parser.py

@@ -0,0 +1,73 @@
+from collections.abc import Callable
+from unittest.mock import Mock, MagicMock
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+
+from country_workspace.contrib.name_parser.parser import LSTM
+
+
+@pytest.fixture
+def embedding_mock(mocker: MockerFixture) -> torch.nn.Embedding:
+    return mocker.patch("country_workspace.contrib.name_parser.parser.nn.Embedding")
+
+
+@pytest.fixture
+def lstm_mock(mocker: MockerFixture) -> torch.nn.LSTM:
+    return mocker.patch("country_workspace.contrib.name_parser.parser.nn.LSTM")
+
+
+@pytest.fixture
+def linear_mock(mocker: MockerFixture) -> torch.nn.Linear:
+    return mocker.patch("country_workspace.contrib.name_parser.parser.nn.Linear")
+
+
+@pytest.fixture
+def log_softmax_mock(mocker: MockerFixture) -> torch.nn.LogSoftmax:
+    return mocker.patch("country_workspace.contrib.name_parser.parser.nn.LogSoftmax")
+
+
+@pytest.fixture
+def zeros_mock(mocker: MockerFixture) -> torch.Tensor:
+    return mocker.patch("country_workspace.contrib.name_parser.parser.torch.zeros")
+
+
+def test_lstm_class_init(
+    embedding_mock: torch.nn.Embedding,
+    lstm_mock: torch.nn.LSTM,
+    linear_mock: torch.nn.Linear,
+    log_softmax_mock: torch.nn.LogSoftmax,
+) -> None:
+    instance = Mock(spec=LSTM)
+    input_size, hidden_size, output_size, num_layers = range(4)
+
+    LSTM.__init__(instance, input_size, hidden_size, output_size, num_layers)
+
+    assert instance.hidden_size == hidden_size
+    assert instance.num_layers == num_layers
+    assert instance.embedding == embedding_mock.return_value
+    assert instance.lstm == lstm_mock.return_value
+    assert instance.fc == linear_mock.return_value
+    assert instance.softmax == log_softmax_mock.return_value
+
+    embedding_mock.assert_called_with(input_size, hidden_size)
+    lstm_mock.assert_called_once_with(hidden_size, hidden_size, num_layers, batch_first=True)
+    linear_mock.assert_called_once_with(hidden_size, output_size)
+    log_softmax_mock.assert_called_once_with(dim=1)
+
+
+def test_lstm_class_forward(
+    embedding_mock: torch.nn.Embedding,
+    lstm_mock: torch.nn.LSTM,
+    linear_mock: torch.nn.Linear,
+    log_softmax_mock: torch.nn.LogSoftmax,
+    zeros_mock: Callable,
+) -> None:
+    instance = Mock(spec=LSTM)
+    input_ = Mock(spec=torch.Tensor)
+    input_size, hidden_size, output_size, num_layers = range(4)
+    lstm_mock.return_value.return_value = MagicMock(), None
+    LSTM.__init__(instance, input_size, hidden_size, output_size, num_layers)
+
+    assert LSTM.forward(instance, input_) == instance.softmax.return_value