Unlock 30% off on Manual Testing Annual Plans this Holiday Season.

Copied to Clipboard!

How to use mock_client method in tempest

Best Python code snippet using tempest_python

magics_impl_test.py

Source:magics_impl_test.py

1#!/usr/bin/env python2import collections3import io4import ipaddress5import time6from unittest import mock7from absl.testing import absltest8import pandas as pd9import grr_colab10from grr_colab import _timeout11from grr_colab import magics_impl12from grr_response_proto import artifact_pb213from grr_response_proto.api import client_pb214class GrrSetNoFlowTimeoutImplTest(absltest.TestCase):15  def testNone(self):16    with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):17      magics_impl.grr_set_no_flow_timeout_impl()18      self.assertIsNone(_timeout._FLOW_TIMEOUT)19class GrrSetDefaultFlowTimeoutImplTest(absltest.TestCase):20  def testDefault(self):21    with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 10):22      magics_impl.grr_set_default_flow_timeout_impl()23      self.assertEqual(_timeout._FLOW_TIMEOUT, 30)24class GrrSetFlowTimeoutImplTest(absltest.TestCase):25  def testNoneValueError(self):26    with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):27      with self.assertRaises(ValueError):28        magics_impl.grr_set_flow_timeout_impl(None)29  def testZero(self):30    with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):31      magics_impl.grr_set_flow_timeout_impl(0)32      self.assertEqual(_timeout._FLOW_TIMEOUT, 0)33  def testNonZero(self):34    with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):35      magics_impl.grr_set_flow_timeout_impl(10)36      self.assertEqual(_timeout._FLOW_TIMEOUT, 10)37  def testNegativeValueError(self):38    with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):39      with self.assertRaises(ValueError):40        magics_impl.grr_set_flow_timeout_impl(-10)41class GrrListArtifactsImplTest(absltest.TestCase):42  def testEmptyResults(self):43    with mock.patch.object(44        grr_colab, 'list_artifacts', return_value=[]) as mock_fn:45      magics_impl.grr_list_artifacts_impl()46      mock_fn.assert_called_once_with()47  def testPriorityColumns(self):48    artifact = artifact_pb2.ArtifactDescriptor()49    artifact.artifact.name = 'foo'50    artifact.artifact.doc = 'bar'51    artifact.is_custom = True52    with mock.patch.object(53        grr_colab, 'list_artifacts', return_value=[artifact]):54      df = magics_impl.grr_list_artifacts_impl()55    self.assertEqual((1, 3), df.shape)56    self.assertEqual(57        list(df.columns), ['artifact.name', 'artifact.doc', 'is_custom'])58    self.assertEqual(df['artifact.name'][0], 'foo')59    self.assertEqual(df['artifact.doc'][0], 'bar')60    self.assertTrue(df['is_custom'][0])61class GrrSearchClientsImplTest(absltest.TestCase):62  def testNoArgs(self):63    with mock.patch.object(64        grr_colab.Client, 'search', return_value=[]) as mock_fn:65      magics_impl.grr_search_clients_impl()66      mock_fn.assert_called_once_with(67          version=None, host=None, labels=None, mac=None, ip=None, user=None)68  def testWithArgs(self):69    with mock.patch.object(70        grr_colab.Client, 'search', return_value=[]) as mock_fn:71      magics_impl.grr_search_clients_impl(72          version='test_v',73          host='test_host',74          labels=['test_label'],75          mac='test_mac',76          ip='test_ip',77          user='test_user')78      mock_fn.assert_called_once_with(79          version='test_v',80          host='test_host',81          labels=['test_label'],82          mac='test_mac',83          ip='test_ip',84          user='test_user')85  def testIsDataframe(self):86    mock_client = _MockClient()87    with mock.patch.object(88        grr_colab.Client, 'search', return_value=[mock_client]):89      df = magics_impl.grr_search_clients_impl()90      self.assertIsInstance(df, pd.DataFrame)91      self.assertEqual(df.shape[0], 1)92  def testSortedByLastSeen(self):93    mock_clients = [94        _MockClient(client_id='foo', last_seen_at=1000),95        _MockClient(client_id='bar', last_seen_at=10),96        _MockClient(client_id='quux', last_seen_at=100),97    ]98    with mock.patch.object(99        grr_colab.Client, 'search', return_value=mock_clients):100      df = magics_impl.grr_search_clients_impl()101    self.assertEqual(df.shape[0], 3)102    self.assertEqual(list(df['client_id']), ['foo', 'quux', 'bar'])103  def testOnlineColumns(self):104    current_time_secs = int(time.time())105    mock_clients = [106        _MockClient(last_seen_at=(current_time_secs - 2 * 60) * (10**6)),107        _MockClient(last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)),108        _MockClient(109            last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)),110    ]111    with mock.patch.object(112        grr_colab.Client, 'search', return_value=mock_clients):113      df = magics_impl.grr_search_clients_impl()114    self.assertEqual(df.shape[0], 3)115    self.assertIn('online', df.columns)116    self.assertEqual(df['online'][0], 'online')117    self.assertEqual(df['online'][1], 'seen-1d')118    self.assertEqual(df['online'][2], 'offline')119    self.assertIn('online.pretty', df.columns)120    self.assertEqual(df['online.pretty'][0], 'ð')121    self.assertEqual(df['online.pretty'][1], 'ð')122    self.assertEqual(df['online.pretty'][2], 'ð')123  def testLastSeenAgoColumn(self):124    current_time_secs = 1560000000125    mock_clients = [126        _MockClient(last_seen_at=(current_time_secs + 1) * (10**6)),127        _MockClient(last_seen_at=(current_time_secs - 1) * (10**6)),128        _MockClient(last_seen_at=(current_time_secs - 2 * 60) * (10**6)),129        _MockClient(last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)),130        _MockClient(131            last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)),132    ]133    with mock.patch.object(time, 'time', return_value=current_time_secs):134      with mock.patch.object(135          grr_colab.Client, 'search', return_value=mock_clients):136        df = magics_impl.grr_search_clients_impl()137    self.assertEqual(df.shape[0], 5)138    self.assertIn('last_seen_ago', df.columns)139    self.assertEqual(df['last_seen_ago'][0], 'in 1 seconds')140    self.assertEqual(df['last_seen_ago'][1], '1 seconds ago')141    self.assertEqual(df['last_seen_ago'][2], '2 minutes ago')142    self.assertEqual(df['last_seen_ago'][3], '3 hours ago')143    self.assertEqual(df['last_seen_ago'][4], '4 days ago')144  def testPriorityColumns(self):145    mock_client = _MockClient(client_id='foo')146    mock_client.set_hostname('test_hostname')147    mock_client.set_os_version('test_version')148    with mock.patch.object(149        grr_colab.Client, 'search', return_value=[mock_client]):150      df = magics_impl.grr_search_clients_impl()151    self.assertGreaterEqual(df.shape[1], 7)152    self.assertEqual(153        list(df)[:7], [154            'online.pretty', 'online', 'client_id', 'last_seen_ago',155            'last_seen_at.pretty', 'knowledge_base.fqdn', 'os_info.version'156        ])157class GrrSearchOnlineClientsImplTest(absltest.TestCase):158  def testContainsOnlineOnly(self):159    current_time_secs = int(time.time())160    mock_clients = [161        _MockClient(162            client_id='foo',163            last_seen_at=(current_time_secs - 2 * 60) * (10**6)),164        _MockClient(165            client_id='bar',166            last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)),167        _MockClient(168            client_id='quux',169            last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)),170    ]171    with mock.patch.object(172        grr_colab.Client, 'search', return_value=mock_clients):173      df = magics_impl.grr_search_online_clients_impl()174    self.assertEqual(df.shape[0], 1)175    self.assertEqual(df['client_id'][0], 'foo')176    self.assertEqual(df['online'][0], 'online')177    self.assertEqual(df['online.pretty'][0], 'ð')178  def testEmptyResults(self):179    current_time_secs = int(time.time())180    mock_clients = [181        _MockClient(last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)),182        _MockClient(183            last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)),184    ]185    with mock.patch.object(186        grr_colab.Client, 'search', return_value=mock_clients):187      df = magics_impl.grr_search_online_clients_impl()188    self.assertEqual(df.shape[0], 0)189  def testSortedByLastSeen(self):190    current_time_secs = int(time.time())191    mock_clients = [192        _MockClient(193            client_id='foo', last_seen_at=(current_time_secs - 5) * (10**6)),194        _MockClient(195            client_id='bar', last_seen_at=(current_time_secs - 3) * (10**6)),196        _MockClient(197            client_id='quux', last_seen_at=(current_time_secs - 10) * (10**6)),198    ]199    with mock.patch.object(200        grr_colab.Client, 'search', return_value=mock_clients):201      df = magics_impl.grr_search_online_clients_impl()202    self.assertEqual(df.shape[0], 3)203    self.assertEqual(list(df['client_id']), ['bar', 'foo', 'quux'])204class GrrSetClientImplTest(absltest.TestCase):205  def testBothHostnameAndId(self):206    with mock.patch.object(grr_colab.Client, 'with_id') as with_id_fn:207      with mock.patch.object(grr_colab.Client,208                             'with_hostname') as with_hostname_fn:209        with mock.patch.object(magics_impl, '_state', magics_impl._State()):210          with self.assertRaises(ValueError):211            magics_impl.grr_set_client_impl('test_hostname', 'test_id')212          self.assertFalse(with_id_fn.called)213          self.assertFalse(with_hostname_fn.called)214  def testNeitherHostnameOrId(self):215    with mock.patch.object(grr_colab.Client, 'with_id') as with_id_fn:216      with mock.patch.object(grr_colab.Client,217                             'with_hostname') as with_hostname_fn:218        with mock.patch.object(magics_impl, '_state', magics_impl._State()):219          with self.assertRaises(ValueError):220            magics_impl.grr_set_client_impl()221          self.assertFalse(with_id_fn.called)222          self.assertFalse(with_hostname_fn.called)223  def testWithHostname(self):224    mock_client = _MockClient(client_id='bar')225    with mock.patch.object(grr_colab.Client, 'with_id') as with_id_fn:226      with mock.patch.object(227          grr_colab.Client, 'with_hostname',228          return_value=mock_client) as with_hostname_fn:229        with mock.patch.object(magics_impl, '_state', magics_impl._State()):230          magics_impl.grr_set_client_impl(hostname='foo')231          self.assertFalse(with_id_fn.called)232          with_hostname_fn.assert_called_once_with('foo')233          self.assertEqual(magics_impl._state.client.id, mock_client.id)234  def testWithId(self):235    mock_client = _MockClient(client_id='foo')236    with mock.patch.object(237        grr_colab.Client, 'with_id', return_value=mock_client) as with_id_fn:238      with mock.patch.object(grr_colab.Client,239                             'with_hostname') as with_hostname_fn:240        with mock.patch.object(magics_impl, '_state', magics_impl._State()):241          magics_impl.grr_set_client_impl(client='foo')242          self.assertFalse(with_hostname_fn.called)243          with_id_fn.assert_called_once_with('foo')244          self.assertEqual(magics_impl._state.client.id, mock_client.id)245class GrrRequestApprovalImplTest(absltest.TestCase):246  def testNoClientSelected(self):247    with self.assertRaises(magics_impl.NoClientSelectedError):248      magics_impl.grr_request_approval_impl('foo', ['bar'])249  def testNoWait(self):250    mock_client = mock.MagicMock()251    with mock.patch.object(magics_impl._state, 'client', mock_client):252      magics_impl.grr_request_approval_impl('foo', ['bar'])253      mock_client.request_approval.assert_called_once_with(254          reason='foo', approvers=['bar'])255      self.assertFalse(mock_client.request_approval_and_wait.called)256  def testWait(self):257    mock_client = mock.MagicMock()258    with mock.patch.object(magics_impl._state, 'client', mock_client):259      magics_impl.grr_request_approval_impl('foo', ['bar'], wait=True)260      mock_client.request_approval_and_wait.assert_called_once_with(261          reason='foo', approvers=['bar'])262      self.assertFalse(mock_client.request_approval.called)263class GrrIdImplTest(absltest.TestCase):264  def testNoClientSelected(self):265    with self.assertRaises(magics_impl.NoClientSelectedError):266      magics_impl.grr_id_impl()267  def testWithClientSet(self):268    mock_client = _MockClient(client_id='foo')269    with mock.patch.object(magics_impl._state, 'client', mock_client):270      client_id = magics_impl.grr_id_impl()271    self.assertEqual(client_id, 'foo')272class GrrCdImplTest(absltest.TestCase):273  def testNoClientSelected(self):274    with self.assertRaises(magics_impl.NoClientSelectedError):275      magics_impl.grr_cd_impl('foo')276  def testRelativePath(self):277    mock_client = _MockClient()278    with mock.patch.object(magics_impl._state, 'client', mock_client):279      with mock.patch.object(magics_impl._state, 'cur_dir', '/foo'):280        magics_impl.grr_cd_impl('bar/quux')281        self.assertEqual(magics_impl._state.cur_dir, '/foo/bar/quux')282  def testGoBack(self):283    mock_client = _MockClient()284    with mock.patch.object(magics_impl._state, 'client', mock_client):285      with mock.patch.object(magics_impl._state, 'cur_dir', '/foo/bar/quux'):286        magics_impl.grr_cd_impl('../..')287        self.assertEqual(magics_impl._state.cur_dir, '/foo')288  def testAbsolutePath(self):289    mock_client = _MockClient()290    with mock.patch.object(magics_impl._state, 'client', mock_client):291      with mock.patch.object(magics_impl._state, 'cur_dir', '/foo'):292        magics_impl.grr_cd_impl('/bar/quux')293        self.assertEqual(magics_impl._state.cur_dir, '/bar/quux')294class GrrPwdImplTest(absltest.TestCase):295  def testNoClientSelected(self):296    with self.assertRaises(magics_impl.NoClientSelectedError):297      magics_impl.grr_pwd_impl()298  def testRoot(self):299    mock_client = _MockClient()300    with mock.patch.object(magics_impl._state, 'client', mock_client):301      cur_dir = magics_impl.grr_pwd_impl()302    self.assertEqual(cur_dir, '/')303  def testNonRootPath(self):304    mock_client = _MockClient()305    with mock.patch.object(magics_impl._state, 'client', mock_client):306      with mock.patch.object(magics_impl._state, 'cur_dir', '/foo/bar'):307        cur_dir = magics_impl.grr_pwd_impl()308    self.assertEqual(cur_dir, '/foo/bar')309class GrrLsImplTest(absltest.TestCase):310  def testNoClientSelected(self):311    with self.assertRaises(magics_impl.NoClientSelectedError):312      magics_impl.grr_ls_impl('/foo')313  def testRelativePath(self):314    mock_client = mock.MagicMock()315    with mock.patch.object(magics_impl._state, 'client', mock_client):316      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):317        magics_impl.grr_ls_impl('foo/bar')318        mock_client.os.ls.assert_called_once_with('/quux/foo/bar')319        self.assertFalse(mock_client.os.cached.ls.called)320  def testAbsolutePath(self):321    mock_client = mock.MagicMock()322    with mock.patch.object(magics_impl._state, 'client', mock_client):323      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):324        magics_impl.grr_ls_impl('/foo/bar')325        mock_client.os.ls.assert_called_once_with('/foo/bar')326        self.assertFalse(mock_client.os.cached.ls.called)327  def testCurrentPath(self):328    mock_client = mock.MagicMock()329    with mock.patch.object(magics_impl._state, 'client', mock_client):330      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):331        magics_impl.grr_ls_impl()332        mock_client.os.ls.assert_called_once_with('/quux')333        self.assertFalse(mock_client.os.cached.ls.called)334  def testCached(self):335    mock_client = mock.MagicMock()336    with mock.patch.object(magics_impl._state, 'client', mock_client):337      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):338        magics_impl.grr_ls_impl('foo/bar', cached=True)339        mock_client.os.cached.ls.assert_called_once_with('/quux/foo/bar')340        self.assertFalse(mock_client.os.ls.called)341  def testTskPathType(self):342    mock_client = mock.MagicMock()343    with mock.patch.object(magics_impl._state, 'client', mock_client):344      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):345        magics_impl.grr_ls_impl('foo/bar', path_type=magics_impl.TSK)346        mock_client.tsk.ls.assert_called_once_with('/quux/foo/bar')347  def testNtfsPathType(self):348    mock_client = mock.MagicMock()349    with mock.patch.object(magics_impl._state, 'client', mock_client):350      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):351        magics_impl.grr_ls_impl('foo/bar', path_type=magics_impl.NTFS)352        mock_client.ntfs.ls.assert_called_once_with('/quux/foo/bar')353  def testRegistryPathType(self):354    mock_client = mock.MagicMock()355    with mock.patch.object(magics_impl._state, 'client', mock_client):356      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):357        magics_impl.grr_ls_impl('foo/bar', path_type=magics_impl.REGISTRY)358        mock_client.registry.ls.assert_called_once_with('/quux/foo/bar')359class GrrStatImplTest(absltest.TestCase):360  def testNoClientSelected(self):361    with self.assertRaises(magics_impl.NoClientSelectedError):362      magics_impl.grr_stat_impl('foo')363  def testRelativePath(self):364    mock_client = mock.MagicMock()365    with mock.patch.object(magics_impl._state, 'client', mock_client):366      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):367        magics_impl.grr_stat_impl('foo/bar')368        mock_client.os.glob.assert_called_once_with('/quux/foo/bar')369  def testAbsolutePath(self):370    mock_client = mock.MagicMock()371    with mock.patch.object(magics_impl._state, 'client', mock_client):372      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):373        magics_impl.grr_stat_impl('/foo/bar')374        mock_client.os.glob.assert_called_once_with('/foo/bar')375  def testTskPathType(self):376    mock_client = mock.MagicMock()377    with mock.patch.object(magics_impl._state, 'client', mock_client):378      magics_impl.grr_stat_impl('/foo/bar', path_type=magics_impl.TSK)379      mock_client.tsk.glob.assert_called_once_with('/foo/bar')380  def testNtfsPathType(self):381    mock_client = mock.MagicMock()382    with mock.patch.object(magics_impl._state, 'client', mock_client):383      magics_impl.grr_stat_impl('/foo/bar', path_type=magics_impl.NTFS)384      mock_client.ntfs.glob.assert_called_once_with('/foo/bar')385  def testRegistryPathType(self):386    mock_client = mock.MagicMock()387    with mock.patch.object(magics_impl._state, 'client', mock_client):388      magics_impl.grr_stat_impl('/foo/bar', path_type=magics_impl.REGISTRY)389      mock_client.registry.glob.assert_called_once_with('/foo/bar')390class GrrHeadImplTest(absltest.TestCase):391  def testNoClientSelected(self):392    with self.assertRaises(magics_impl.NoClientSelectedError):393      magics_impl.grr_head_impl('foo')394  def testRelativePath(self):395    mock_client = _MockClient()396    mock_client.os.add_file('/quux/foo/bar', b'foo bar')397    mock_client.os.cached.add_file('/quux/foo/bar', b'bar foo')398    with mock.patch.object(magics_impl._state, 'client', mock_client):399      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):400        data = magics_impl.grr_head_impl('foo/bar')401    self.assertEqual(data, b'foo bar')402  def testAbsolutePath(self):403    mock_client = _MockClient()404    mock_client.os.add_file('/foo/bar', b'foo bar')405    mock_client.os.cached.add_file('/foo/bar', b'bar foo')406    with mock.patch.object(magics_impl._state, 'client', mock_client):407      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):408        data = magics_impl.grr_head_impl('/foo/bar')409    self.assertEqual(data, b'foo bar')410  def testCached(self):411    mock_client = _MockClient()412    mock_client.os.add_file('/foo/bar', b'foo bar')413    mock_client.os.cached.add_file('/foo/bar', b'bar foo')414    with mock.patch.object(magics_impl._state, 'client', mock_client):415      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):416        data = magics_impl.grr_head_impl('/foo/bar', cached=True)417    self.assertEqual(data, b'bar foo')418  def testTskPathType(self):419    mock_client = _MockClient()420    mock_client.tsk.add_file('/foo/bar', b'foo bar')421    with mock.patch.object(magics_impl._state, 'client', mock_client):422      data = magics_impl.grr_head_impl('/foo/bar', path_type=magics_impl.TSK)423    self.assertEqual(data, b'foo bar')424  def testNtfsPathType(self):425    mock_client = _MockClient()426    mock_client.ntfs.add_file('/foo/bar', b'foo bar')427    with mock.patch.object(magics_impl._state, 'client', mock_client):428      data = magics_impl.grr_head_impl('/foo/bar', path_type=magics_impl.NTFS)429    self.assertEqual(data, b'foo bar')430  def testRegistryPathType(self):431    mock_client = _MockClient()432    mock_client.registry.add_file('/foo/bar', b'foo bar')433    with mock.patch.object(magics_impl._state, 'client', mock_client):434      data = magics_impl.grr_head_impl(435          '/foo/bar', path_type=magics_impl.REGISTRY)436    self.assertEqual(data, b'foo bar')437  def testBytes(self):438    mock_client = _MockClient()439    mock_client.os.add_file('/foo/bar', b'foo bar')440    mock_client.os.cached.add_file('/foo/bar', b'bar foo')441    with mock.patch.object(magics_impl._state, 'client', mock_client):442      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):443        data = magics_impl.grr_head_impl('/foo/bar', bytes=3)444    self.assertEqual(data, b'foo')445  def testOffset(self):446    mock_client = _MockClient()447    mock_client.os.add_file('/foo/bar', b'foo bar')448    mock_client.os.cached.add_file('/foo/bar', b'bar foo')449    with mock.patch.object(magics_impl._state, 'client', mock_client):450      with mock.patch.object(magics_impl._state, 'cur_dir', '/'):451        data = magics_impl.grr_head_impl('/foo/bar', bytes=3, offset=4)452    self.assertEqual(data, b'bar')453class GrrGrepImplTest(absltest.TestCase):454  def testNoClientSelected(self):455    with self.assertRaises(magics_impl.NoClientSelectedError):456      magics_impl.grr_grep_impl('foo', 'bar')457  def testRelativePath(self):458    mock_client = mock.MagicMock()459    with mock.patch.object(magics_impl._state, 'client', mock_client):460      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):461        magics_impl.grr_grep_impl('foo bar', 'foo/bar')462        mock_client.os.grep.assert_called_once_with('/quux/foo/bar', b'foo bar')463  def testAbsolutePath(self):464    mock_client = mock.MagicMock()465    with mock.patch.object(magics_impl._state, 'client', mock_client):466      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):467        magics_impl.grr_grep_impl('foo bar', '/foo/bar')468        mock_client.os.grep.assert_called_once_with('/foo/bar', b'foo bar')469  def testFixedStrings(self):470    mock_client = mock.MagicMock()471    with mock.patch.object(magics_impl._state, 'client', mock_client):472      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):473        magics_impl.grr_grep_impl('foo bar', '/foo/bar', fixed_strings=True)474        mock_client.os.fgrep.assert_called_once_with('/foo/bar', b'foo bar')475        self.assertFalse(mock_client.grep.called)476  def testTskPathType(self):477    mock_client = mock.MagicMock()478    with mock.patch.object(magics_impl._state, 'client', mock_client):479      magics_impl.grr_grep_impl(480          'foo bar', '/foo/bar', path_type=magics_impl.TSK)481      mock_client.tsk.grep.assert_called_once_with('/foo/bar', b'foo bar')482  def testNtfsPathType(self):483    mock_client = mock.MagicMock()484    with mock.patch.object(magics_impl._state, 'client', mock_client):485      magics_impl.grr_grep_impl(486          'foo bar', '/foo/bar', path_type=magics_impl.NTFS)487      mock_client.ntfs.grep.assert_called_once_with('/foo/bar', b'foo bar')488  def testRegistryPathType(self):489    mock_client = mock.MagicMock()490    with mock.patch.object(magics_impl._state, 'client', mock_client):491      magics_impl.grr_grep_impl(492          'foo bar', '/foo/bar', path_type=magics_impl.REGISTRY)493      mock_client.registry.grep.assert_called_once_with('/foo/bar', b'foo bar')494  def testHexString(self):495    mock_client = mock.MagicMock()496    with mock.patch.object(magics_impl._state, 'client', mock_client):497      magics_impl.grr_grep_impl('ffac90', '/foo/bar', hex_string=True)498      mock_client.os.grep.assert_called_once_with('/foo/bar', b'\xff\xac\x90')499class GrrFgrepImplTest(absltest.TestCase):500  def testNoClientSelected(self):501    with self.assertRaises(magics_impl.NoClientSelectedError):502      magics_impl.grr_fgrep_impl('foo', 'bar')503  def testRelativePath(self):504    mock_client = mock.MagicMock()505    with mock.patch.object(magics_impl._state, 'client', mock_client):506      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):507        magics_impl.grr_fgrep_impl('foo bar', 'foo/bar')508        mock_client.os.fgrep.assert_called_once_with('/quux/foo/bar',509                                                     b'foo bar')510  def testAbsolutePath(self):511    mock_client = mock.MagicMock()512    with mock.patch.object(magics_impl._state, 'client', mock_client):513      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):514        magics_impl.grr_fgrep_impl('foo bar', '/foo/bar')515        mock_client.os.fgrep.assert_called_once_with('/foo/bar', b'foo bar')516  def testTskPathType(self):517    mock_client = mock.MagicMock()518    with mock.patch.object(magics_impl._state, 'client', mock_client):519      magics_impl.grr_fgrep_impl(520          'foo bar', '/foo/bar', path_type=magics_impl.TSK)521      mock_client.tsk.fgrep.assert_called_once_with('/foo/bar', b'foo bar')522  def testNtfsPathType(self):523    mock_client = mock.MagicMock()524    with mock.patch.object(magics_impl._state, 'client', mock_client):525      magics_impl.grr_fgrep_impl(526          'foo bar', '/foo/bar', path_type=magics_impl.NTFS)527      mock_client.ntfs.fgrep.assert_called_once_with('/foo/bar', b'foo bar')528  def testRegistryPathType(self):529    mock_client = mock.MagicMock()530    with mock.patch.object(magics_impl._state, 'client', mock_client):531      magics_impl.grr_fgrep_impl(532          'foo bar', '/foo/bar', path_type=magics_impl.REGISTRY)533      mock_client.registry.fgrep.assert_called_once_with('/foo/bar', b'foo bar')534  def testHexString(self):535    mock_client = mock.MagicMock()536    with mock.patch.object(magics_impl._state, 'client', mock_client):537      magics_impl.grr_fgrep_impl('ffac90', '/foo/bar', hex_string=True)538      mock_client.os.fgrep.assert_called_once_with('/foo/bar', b'\xff\xac\x90')539class GrrInterrogateImplTest(absltest.TestCase):540  def testNoClientSelected(self):541    with self.assertRaises(magics_impl.NoClientSelectedError):542      magics_impl.grr_interrogate_impl()543  def testWithClientSet(self):544    mock_client = mock.MagicMock()545    with mock.patch.object(magics_impl._state, 'client', mock_client):546      magics_impl.grr_interrogate_impl()547      mock_client.interrogate.assert_called_once_with()548class GrrHostnameImplTest(absltest.TestCase):549  def testNoClientSelected(self):550    with self.assertRaises(magics_impl.NoClientSelectedError):551      magics_impl.grr_hostname_impl()552  def testWithClientSet(self):553    mock_client = _MockClient()554    mock_client.set_hostname('foobar')555    with mock.patch.object(magics_impl._state, 'client', mock_client):556      hostname = magics_impl.grr_hostname_impl()557    self.assertEqual(hostname, 'foobar')558class GrrIfconfigImplTest(absltest.TestCase):559  def testNoClientSelected(self):560    with self.assertRaises(magics_impl.NoClientSelectedError):561      magics_impl.grr_ifconfig_impl()562  def testWithClientSet(self):563    mock_client = _MockClient(ifaces=['foo', 'bar'])564    with mock.patch.object(magics_impl._state, 'client', mock_client):565      df = magics_impl.grr_ifconfig_impl()566    self.assertEqual(df.shape, (2, 1))567    self.assertEqual(list(df['ifname']), ['foo', 'bar'])568  def testPrettyIpAddress(self):569    mock_client = _MockClient()570    ipv4 = ipaddress.IPv4Address('42.0.255.32')571    ipv6 = ipaddress.IPv6Address('2001:db8::1000')572    mock_client.add_iface('foo', [ipv4.packed, ipv6.packed])573    mock_client.add_iface('bar', [])574    with mock.patch.object(magics_impl._state, 'client', mock_client):575      df = magics_impl.grr_ifconfig_impl()576    self.assertEqual(df.shape, (2, 2))577    self.assertEqual(list(df['ifname']), ['foo', 'bar'])578    self.assertTrue(pd.isna(df['addresses'][1]))579    addresses = df['addresses'][0]580    self.assertEqual(581        list(addresses['packed_bytes']), [ipv4.packed, ipv6.packed])582    self.assertEqual(583        list(addresses['packed_bytes.pretty']),584        [str(ipv4), str(ipv6)])585  def testPrettyMacAddress(self):586    mock_client = _MockClient()587    mock_client.add_iface('foo', mac=b'\xaa\x12\x42\xff\xa5\xd0')588    mock_client.add_iface('bar')589    with mock.patch.object(magics_impl._state, 'client', mock_client):590      df = magics_impl.grr_ifconfig_impl()591    self.assertEqual(df.shape, (2, 3))592    self.assertEqual(list(df['ifname']), ['foo', 'bar'])593    self.assertTrue(pd.isna(df['mac_address'][1]))594    self.assertTrue(pd.isna(df['mac_address.pretty'][1]))595    self.assertEqual(df['mac_address'][0], b'\xaa\x12\x42\xff\xa5\xd0')596    self.assertEqual(df['mac_address.pretty'][0], 'aa:12:42:ff:a5:d0')597class GrrUnameImplTest(absltest.TestCase):598  def testNoClientSelected(self):599    with self.assertRaises(magics_impl.NoClientSelectedError):600      magics_impl.grr_uname_impl()601  def testKernelRelease(self):602    mock_client = _MockClient(kernel='foobar')603    with mock.patch.object(magics_impl._state, 'client', mock_client):604      kernel_release = magics_impl.grr_uname_impl(kernel_release=True)605    self.assertEqual(kernel_release, 'foobar')606  def testMachine(self):607    mock_client = _MockClient(arch='foobar')608    with mock.patch.object(magics_impl._state, 'client', mock_client):609      machine = magics_impl.grr_uname_impl(machine=True)610    self.assertEqual(machine, 'foobar')611  def testNoOptionsProvided(self):612    mock_client = _MockClient()613    with mock.patch.object(magics_impl._state, 'client', mock_client):614      with self.assertRaises(ValueError):615        magics_impl.grr_uname_impl()616class GrrPsImplTest(absltest.TestCase):617  def testNoClientSelected(self):618    with self.assertRaises(magics_impl.NoClientSelectedError):619      magics_impl.grr_ps_impl()620  def testWithClientSet(self):621    mock_client = mock.MagicMock()622    with mock.patch.object(magics_impl._state, 'client', mock_client):623      magics_impl.grr_ps_impl()624      mock_client.ps.assert_called_once_with()625class GrrOsqueryiImplTest(absltest.TestCase):626  def testNoClientSelected(self):627    with self.assertRaises(magics_impl.NoClientSelectedError):628      magics_impl.grr_osqueryi_impl('foo bar')629  def testWithClientSet(self):630    mock_client = mock.MagicMock()631    with mock.patch.object(magics_impl._state, 'client', mock_client):632      magics_impl.grr_osqueryi_impl('foo bar')633      mock_client.osquery.assert_called_once_with('foo bar')634class GrrCollectImplTest(absltest.TestCase):635  def testNoClientSelected(self):636    with self.assertRaises(magics_impl.NoClientSelectedError):637      magics_impl.grr_collect_impl('FakeArtifact')638  def testWithClientSet(self):639    mock_client = mock.MagicMock()640    with mock.patch.object(magics_impl._state, 'client', mock_client):641      magics_impl.grr_collect_impl('FakeArtifact')642      mock_client.collect.assert_called_once_with('FakeArtifact')643class GrrYaraImplTest(absltest.TestCase):644  def testNoClientSelected(self):645    with self.assertRaises(magics_impl.NoClientSelectedError):646      magics_impl.grr_yara_impl('foo')647  def testWithClientSet(self):648    mock_client = mock.MagicMock()649    with mock.patch.object(magics_impl._state, 'client', mock_client):650      magics_impl.grr_yara_impl('foo', [42], 'bar')651      mock_client.yara.assert_called_once_with('foo', [42], 'bar')652class GrrWgetImplTest(absltest.TestCase):653  def testNoClientSelected(self):654    with self.assertRaises(magics_impl.NoClientSelectedError):655      magics_impl.grr_wget_impl('foo')656  def testRelativePath(self):657    mock_client = mock.MagicMock()658    with mock.patch.object(magics_impl._state, 'client', mock_client):659      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):660        magics_impl.grr_wget_impl('foo/bar')661        mock_client.os.wget.assert_called_once_with('/quux/foo/bar')662  def testAbsolutePath(self):663    mock_client = mock.MagicMock()664    with mock.patch.object(magics_impl._state, 'client', mock_client):665      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):666        magics_impl.grr_wget_impl('/foo/bar')667        mock_client.os.wget.assert_called_once_with('/foo/bar')668  def testCached(self):669    mock_client = mock.MagicMock()670    with mock.patch.object(magics_impl._state, 'client', mock_client):671      with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):672        magics_impl.grr_wget_impl('/foo/bar', cached=True)673        mock_client.os.cached.wget.assert_called_once_with('/foo/bar')674  def testTskPathType(self):675    mock_client = mock.MagicMock()676    with mock.patch.object(magics_impl._state, 'client', mock_client):677      magics_impl.grr_wget_impl('/foo/bar', path_type=magics_impl.TSK)678      mock_client.tsk.wget.assert_called_once_with('/foo/bar')679  def testNtfsPathType(self):680    mock_client = mock.MagicMock()681    with mock.patch.object(magics_impl._state, 'client', mock_client):682      magics_impl.grr_wget_impl('/foo/bar', path_type=magics_impl.NTFS)683      mock_client.ntfs.wget.assert_called_once_with('/foo/bar')684  def testRegistryPathType(self):685    mock_client = mock.MagicMock()686    with mock.patch.object(magics_impl._state, 'client', mock_client):687      magics_impl.grr_wget_impl('/foo/bar', path_type=magics_impl.REGISTRY)688      mock_client.registry.wget.assert_called_once_with('/foo/bar')689class _MockClient(grr_colab.Client):690  class MockInnerClient(object):691    def __init__(self):692      self.data = client_pb2.ApiClient()693  class MockVFS(object):694    def __init__(self):695      self.files = collections.defaultdict(bytes)696    def open(self, path):697      data = self.files[path]698      return io.BytesIO(data)699    def add_file(self, path, data):700      self.files[path] = data701  class MockFileSystem(object):702    def __init__(self):703      self._cached = _MockClient.MockVFS()704      self.files = collections.defaultdict(bytes)705    @property706    def cached(self):707      return self._cached708    def open(self, path):709      data = self.files[path]710      return io.BytesIO(data)711    def add_file(self, path, data):712      self.files[path] = data713  def __init__(self,714               client_id='',715               last_seen_at=0,716               ifaces=None,717               kernel='',718               arch=''):719    if ifaces is None:720      ifaces = []721    self._client = _MockClient.MockInnerClient()722    self._summary = None723    self._os = _MockClient.MockFileSystem()724    self._tsk = _MockClient.MockFileSystem()725    self._ntfs = _MockClient.MockFileSystem()726    self._registry = _MockClient.MockFileSystem()727    self.set_id(client_id)728    self.set_last_seen_at(last_seen_at)729    self.set_ifaces(ifaces)730    self.set_kernel(kernel)731    self.set_arch(arch)732  @property733  def os(self):734    return self._os735  @property736  def tsk(self):737    return self._tsk738  @property739  def ntfs(self):740    return self._ntfs741  @property742  def registry(self):743    return self._registry744  def set_id(self, client_id):745    self._client.data.client_id = client_id746    self._client.client_id = client_id747  def set_last_seen_at(self, ms):748    self._client.data.last_seen_at = ms749  def set_hostname(self, hostname):750    self._client.data.knowledge_base.fqdn = hostname751  def set_os_version(self, os_version):752    self._client.data.os_info.version = os_version753  def set_kernel(self, kernel):754    self._client.data.os_info.kernel = kernel755  def set_arch(self, arch):756    self._client.data.os_info.machine = arch757  def set_ifaces(self, ifnames):758    for ifname in ifnames:759      self.add_iface(ifname, [])760  def add_iface(self, ifname, addresses_bytes=None, mac=None):761    if addresses_bytes is None:762      addresses_bytes = []763    iface = self._client.data.interfaces.add()764    iface.ifname = ifname765    if mac is not None:766      iface.mac_address = mac767    for packed_bytes in addresses_bytes:768      addr = iface.addresses.add()769      addr.packed_bytes = packed_bytes770if __name__ == '__main__':...

