from unittest import TestCase
import torch

import transformers

from model.bert_model import BertModel


class TestBertModel(TestCase):

    def test_forward(self):
        # Bert Config
        vocab_size = 10
        sequence_len = 20
        batch = 32
        num_classes = 3

        expected_shape = (batch, sequence_len, num_classes)

        input_batch = torch.randint(low=0, high=vocab_size-1, size=(batch, sequence_len))
        config= transformers.BertConfig(vocab_size=vocab_size,hidden_size=10, num_hidden_layers=1, num_attention_heads=1,num_labels=num_classes)
        sut = BertModel(None, None, bert_config=config)

        # Act
        actual = sut.forward(input_batch)[0]

        # Assert
        self.assertEqual(expected_shape, actual.shape)

[(20, 'torch.randint', 'torch.randint', False, 'import torch\n'), (21, 'transformers.BertConfig', 'transformers.BertConfig', False, 'import transformers\n'), (22, 'model.bert_model.BertModel', 'BertModel', False, 'from model.bert_model import BertModel\n')]