More tests

test_urlchecker_db.py

+import unittest
+from urlchecker_db import LogDB
+import peewee
+
+
+class UrlCheckerNetworkTestCase(unittest.TestCase):
+    def setUp(self):
+        self.db = LogDB(":memory:")
+
+    def test_createTable(self):
+        assert self.db._db.pragma("user_version") == 0
+        self.db.createTables()
+        assert self.db._db.pragma("user_version") == 1
+
+
+def suite():
+    suite = unittest.TestSuite()
+    suite.addTest(unittest.makeSuite(UrlCheckerDBTestCase))
+    return suite
+
+
+if __name__ == "__main__":
+    unittest.main(defaultTest="suite")

urlchecker_db.py

@@ -47,8 +47,11 @@ class LogDB:
             transport = peewee.TextField()
             port = peewee.IntegerField()
             state = peewee.TextField()
+
             class Meta:
-                primary_key = peewee.CompositeKey("status", "ip", "transport", "port")
+                primary_key = peewee.CompositeKey(
+                    "status", "ip", "transport", "port"
+                )
 
         class DnsChange(BaseModel):
             status = peewee.ForeignKeyField(Status)
@@ -62,6 +65,7 @@ class LogDB:
             ip = peewee.TextField()
             url = peewee.TextField()
             status_code = peewee.IntegerField()
+
             class Meta:
                 primary_key = peewee.CompositeKey("status", "ip", "url")
 

urlchecker_http.py

@@ -21,12 +21,7 @@ def getUserAgent(version):
     )
 
 
-def request(
-    url,
-    headers=None,
-    session=requests,
-    version=0
-):
+def request(url, headers=None, session=requests, version=0):
 
     if headers is None:
         headers = {}
@@ -99,19 +94,16 @@ def checkHttpStatus(db, status_id, url, ip, bot_version):
         )
         session = requests.Session()
         session.mount(base_url, ForcedIPHTTPSAdapter(dest_ip=ip))
-        request_kw['session'] = session
+        request_kw["session"] = session
         ip_url = url
     elif parsed_url.scheme == "http":
         # Force IP location
         parsed_url = parsed_url._replace(netloc=ip)
         ip_url = parsed_url.geturl()
     else:
-        raise NotImplementedError('Unhandled url: %s' % url)
+        raise NotImplementedError("Unhandled url: %s" % url)
 
     response = request(
-        ip_url,
-        headers={"Host": hostname},
-        version=bot_version,
-        **request_kw
+        ip_url, headers={"Host": hostname}, version=bot_version, **request_kw
     )
     logHttpStatus(db, ip, url, response.status_code, status_id)

urlchecker_network.py

@@ -56,8 +56,8 @@ def isTcpPortOpen(db, ip, port, status_id):
             sock.close()
             raise
     except:
-      sock.close()
-      raise
+        sock.close()
+        raise
 
     sock.close()