test_influxdb.py

Source:test_influxdb.py

1"""The tests for the InfluxDB component."""2import datetime3import unittest4from unittest import mock5import influxdb as influx_client6from homeassistant.setup import setup_component7import homeassistant.components.influxdb as influxdb8from homeassistant.const import EVENT_STATE_CHANGED, STATE_OFF, STATE_ON, \9                                STATE_STANDBY10from tests.common import get_test_home_assistant11@mock.patch('influxdb.InfluxDBClient')12@mock.patch(13    'homeassistant.components.influxdb.InfluxThread.batch_timeout',14    mock.Mock(return_value=0))15class TestInfluxDB(unittest.TestCase):16    """Test the InfluxDB component."""17    def setUp(self):18        """Set up things to be run when tests are started."""19        self.hass = get_test_home_assistant()20        self.handler_method = None21        self.hass.bus.listen = mock.Mock()22    def tearDown(self):23        """Clear data."""24        self.hass.stop()25    def test_setup_config_full(self, mock_client):26        """Test the setup with full configuration."""27        config = {28            'influxdb': {29                'host': 'host',30                'port': 123,31                'database': 'db',32                'username': 'user',33                'password': 'password',34                'max_retries': 4,35                'ssl': 'False',36                'verify_ssl': 'False',37            }38        }39        assert setup_component(self.hass, influxdb.DOMAIN, config)40        self.assertTrue(self.hass.bus.listen.called)41        self.assertEqual(42            EVENT_STATE_CHANGED, self.hass.bus.listen.call_args_list[0][0][0])43        self.assertTrue(mock_client.return_value.query.called)44    def test_setup_config_defaults(self, mock_client):45        """Test the setup with default configuration."""46        config = {47            'influxdb': {48                'host': 'host',49                'username': 'user',50                'password': 'pass',51            }52        }53        assert setup_component(self.hass, influxdb.DOMAIN, config)54        self.assertTrue(self.hass.bus.listen.called)55        self.assertEqual(56            EVENT_STATE_CHANGED, self.hass.bus.listen.call_args_list[0][0][0])57    def test_setup_minimal_config(self, mock_client):58        """Test the setup with minimal configuration."""59        config = {60            'influxdb': {}61        }62        assert setup_component(self.hass, influxdb.DOMAIN, config)63    def test_setup_missing_password(self, mock_client):64        """Test the setup with existing username and missing password."""65        config = {66            'influxdb': {67                'username': 'user'68            }69        }70        assert not setup_component(self.hass, influxdb.DOMAIN, config)71    def test_setup_query_fail(self, mock_client):72        """Test the setup for query failures."""73        config = {74            'influxdb': {75                'host': 'host',76                'username': 'user',77                'password': 'pass',78            }79        }80        mock_client.return_value.query.side_effect = \81            influx_client.exceptions.InfluxDBClientError('fake')82        assert not setup_component(self.hass, influxdb.DOMAIN, config)83    def _setup(self, **kwargs):84        """Set up the client."""85        config = {86            'influxdb': {87                'host': 'host',88                'username': 'user',89                'password': 'pass',90                'exclude': {91                    'entities': ['fake.blacklisted'],92                    'domains': ['another_fake']93                }94            }95        }96        config['influxdb'].update(kwargs)97        assert setup_component(self.hass, influxdb.DOMAIN, config)98        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]99    def test_event_listener(self, mock_client):100        """Test the event listener."""101        self._setup()102        # map of HA State to valid influxdb [state, value] fields103        valid = {104            '1': [None, 1],105            '1.0': [None, 1.0],106            STATE_ON: [STATE_ON, 1],107            STATE_OFF: [STATE_OFF, 0],108            STATE_STANDBY: [STATE_STANDBY, None],109            'foo': ['foo', None]110        }111        for in_, out in valid.items():112            attrs = {113                'unit_of_measurement': 'foobars',114                'longitude': '1.1',115                'latitude': '2.2',116                'battery_level': '99%',117                'temperature': '20c',118                'last_seen': 'Last seen 23 minutes ago',119                'updated_at': datetime.datetime(2017, 1, 1, 0, 0),120                'multi_periods': '0.120.240.2023873'121            }122            state = mock.MagicMock(123                state=in_, domain='fake', entity_id='fake.entity-id',124                object_id='entity', attributes=attrs)125            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)126            body = [{127                'measurement': 'foobars',128                'tags': {129                    'domain': 'fake',130                    'entity_id': 'entity',131                },132                'time': 12345,133                'fields': {134                    'longitude': 1.1,135                    'latitude': 2.2,136                    'battery_level_str': '99%',137                    'battery_level': 99.0,138                    'temperature_str': '20c',139                    'temperature': 20.0,140                    'last_seen_str': 'Last seen 23 minutes ago',141                    'last_seen': 23.0,142                    'updated_at_str': '2017-01-01 00:00:00',143                    'updated_at': 20170101000000,144                    'multi_periods_str': '0.120.240.2023873'145                },146            }]147            if out[0] is not None:148                body[0]['fields']['state'] = out[0]149            if out[1] is not None:150                body[0]['fields']['value'] = out[1]151            self.handler_method(event)152            self.hass.data[influxdb.DOMAIN].block_till_done()153            self.assertEqual(154                mock_client.return_value.write_points.call_count, 1155            )156            self.assertEqual(157                mock_client.return_value.write_points.call_args,158                mock.call(body)159            )160            mock_client.return_value.write_points.reset_mock()161    def test_event_listener_no_units(self, mock_client):162        """Test the event listener for missing units."""163        self._setup()164        for unit in (None, ''):165            if unit:166                attrs = {'unit_of_measurement': unit}167            else:168                attrs = {}169            state = mock.MagicMock(170                state=1, domain='fake', entity_id='fake.entity-id',171                object_id='entity', attributes=attrs)172            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)173            body = [{174                'measurement': 'fake.entity-id',175                'tags': {176                    'domain': 'fake',177                    'entity_id': 'entity',178                },179                'time': 12345,180                'fields': {181                    'value': 1,182                },183            }]184            self.handler_method(event)185            self.hass.data[influxdb.DOMAIN].block_till_done()186            self.assertEqual(187                mock_client.return_value.write_points.call_count, 1188            )189            self.assertEqual(190                mock_client.return_value.write_points.call_args,191                mock.call(body)192            )193            mock_client.return_value.write_points.reset_mock()194    def test_event_listener_inf(self, mock_client):195        """Test the event listener for missing units."""196        self._setup()197        attrs = {'bignumstring':  '9' * 999, 'nonumstring': 'nan'}198        state = mock.MagicMock(199            state=8, domain='fake', entity_id='fake.entity-id',200            object_id='entity', attributes=attrs)201        event = mock.MagicMock(data={'new_state': state}, time_fired=12345)202        body = [{203            'measurement': 'fake.entity-id',204            'tags': {205                'domain': 'fake',206                'entity_id': 'entity',207            },208            'time': 12345,209            'fields': {210                'value': 8,211            },212        }]213        self.handler_method(event)214        self.hass.data[influxdb.DOMAIN].block_till_done()215        self.assertEqual(216            mock_client.return_value.write_points.call_count, 1217        )218        self.assertEqual(219            mock_client.return_value.write_points.call_args,220            mock.call(body)221        )222        mock_client.return_value.write_points.reset_mock()223    def test_event_listener_states(self, mock_client):224        """Test the event listener against ignored states."""225        self._setup()226        for state_state in (1, 'unknown', '', 'unavailable'):227            state = mock.MagicMock(228                state=state_state, domain='fake', entity_id='fake.entity-id',229                object_id='entity', attributes={})230            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)231            body = [{232                'measurement': 'fake.entity-id',233                'tags': {234                    'domain': 'fake',235                    'entity_id': 'entity',236                },237                'time': 12345,238                'fields': {239                    'value': 1,240                },241            }]242            self.handler_method(event)243            self.hass.data[influxdb.DOMAIN].block_till_done()244            if state_state == 1:245                self.assertEqual(246                    mock_client.return_value.write_points.call_count, 1247                )248                self.assertEqual(249                    mock_client.return_value.write_points.call_args,250                    mock.call(body)251                )252            else:253                self.assertFalse(mock_client.return_value.write_points.called)254            mock_client.return_value.write_points.reset_mock()255    def test_event_listener_blacklist(self, mock_client):256        """Test the event listener against a blacklist."""257        self._setup()258        for entity_id in ('ok', 'blacklisted'):259            state = mock.MagicMock(260                state=1, domain='fake', entity_id='fake.{}'.format(entity_id),261                object_id=entity_id, attributes={})262            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)263            body = [{264                'measurement': 'fake.{}'.format(entity_id),265                'tags': {266                    'domain': 'fake',267                    'entity_id': entity_id,268                },269                'time': 12345,270                'fields': {271                    'value': 1,272                },273            }]274            self.handler_method(event)275            self.hass.data[influxdb.DOMAIN].block_till_done()276            if entity_id == 'ok':277                self.assertEqual(278                    mock_client.return_value.write_points.call_count, 1279                )280                self.assertEqual(281                    mock_client.return_value.write_points.call_args,282                    mock.call(body)283                )284            else:285                self.assertFalse(mock_client.return_value.write_points.called)286            mock_client.return_value.write_points.reset_mock()287    def test_event_listener_blacklist_domain(self, mock_client):288        """Test the event listener against a blacklist."""289        self._setup()290        for domain in ('ok', 'another_fake'):291            state = mock.MagicMock(292                state=1, domain=domain,293                entity_id='{}.something'.format(domain),294                object_id='something', attributes={})295            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)296            body = [{297                'measurement': '{}.something'.format(domain),298                'tags': {299                    'domain': domain,300                    'entity_id': 'something',301                },302                'time': 12345,303                'fields': {304                    'value': 1,305                },306            }]307            self.handler_method(event)308            self.hass.data[influxdb.DOMAIN].block_till_done()309            if domain == 'ok':310                self.assertEqual(311                    mock_client.return_value.write_points.call_count, 1312                )313                self.assertEqual(314                    mock_client.return_value.write_points.call_args,315                    mock.call(body)316                )317            else:318                self.assertFalse(mock_client.return_value.write_points.called)319            mock_client.return_value.write_points.reset_mock()320    def test_event_listener_whitelist(self, mock_client):321        """Test the event listener against a whitelist."""322        config = {323            'influxdb': {324                'host': 'host',325                'username': 'user',326                'password': 'pass',327                'include': {328                    'entities': ['fake.included'],329                }330            }331        }332        assert setup_component(self.hass, influxdb.DOMAIN, config)333        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]334        for entity_id in ('included', 'default'):335            state = mock.MagicMock(336                state=1, domain='fake', entity_id='fake.{}'.format(entity_id),337                object_id=entity_id, attributes={})338            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)339            body = [{340                'measurement': 'fake.{}'.format(entity_id),341                'tags': {342                    'domain': 'fake',343                    'entity_id': entity_id,344                },345                'time': 12345,346                'fields': {347                    'value': 1,348                },349            }]350            self.handler_method(event)351            self.hass.data[influxdb.DOMAIN].block_till_done()352            if entity_id == 'included':353                self.assertEqual(354                    mock_client.return_value.write_points.call_count, 1355                )356                self.assertEqual(357                    mock_client.return_value.write_points.call_args,358                    mock.call(body)359                )360            else:361                self.assertFalse(mock_client.return_value.write_points.called)362            mock_client.return_value.write_points.reset_mock()363    def test_event_listener_whitelist_domain(self, mock_client):364        """Test the event listener against a whitelist."""365        config = {366            'influxdb': {367                'host': 'host',368                'username': 'user',369                'password': 'pass',370                'include': {371                    'domains': ['fake'],372                }373            }374        }375        assert setup_component(self.hass, influxdb.DOMAIN, config)376        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]377        for domain in ('fake', 'another_fake'):378            state = mock.MagicMock(379                state=1, domain=domain,380                entity_id='{}.something'.format(domain),381                object_id='something', attributes={})382            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)383            body = [{384                'measurement': '{}.something'.format(domain),385                'tags': {386                    'domain': domain,387                    'entity_id': 'something',388                },389                'time': 12345,390                'fields': {391                    'value': 1,392                },393            }]394            self.handler_method(event)395            self.hass.data[influxdb.DOMAIN].block_till_done()396            if domain == 'fake':397                self.assertEqual(398                    mock_client.return_value.write_points.call_count, 1399                )400                self.assertEqual(401                    mock_client.return_value.write_points.call_args,402                    mock.call(body)403                )404            else:405                self.assertFalse(mock_client.return_value.write_points.called)406            mock_client.return_value.write_points.reset_mock()407    def test_event_listener_invalid_type(self, mock_client):408        """Test the event listener when an attribute has an invalid type."""409        self._setup()410        # map of HA State to valid influxdb [state, value] fields411        valid = {412            '1': [None, 1],413            '1.0': [None, 1.0],414            STATE_ON: [STATE_ON, 1],415            STATE_OFF: [STATE_OFF, 0],416            STATE_STANDBY: [STATE_STANDBY, None],417            'foo': ['foo', None]418        }419        for in_, out in valid.items():420            attrs = {421                'unit_of_measurement': 'foobars',422                'longitude': '1.1',423                'latitude': '2.2',424                'invalid_attribute': ['value1', 'value2']425            }426            state = mock.MagicMock(427                state=in_, domain='fake', entity_id='fake.entity-id',428                object_id='entity', attributes=attrs)429            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)430            body = [{431                'measurement': 'foobars',432                'tags': {433                    'domain': 'fake',434                    'entity_id': 'entity',435                },436                'time': 12345,437                'fields': {438                    'longitude': 1.1,439                    'latitude': 2.2,440                    'invalid_attribute_str': "['value1', 'value2']"441                },442            }]443            if out[0] is not None:444                body[0]['fields']['state'] = out[0]445            if out[1] is not None:446                body[0]['fields']['value'] = out[1]447            self.handler_method(event)448            self.hass.data[influxdb.DOMAIN].block_till_done()449            self.assertEqual(450                mock_client.return_value.write_points.call_count, 1451            )452            self.assertEqual(453                mock_client.return_value.write_points.call_args,454                mock.call(body)455            )456            mock_client.return_value.write_points.reset_mock()457    def test_event_listener_default_measurement(self, mock_client):458        """Test the event listener with a default measurement."""459        config = {460            'influxdb': {461                'host': 'host',462                'username': 'user',463                'password': 'pass',464                'default_measurement': 'state',465                'exclude': {466                    'entities': ['fake.blacklisted']467                }468            }469        }470        assert setup_component(self.hass, influxdb.DOMAIN, config)471        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]472        for entity_id in ('ok', 'blacklisted'):473            state = mock.MagicMock(474                state=1, domain='fake', entity_id='fake.{}'.format(entity_id),475                object_id=entity_id, attributes={})476            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)477            body = [{478                'measurement': 'state',479                'tags': {480                    'domain': 'fake',481                    'entity_id': entity_id,482                },483                'time': 12345,484                'fields': {485                    'value': 1,486                },487            }]488            self.handler_method(event)489            self.hass.data[influxdb.DOMAIN].block_till_done()490            if entity_id == 'ok':491                self.assertEqual(492                    mock_client.return_value.write_points.call_count, 1493                )494                self.assertEqual(495                    mock_client.return_value.write_points.call_args,496                    mock.call(body)497                )498            else:499                self.assertFalse(mock_client.return_value.write_points.called)500            mock_client.return_value.write_points.reset_mock()501    def test_event_listener_unit_of_measurement_field(self, mock_client):502        """Test the event listener for unit of measurement field."""503        config = {504            'influxdb': {505                'host': 'host',506                'username': 'user',507                'password': 'pass',508                'override_measurement': 'state',509            }510        }511        assert setup_component(self.hass, influxdb.DOMAIN, config)512        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]513        attrs = {514            'unit_of_measurement': 'foobars',515        }516        state = mock.MagicMock(517            state='foo', domain='fake', entity_id='fake.entity-id',518            object_id='entity', attributes=attrs)519        event = mock.MagicMock(data={'new_state': state}, time_fired=12345)520        body = [{521            'measurement': 'state',522            'tags': {523                'domain': 'fake',524                'entity_id': 'entity',525            },526            'time': 12345,527            'fields': {528                'state': 'foo',529                'unit_of_measurement_str': 'foobars',530            },531        }]532        self.handler_method(event)533        self.hass.data[influxdb.DOMAIN].block_till_done()534        self.assertEqual(535            mock_client.return_value.write_points.call_count, 1536        )537        self.assertEqual(538            mock_client.return_value.write_points.call_args,539            mock.call(body)540        )541        mock_client.return_value.write_points.reset_mock()542    def test_event_listener_tags_attributes(self, mock_client):543        """Test the event listener when some attributes should be tags."""544        config = {545            'influxdb': {546                'host': 'host',547                'username': 'user',548                'password': 'pass',549                'tags_attributes': ['friendly_fake']550            }551        }552        assert setup_component(self.hass, influxdb.DOMAIN, config)553        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]554        attrs = {555            'friendly_fake': 'tag_str',556            'field_fake': 'field_str',557        }558        state = mock.MagicMock(559            state=1, domain='fake',560            entity_id='fake.something',561            object_id='something', attributes=attrs)562        event = mock.MagicMock(data={'new_state': state}, time_fired=12345)563        body = [{564            'measurement': 'fake.something',565            'tags': {566                'domain': 'fake',567                'entity_id': 'something',568                'friendly_fake': 'tag_str'569            },570            'time': 12345,571            'fields': {572                'value': 1,573                'field_fake_str': 'field_str'574            },575        }]576        self.handler_method(event)577        self.hass.data[influxdb.DOMAIN].block_till_done()578        self.assertEqual(579            mock_client.return_value.write_points.call_count, 1580        )581        self.assertEqual(582            mock_client.return_value.write_points.call_args,583            mock.call(body)584        )585        mock_client.return_value.write_points.reset_mock()586    def test_event_listener_component_override_measurement(self, mock_client):587        """Test the event listener with overridden measurements."""588        config = {589            'influxdb': {590                'host': 'host',591                'username': 'user',592                'password': 'pass',593                'component_config': {594                    'sensor.fake_humidity': {595                        'override_measurement': 'humidity'596                    }597                },598                'component_config_glob': {599                    'binary_sensor.*motion': {600                        'override_measurement': 'motion'601                    }602                },603                'component_config_domain': {604                    'climate': {605                        'override_measurement': 'hvac'606                    }607                }608            }609        }610        assert setup_component(self.hass, influxdb.DOMAIN, config)611        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]612        test_components = [613            {'domain': 'sensor', 'id': 'fake_humidity', 'res': 'humidity'},614            {'domain': 'binary_sensor', 'id': 'fake_motion', 'res': 'motion'},615            {'domain': 'climate', 'id': 'fake_thermostat', 'res': 'hvac'},616            {'domain': 'other', 'id': 'just_fake', 'res': 'other.just_fake'},617        ]618        for comp in test_components:619            state = mock.MagicMock(620                state=1, domain=comp['domain'],621                entity_id=comp['domain'] + '.' + comp['id'],622                object_id=comp['id'], attributes={})623            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)624            body = [{625                'measurement': comp['res'],626                'tags': {627                    'domain': comp['domain'],628                    'entity_id': comp['id']629                },630                'time': 12345,631                'fields': {632                    'value': 1,633                },634            }]635            self.handler_method(event)636            self.hass.data[influxdb.DOMAIN].block_till_done()637            self.assertEqual(638                mock_client.return_value.write_points.call_count, 1639            )640            self.assertEqual(641                mock_client.return_value.write_points.call_args,642                mock.call(body)643            )644            mock_client.return_value.write_points.reset_mock()645    def test_scheduled_write(self, mock_client):646        """Test the event listener to retry after write failures."""647        config = {648            'influxdb': {649                'host': 'host',650                'username': 'user',651                'password': 'pass',652                'max_retries': 1653            }654        }655        assert setup_component(self.hass, influxdb.DOMAIN, config)656        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]657        state = mock.MagicMock(658            state=1, domain='fake', entity_id='entity.id', object_id='entity',659            attributes={})660        event = mock.MagicMock(data={'new_state': state}, time_fired=12345)661        mock_client.return_value.write_points.side_effect = \662            IOError('foo')663        # Write fails664        with mock.patch.object(influxdb.time, 'sleep') as mock_sleep:665            self.handler_method(event)666            self.hass.data[influxdb.DOMAIN].block_till_done()667            assert mock_sleep.called668        json_data = mock_client.return_value.write_points.call_args[0][0]669        self.assertEqual(mock_client.return_value.write_points.call_count, 2)670        mock_client.return_value.write_points.assert_called_with(json_data)671        # Write works again672        mock_client.return_value.write_points.side_effect = None673        with mock.patch.object(influxdb.time, 'sleep') as mock_sleep:674            self.handler_method(event)675            self.hass.data[influxdb.DOMAIN].block_till_done()676            assert not mock_sleep.called677        self.assertEqual(mock_client.return_value.write_points.call_count, 3)678    def test_queue_backlog_full(self, mock_client):679        """Test the event listener to drop old events."""680        self._setup()681        state = mock.MagicMock(682            state=1, domain='fake', entity_id='entity.id', object_id='entity',683            attributes={})684        event = mock.MagicMock(data={'new_state': state}, time_fired=12345)685        monotonic_time = 0686        def fast_monotonic():687            """Monotonic time that ticks fast enough to cause a timeout."""688            nonlocal monotonic_time689            monotonic_time += 60690            return monotonic_time691        with mock.patch('homeassistant.components.influxdb.time.monotonic',692                        new=fast_monotonic):693            self.handler_method(event)694            self.hass.data[influxdb.DOMAIN].block_till_done()695            self.assertEqual(696                mock_client.return_value.write_points.call_count, 0697            )...

