from unittest import TestCase

from scripts.migrate_to_whatsapp_templates.pmtct_postbirth2 import (
    PMTCTPostbirth2Migration,
)


class PMTCTPostbirth2MigrationTests(TestCase):
    def setUp(self):
        self.pmtct_postbirth2 = PMTCTPostbirth2Migration()

    def test_get_template_name(self):
        """
        If the sequence number is < 7, should be the first postbirth
        template, else should be the second postbirth template.
        """
        self.assertEqual(
            self.pmtct_postbirth2.get_template_name({"sequence_number": "6"}),
            "mc_postbirth1",
        )
        self.assertEqual(
            self.pmtct_postbirth2.get_template_name({"sequence_number": "7"}),
            "mc_postbirth2",
        )

    def test_get_template_variables_up_to_7_weeks(self):
        """
        Up to 7 weeks, the first variable should be the number of weeks, and the second
        should be the message content.
        """
        self.assertEqual(
            self.pmtct_postbirth2.get_template_variables(
                {"sequence_number": "5", "text_content": "Test message"}
            ),
            ["6", "Test message"],
        )

    def test_get_template_variables_over_7_weeks(self):
        """
        From 8 weeks onwards, the first variable should be the number of months, and
        the second should be the message content.
        """
        self.assertEqual(
            self.pmtct_postbirth2.get_template_variables(
                {"sequence_number": "7", "text_content": "Test message"}
            ),
            ["1", "Test message"],
        )

    def test_get_weeks_from_sequence_number(self):
        """
        Should return the correct number of weeks for the sequence number of the message
        """
        self.assertEqual(self.pmtct_postbirth2.get_weeks_from_sequence_number(1), 2)
        self.assertEqual(self.pmtct_postbirth2.get_weeks_from_sequence_number(2), 3)
        self.assertEqual(self.pmtct_postbirth2.get_weeks_from_sequence_number(6), 7)

    def test_get_months_from_sequence_number(self):
        """
        Returns correct number of months for the given sequence number
        """
        self.assertEqual(self.pmtct_postbirth2.get_months_from_sequence_number(1), 0)
        self.assertEqual(self.pmtct_postbirth2.get_months_from_sequence_number(7), 1)
        self.assertEqual(self.pmtct_postbirth2.get_months_from_sequence_number(50), 11)

from unittest import TestCase

from scripts.migrate_to_whatsapp_templates.pmtct_postbirth2 import (
    PMTCTPostbirth2Migration,
)


class PMTCTPostbirth2MigrationTests(TestCase):
    def setUp(self):
        self.pmtct_postbirth2 = PMTCTPostbirth2Migration()

    def test_get_template_name(self):
        """
        If the sequence number is < 7, should be the first postbirth
        template, else should be the second postbirth template.
        """
        self.assertEqual(
            self.pmtct_postbirth2.get_template_name({"sequence_number": "6"}),
            "mc_postbirth1",
        )
        self.assertEqual(
            self.pmtct_postbirth2.get_template_name({"sequence_number": "7"}),
            "mc_postbirth2",
        )

    def test_get_template_variables_up_to_7_weeks(self):
        """
        Up to 7 weeks, the first variable should be the number of weeks, and the second
        should be the message content.
        """
        self.assertEqual(
            self.pmtct_postbirth2.get_template_variables(
                {"sequence_number": "5", "text_content": "Test message"}
            ),
            ["6", "Test message"],
        )

    def test_get_template_variables_over_7_weeks(self):
        """
        From 8 weeks onwards, the first variable should be the number of months, and
        the second should be the message content.
        """
        self.assertEqual(
            self.pmtct_postbirth2.get_template_variables(
                {"sequence_number": "7", "text_content": "Test message"}
            ),
            ["1", "Test message"],
        )

    def test_get_weeks_from_sequence_number(self):
        """
        Should return the correct number of weeks for the sequence number of the message
        """
        self.assertEqual(self.pmtct_postbirth2.get_weeks_from_sequence_number(1), 2)
        self.assertEqual(self.pmtct_postbirth2.get_weeks_from_sequence_number(2), 3)
        self.assertEqual(self.pmtct_postbirth2.get_weeks_from_sequence_number(6), 7)

    def test_get_months_from_sequence_number(self):
        """
        Returns correct number of months for the given sequence number
        """
        self.assertEqual(self.pmtct_postbirth2.get_months_from_sequence_number(1), 0)
        self.assertEqual(self.pmtct_postbirth2.get_months_from_sequence_number(7), 1)
        self.assertEqual(self.pmtct_postbirth2.get_months_from_sequence_number(50), 11)

If the sequence number is < 7, should be the first postbirth         template, else should be the second postbirth template. Up to 7 weeks, the first variable should be the number of weeks, and the second         should be the message content. From 8 weeks onwards, the first variable should be the number of months, and         the second should be the message content. Should return the correct number of weeks for the sequence number of the message Returns correct number of months for the given sequence number