#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Test for WorkerIdInterceptor"""
from __future__ import absolute_import

import collections
import logging
import unittest

import grpc

from apache_beam.runners.worker.token_auth_interceptor import TokenAuthInterceptor


class _ClientCallDetails(
    collections.namedtuple('_ClientCallDetails',
                           ('method', 'timeout', 'metadata', 'credentials')),
    grpc.ClientCallDetails):
    pass


class TokenAuthInterceptorTest(unittest.TestCase):

    def test_fs_token_insertion(self):
        token_key = 'fs_token'
        headers_holder = {}

        def continuation(client_details, request_iterator):
            headers_holder.update({
                token_key: dict(client_details.metadata).get(token_key)
            })

        TokenAuthInterceptor(token="test_token").intercept_stream_stream(continuation,
                                                      _ClientCallDetails(
                                                          None, None, None, None),
                                                      [])
        self.assertEqual(headers_holder[token_key], 'test_token',
                         'fs_token not set')

    def test_failure_when_fs_token_exists(self):
        token_key = 'fs_token'
        headers_holder = {}

        def continuation(client_details, request_iterator):
            headers_holder.update({
                token_key: dict(client_details.metadata).get(token_key)
            })

        with self.assertRaises(RuntimeError):
            TokenAuthInterceptor().intercept_stream_stream(
                continuation, _ClientCallDetails(None, None, {'fs_token': 'test_token'},
                                                 None), [])


if __name__ == '__main__':
    logging.getLogger().setLevel(logging.INFO)
    unittest.main()
