From abeb696a9dc7fef6a9b6ae334060ea1248fa6366 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Doma=C5=84ski?= <jandom@gmail.com>
Date: Tue, 20 Jul 2021 12:43:10 +0100
Subject: [PATCH] Refactor to use pytest over unittest (#1037)

---
 testing-unit-py/README.md        |  5 +--
 testing-unit-py/requirements.txt |  1 +
 testing-unit-py/test_ec2.py      | 60 ++++++++++++++++----------------
 3 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/testing-unit-py/README.md b/testing-unit-py/README.md
index 9c5479323..952caff85 100644
--- a/testing-unit-py/README.md
+++ b/testing-unit-py/README.md
@@ -9,13 +9,14 @@ An example of writing mock-based unit tests with both infrastructure definition
    ```bash
    $ python3 -m venv venv
    $ source venv/bin/activate
-   $ pip3 install -r requirements.txt
+   $ python -m pip install --upgrade pip
+   $ python -m pip install -r requirements.txt
    ```
 
 2.  Run the tests:
 
     ```
-    $ python -m unittest
+    $ python -m pytest # or simply `pytest`
 
     ------------------------------------------------------------
     Ran 2 tests in 0.004s
diff --git a/testing-unit-py/requirements.txt b/testing-unit-py/requirements.txt
index f676e7b4d..547b7addd 100644
--- a/testing-unit-py/requirements.txt
+++ b/testing-unit-py/requirements.txt
@@ -1,2 +1,3 @@
 pulumi>=3.5.1,<4.0.0
 pulumi-aws>=4.0.0,<5.0.0
+pytest
\ No newline at end of file
diff --git a/testing-unit-py/test_ec2.py b/testing-unit-py/test_ec2.py
index 742caf57b..4cc2a5690 100644
--- a/testing-unit-py/test_ec2.py
+++ b/testing-unit-py/test_ec2.py
@@ -24,33 +24,33 @@ def call(self, args: pulumi.runtime.MockCallArgs):
 # Now actually import the code that creates resources, and then test it.
 import infra
 
-class TestingWithMocks(unittest.TestCase):
-    # Test if the service has tags and a name tag.
-    @pulumi.runtime.test
-    def test_server_tags(self):
-        def check_tags(args):
-            urn, tags = args
-            self.assertIsNotNone(tags, f'server {urn} must have tags')
-            self.assertIn('Name', tags, 'server {urn} must have a name tag')
-
-        return pulumi.Output.all(infra.server.urn, infra.server.tags).apply(check_tags)
-
-    # Test if the instance is configured with user_data.
-    @pulumi.runtime.test
-    def test_server_userdata(self):
-        def check_user_data(args):
-            urn, user_data = args
-            self.assertFalse(user_data, f'illegal use of user_data on server {urn}')
-
-        return pulumi.Output.all(infra.server.urn, infra.server.user_data).apply(check_user_data)
-
-    # Test if port 22 for ssh is exposed.
-    @pulumi.runtime.test
-    def test_security_group_rules(self):
-        def check_security_group_rules(args):
-            urn, ingress = args
-            ssh_open = any([rule['from_port'] == 22 and any([block == "0.0.0.0/0" for block in rule['cidr_blocks']]) for rule in ingress])
-            self.assertFalse(ssh_open, f'security group {urn} exposes port 22 to the Internet (CIDR 0.0.0.0/0)')
-
-        # Return the results of the unit tests.
-        return pulumi.Output.all(infra.group.urn, infra.group.ingress).apply(check_security_group_rules)
+
+# Test if the service has tags and a name tag.
+@pulumi.runtime.test
+def test_server_tags():
+    def check_tags(args):
+        urn, tags = args
+        assert tags, f'server {urn} must have tags'
+        assert 'Name' in tags, 'server {urn} must have a name tag'
+
+    return pulumi.Output.all(infra.server.urn, infra.server.tags).apply(check_tags)
+
+# Test if the instance is configured with user_data.
+@pulumi.runtime.test
+def test_server_userdata():
+    def check_user_data(args):
+        urn, user_data = args
+        assert user_data == None, f'illegal use of user_data on server {urn}'
+
+    return pulumi.Output.all(infra.server.urn, infra.server.user_data).apply(check_user_data)
+
+# Test if port 22 for ssh is exposed.
+@pulumi.runtime.test
+def test_security_group_rules():
+    def check_security_group_rules(args):
+        urn, ingress = args
+        ssh_open = any([rule['from_port'] == 22 and any([block == "0.0.0.0/0" for block in rule['cidr_blocks']]) for rule in ingress])
+        assert ssh_open == False, f'security group {urn} exposes port 22 to the Internet (CIDR 0.0.0.0/0)'
+
+    # Return the results of the unit tests.
+    return pulumi.Output.all(infra.group.urn, infra.group.ingress).apply(check_security_group_rules)