benchmark_uploader_test.py

Source:benchmark_uploader_test.py

1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14# ==============================================================================15"""Tests for benchmark_uploader."""16from __future__ import absolute_import17from __future__ import division18from __future__ import print_function19import json20import os21import tempfile22import unittest23from mock import MagicMock24from mock import patch25import tensorflow as tf  # pylint: disable=g-bad-import-order26try:27  from google.cloud import bigquery28  from official.benchmark import benchmark_uploader29except ImportError:30  bigquery = None31  benchmark_uploader = None32@unittest.skipIf(bigquery is None, "Bigquery dependency is not installed.")33class BigQueryUploaderTest(tf.test.TestCase):34  @patch.object(bigquery, "Client")35  def setUp(self, mock_bigquery):36    self.mock_client = mock_bigquery.return_value37    self.mock_dataset = MagicMock(name="dataset")38    self.mock_table = MagicMock(name="table")39    self.mock_client.dataset.return_value = self.mock_dataset40    self.mock_dataset.table.return_value = self.mock_table41    self.mock_client.insert_rows_json.return_value = []42    self.benchmark_uploader = benchmark_uploader.BigQueryUploader()43    self.benchmark_uploader._bq_client = self.mock_client44    self.log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())45    with open(os.path.join(self.log_dir, "metric.log"), "a") as f:46      json.dump({"name": "accuracy", "value": 1.0}, f)47      f.write("\n")48      json.dump({"name": "loss", "value": 0.5}, f)49      f.write("\n")50    with open(os.path.join(self.log_dir, "run.log"), "w") as f:51      json.dump({"model_name": "value"}, f)52  def tearDown(self):53    tf.gfile.DeleteRecursively(self.get_temp_dir())54  def test_upload_benchmark_run_json(self):55    self.benchmark_uploader.upload_benchmark_run_json(56        "dataset", "table", "run_id", {"model_name": "value"})57    self.mock_client.insert_rows_json.assert_called_once_with(58        self.mock_table, [{"model_name": "value", "model_id": "run_id"}])59  def test_upload_benchmark_metric_json(self):60    metric_json_list = [61        {"name": "accuracy", "value": 1.0},62        {"name": "loss", "value": 0.5}63    ]64    expected_params = [65        {"run_id": "run_id", "name": "accuracy", "value": 1.0},66        {"run_id": "run_id", "name": "loss", "value": 0.5}67    ]68    self.benchmark_uploader.upload_benchmark_metric_json(69        "dataset", "table", "run_id", metric_json_list)70    self.mock_client.insert_rows_json.assert_called_once_with(71        self.mock_table, expected_params)72  def test_upload_benchmark_run_file(self):73    self.benchmark_uploader.upload_benchmark_run_file(74        "dataset", "table", "run_id", os.path.join(self.log_dir, "run.log"))75    self.mock_client.insert_rows_json.assert_called_once_with(76        self.mock_table, [{"model_name": "value", "model_id": "run_id"}])77  def test_upload_metric_file(self):78    self.benchmark_uploader.upload_metric_file(79        "dataset", "table", "run_id",80        os.path.join(self.log_dir, "metric.log"))81    expected_params = [82        {"run_id": "run_id", "name": "accuracy", "value": 1.0},83        {"run_id": "run_id", "name": "loss", "value": 0.5}84    ]85    self.mock_client.insert_rows_json.assert_called_once_with(86        self.mock_table, expected_params)87  def test_insert_run_status(self):88    self.benchmark_uploader.insert_run_status(89        "dataset", "table", "run_id", "status")90    expected_query = ("INSERT dataset.table "91                      "(run_id, status) "92                      "VALUES('run_id', 'status')")93    self.mock_client.query.assert_called_once_with(query=expected_query)94  def test_update_run_status(self):95    self.benchmark_uploader.update_run_status(96        "dataset", "table", "run_id", "status")97    expected_query = ("UPDATE dataset.table "98                      "SET status = 'status' "99                      "WHERE run_id = 'run_id'")100    self.mock_client.query.assert_called_once_with(query=expected_query)101if __name__ == "__main__":...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.