def test_index_objects_public(self, get_es_mock):
user_1 = UserFactory.create()
user_2 = UserFactory.create()
model = MagicMock()
model.objects.privacy_level().filter.return_value = [
user_1.userprofile, user_2.userprofile
]
index_objects(model, [user_1.userprofile.id, user_2.userprofile.id],
True)
model.objects.assert_has_calls([
call.filter(id__in=[user_1.userprofile.id, user_2.userprofile.id]),
call.privacy_level(PUBLIC)
])
model.index.assert_has_calls([
call(model.extract_document(),
bulk=True,
id_=user_1.userprofile.id,
es=get_es_mock(),
public_index=True),
call(model.extract_document(),
bulk=True,
id_=user_2.userprofile.id,
es=get_es_mock(),
public_index=True)
])
model.refresh_index.assert_has_calls(
[call(es=get_es_mock()),
call(es=get_es_mock())])
def test_invalid_field(self):
"""
Test queryset when both includes and excludes are provided,
and some of the fields are not valid.
"""
# Given
request = self.factory.get(
self.url, {"q": "gold=5,invalid=5,gold!=5,invalid<>5"})
self.view.request = request
self.view.queryset = self.queryset_mock
expected_result = MagicMock()
self.queryset_mock.filter.return_value = self.queryset_mock
self.queryset_mock.exclude.return_value = expected_result
# When
result = LatestUnitVersionViewSet.get_queryset(self.view)
# Then
self.assertEqual(result, expected_result)
self.queryset_mock.assert_has_calls([
call.filter(gold=5),
call.exclude(gold=5),
])
def test_init(self):
with patch('airflow.settings.Session') as mock_session:
mock_session.return_value = UnifiedAlchemyMagicMock(
data=[([
call.query(Connection),
call.filter(Connection.conn_id == AZURE_CONN_ID)
], [AZURE_CONN]),
([
call.query(Connection),
call.filter(Connection.conn_id == HDI_CONN_ID)
], [HDI_CONN])])
op = ConnectedAzureHDInsightCreateClusterOperator(
cluster_name=CLUSTER_NAME,
azure_conn_id=AZURE_CONN_ID,
hdi_conn_id=HDI_CONN_ID,
task_id='foo',
dag=self.dag)
self.assertEqual(op.cluster_params, CLUSTER_PARAMS)
self.assertEqual(op.params, AZURE_CONN.extra_dejson)
def test_index_objects_public(self, get_es_mock):
user_1 = UserFactory.create()
user_2 = UserFactory.create()
mapping_type = MagicMock()
model = MagicMock()
mapping_type.get_model.return_value = model
qs = model.objects.filter().public_indexable().privacy_level
qs.return_value = [user_1.userprofile, user_2.userprofile]
mapping_type.extract_document.return_value = 'foo'
index_objects(mapping_type,
[user_1.userprofile.id, user_2.userprofile.id],
public_index=True)
model.objects.assert_has_calls([
call.filter(id__in=(user_1.userprofile.id, user_2.userprofile.id)),
call.filter().public_indexable(),
call.filter().public_indexable().privacy_level(PUBLIC),
])
mapping_type.bulk_index.assert_has_calls([
call(['foo', 'foo'], id_field='id', es=get_es_mock(),
index=mapping_type.get_index(True))])
def test_index_objects_public(self, get_es_mock):
user_1 = UserFactory.create()
user_2 = UserFactory.create()
mapping_type = MagicMock()
model = MagicMock()
mapping_type.get_model.return_value = model
qs = model.objects.filter().public_indexable().privacy_level
qs.return_value = [user_1.userprofile, user_2.userprofile]
mapping_type.extract_document.return_value = 'foo'
index_objects(mapping_type,
[user_1.userprofile.id, user_2.userprofile.id],
public_index=True)
model.objects.assert_has_calls([
call.filter(id__in=(user_1.userprofile.id, user_2.userprofile.id)),
call.filter().public_indexable(),
call.filter().public_indexable().privacy_level(PUBLIC),
])
mapping_type.bulk_index.assert_has_calls([
call(['foo', 'foo'], id_field='id', es=get_es_mock(),
index=mapping_type.get_index(True))])
def test_index_objects(self, get_es_mock):
user_1 = UserFactory.create()
user_2 = UserFactory.create()
model = MagicMock()
model.objects.filter.return_value = [user_1.userprofile, user_2.userprofile]
index_objects(model, [user_1.userprofile.id, user_2.userprofile.id], False)
model.objects.assert_has_calls([call.filter(id__in=[user_1.userprofile.id, user_2.userprofile.id])])
model.index.assert_has_calls(
[
call(
model.extract_document(), bulk=True, id_=user_1.userprofile.id, es=get_es_mock(), public_index=False
),
call(
model.extract_document(), bulk=True, id_=user_2.userprofile.id, es=get_es_mock(), public_index=False
),
]
)
model.refresh_index.assert_has_calls([call(es=get_es_mock()), call(es=get_es_mock())])
def test_excludes(self):
""" Test queryset when only excludes are provided. """
# Given
request = self.factory.get(self.url, {"q": "gold!=5"})
self.view.request = request
self.view.queryset = self.queryset_mock
expected_result = MagicMock()
self.queryset_mock.filter.return_value = self.queryset_mock
self.queryset_mock.exclude.return_value = expected_result
# When
result = LatestUnitVersionViewSet.get_queryset(self.view)
# Then
self.assertEqual(result, expected_result)
self.queryset_mock.assert_has_calls([
call.filter(),
call.exclude(gold=5),
])
def test_execute(self):
with patch('airflow.settings.Session') as mock_session:
with patch('airflowhdi.operators.azure_hdinsight_ssh_operator.SSHHook') as ssh_hook:
with patch('airflowhdi.operators.azure_hdinsight_ssh_operator.AzureHDInsightHook') as mock_hdi_hook:
with patch.object(SSHOperator, 'execute') as mock_ssh_op_exec:
ssh_endpoint = 'loc'
ssh_port = 523
mock_session.return_value = UnifiedAlchemyMagicMock(data=[
(
[call.query(Connection),
call.filter(Connection.conn_id == AZURE_CONN_ID)],
[AZURE_CONN]
)
])
mock_hdi_hook.return_value.get_cluster_state.return_value = ClusterGetProperties(
cluster_definition=None,
connectivity_endpoints=[ConnectivityEndpoint(
name='SSH',
location=ssh_endpoint,
port=ssh_port
)]
)
mock_hdi_hook.return_value.get_connection.return_value = AZURE_CONN
op = AzureHDInsightSshOperator(
command='date',
cluster_name=CLUSTER_NAME,
azure_conn_id=AZURE_CONN_ID,
task_id='foo',
dag=self.dag)
op.execute(None)
ssh_hook.assert_called_once_with(
remote_host=ssh_endpoint,
port=ssh_port,
username=AZURE_CONN.extra_dejson['SSH_USER_NAME'],
password=AZURE_CONN.extra_dejson['SSH_PASSWORD']
)
mock_ssh_op_exec.assert_called_once()
