Added network module and tests

2 files changed, 82 insertions, 1 deletions

diff --git a/machammer/network.py b/machammer/network.py
new file mode 100644
index 0000000..17608ce
--- /dev/null
+++ b/machammer/network.py
@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+"""Network-related machammer functions."""
+
+from .functions import call, check_output
+from .system_profiler import SystemProfile
+
+
+def get_ports(type='Ethernet'):
+    """Return all devices of type (Ethernet, AirPort)"""
+    profile = SystemProfile('Network')
+    return [x for x in profile if x['type'] == type]
+
+
+def set_wifi_power(on=True):
+    """Set AirPort power to on (True) or off (False)"""
+    state = 'on' if on else 'off'
+    for i in get_ports('AirPort'):
+        call('/usr/sbin/networksetup', '-setairportpower', i['interface'], state)
+
+
+def get_wifi_power():
+    """Get AirPort power state"""
+    results = []
+    for i in get_ports('AirPort'):
+        iface = i['interface']
+        r = check_output('/usr/sbin/networksetup', '-getairportpower', iface)
+        results.append(r.split(': ')[1])
+
+    return 'On' in results
+
+
+def is_wired(primary=False):
+    """Do we have an "active" Ethernet connection?"""
+    for p in get_ports('Ethernet'):
+        r = check_output('/sbin/ifconfig', p['interface'])
+        active = 'status: active' in r
+        return get_primary(p['interface']) if primary else active
+
+
+def is_wireless():
+    return not is_wired()
+
+
+def get_primary(port=None):
+    """Return device node of primary network interface"""
+    try:
+        route = check_output('/sbin/route', 'get', 'default').split('\n')
+    except Exception as e:
+        raise Exception('Failed to get default route: %s' % e)
+
+    for i in [x for x in route if 'interface: ' in x]:
+        p = i.split(': ')[1]
+        return p == port if port else p
diff --git a/tests.py b/tests.py
index bec8125..02bcaf3 100755
--- a/tests.py
+++ b/tests.py
@@ -1,11 +1,13 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+import time
 import logging
 import subprocess
 from unittest import main, skip, TestCase
 
-import machammer.functions as mh
+from machammer import network
+from machammer import functions as mh
 from machammer import system_profiler, screensaver
 
 
@@ -38,6 +40,28 @@ class SystemProfilerTestCase(TestCase):
         self.assertTrue(build in system_profiler.get('Software', 'os_version'))
 
 
+class NetworkTestCase(TestCase):
+    def test_wifi_enable(self):
+        """Turn wifi power on, check that it's on"""
+        network.set_wifi_power(True)
+        time.sleep(3)
+        self.assertTrue(network.get_wifi_power())
+
+    def test_wired(self):
+        self.assertTrue(network.is_wired())
+
+    def test_wifi_disable(self):
+        network.set_wifi_power(False)
+        time.sleep(3)
+        self.assertTrue(not network.get_wifi_power())
+
+    def test_primary(self):
+        self.assertEqual(network.get_primary(), 'en4')
+
+    def test_primary_wired(self):
+        self.assertTrue(network.is_wired(True))
+
+
 class AppsTestCase(TestCase):
     def setUp(self):
         self.profile = system_profiler.SystemProfile('Applications')
@@ -72,6 +96,10 @@ class FunctionsTestCase(TestCase):
         p = mh.mount_image('/Users/filipp/Downloads/AdobeFlashPlayer_22au_a_install.dmg')
         self.assertEquals(p, '/Volumes/Adobe Flash Player Installer')
 
+    @skip('This works, trust me.')
+    def test_sleep(self):
+        mh.sleep()
+
 
 class ScreenSaverTestCase(TestCase):
     def test_set_invalid(self):