Improve and expand unit tests for Person class

Herd_Immunity_Project/person.py

@@ -1,6 +1,6 @@
 import random
-import pytest
-# TODO: Import the virus class
+random.seed(42)
+from virus import Virus
 
 
 class Person(object):
@@ -14,45 +14,65 @@ def __init__(self, _id, is_vaccinated, infection=None):
         should instantiate a Virus object and set it as the value
         self.infection. Otherwise, self.infection should be set to None.
         '''
-        self._id = None
-        self.is_alive = True # Boolean
-        self.is_vaccinated = None # Boolean
-        self.infection = None # Virus or None
+        self._id = None  # int
+        self.is_alive = True  # boolean
+        self.is_vaccinated = None  # boolean
+        self.infection = None  # Virus object or None
 
     def did_survive_infection(self):
         ''' Generate a random number and compare to virus's mortality_rate.
         If random number is smaller, person dies from the disease.
         If Person survives, they become vaccinated and they have no infection.
+        Return a boolean value indicating whether they survived the infection.
         '''
         # Only called if infection attribute is not None.
         # TODO:  Finish this method. Should return a Boolean
         pass
 
-    
-''' These are simple tests to ensure that you are instantiating your Person class correctly. '''
+
+# simple unit tests to ensure your Person class initializer works correctly
 def test_vacc_person_instantiation():
     # create some people to test if our init method works as expected
     person = Person(1, True)
-
     assert person._id == 1
-    assert person.is_alive == True
-    assert person.is_vaccinated == True
-    assert person.infection == False
+    assert person.is_alive is True
+    assert person.is_vaccinated is True
+    assert person.infection is None
 
 
 def test_not_vacc_person_instantiation():
-    person = Person(1, False)
+    person = Person(2, False)
     # TODO: complete your own assert statements that test
     # the values at each attribute
-    pass
+    # assert ...
 
 
 def test_sick_person_instantiation():
-    # import the virus class for the testing
-    from virus import Virus
+    # Create a Virus object to give a Person object an infection
     virus = Virus("Dysentery", 0.7, 0.2)
-
+    # Create a Person object and give them the virus infection
+    person = Person(3, False, virus)
     # TODO: complete your own assert statements that test
     # the values at each attribute
-    person = Person(1, False, virus)
-    pass
+    # assert ...
+
+
+def test_did_survive_infection():
+    # TODO: Create a Virus object to give a Person object an infection
+    virus = Virus("Dysentery", 0.7, 0.2)
+    # TODO: Create a Person object and give them the virus infection
+    person = Person(4, False, virus)
+
+    # Resolve whether the Person survives the infection or not
+    survived = person.did_survive_infection()
+    # Check if the Person survived or not
+    if survived:
+        assert person.is_alive is True
+        # TODO: Write your own assert statements that test
+        # the values of each attribute for a Person who survived
+        # assert ...
+    else:
+        assert person.is_alive is False
+        # TODO: Write your own assert statements that test
+        # the values of each attribute for a Person who did not survive
+        # assert ...

Herd_Immunity_Project/virus.py

@@ -1,6 +1,3 @@
-#!/usr/bin/env python
-import pytest
-
 class Virus(object):
     '''Properties and attributes of the virus used in Simulation.'''