tests/test_graph_client.py — Data Room Toolkit

# By Zero Width Lab (zerowidthlab.com)
from __future__ import annotations
import contextlib
import io
import json
import os
import sys
import tempfile
import unittest
import urllib.error
from email.message import Message
from pathlib import Path
from unittest.mock import patch
__author__ = "Zero Width Lab"
TEST_URL ='https://graph.microsoft.com/v1.0/test'
PACKAGE_ROOT =Path(__file__).resolve().parent.parent
sys.path.insert(0,str(PACKAGE_ROOT / 'scripts'))
import graph_client as gc
from graph_client import GraphClientError as GraphError

class PathAndPageTests(unittest.TestCase):

 def test_paths(self):
  bad =('..','.','a//b','//a','a//','//','C:a','a\\b','%2e%2e/a','a/%2f../b','a/%5cb','a\n','a/%00')
  for path in bad:
   with self.subTest(path=path),self.assertRaises(gc.ConfigError):
    gc.load_config(None,env={},overrides={'root_path':path},require_target=False)
  self.assertEqual(gc.rooted_remote_path('Room','a b/file'),'Room/a b/file')
  for root in ('Room','/Room','Room/','/Room/'):
   self.assertEqual(gc.load_config(None,env={},overrides={'root_path':root},require_target=False)['root_path'],'Room')
  with self.assertRaisesRegex(gc.ConfigError,'folder.*\\\\u200d.*file') as raised:
   gc.rooted_remote_path('Room','folder/\u200dfile')
  self.assertNotIn('\u200d',str(raised.exception))

 def test_pages(self):
  url ='https://graph.microsoft.com/v1.0/drives/d/root:/Room:/children'
  for payload in (None,[],{},{'value':{}},{'value':[None]},{'value':[7]},{'value':[],'@odata.nextLink':7},{'value':[],'@odata.nextLink':url.replace('Room','Outside')}):
   with self.subTest(payload=payload),self.assertRaises(GraphError):
    gc.collection_page(payload,url)
  self.assertEqual(gc.collection_page({'value':[],'@odata.nextLink':url + '?page=2'},url),([],url + '?page=2'))

 def test_bounded_pages_private_outputs(self):
  seen =set()
  gc.check_collection_url('page1',seen)
  with self.assertRaisesRegex(GraphError,'repeated'):
   gc.check_collection_url('page1',seen)
  with patch.object(gc,'MAX_COLLECTION_PAGES',1),self.assertRaisesRegex(GraphError,'exceeds 1 pages'):
   gc.check_collection_url('page2',seen)
  with tempfile.TemporaryDirectory() as folder:
   root =Path(folder).resolve()
   output =root / 'report'
   gc.write_new_text(output,'café')
   self.assertEqual(output.read_text(),'café')
   self.assertEqual(output.stat().st_mode & 511,384)
   (root / 'link').symlink_to(root,target_is_directory=True)
   (root / 'file-link').symlink_to(output)
   for path in (output,root / 'file-link',root / 'link' / 'new',root / 'missing' / 'new'):
    with self.subTest(path=path),self.assertRaisesRegex(GraphError,'cannot create new output'):
     gc.write_new_text(path,'replacement')
   self.assertEqual(output.read_text(),'café')
   with self.assertRaisesRegex(GraphError,'use the physical path.*private/tmp'):
    gc.write_new_text(root / 'link' / 'new','x')
   self.assertFalse((root / 'new').exists())
   self.assertFalse((root / 'missing').exists())

class FakeResponse:

 def __init__(self,status,payload=None,headers=None):
  self.status =status
  self.headers =headers or {}
  self._body =payload if isinstance(payload,bytes) else b'' if payload is None else json.dumps(payload).encode('utf-8')

 def read(self):
  return self._body

 def getcode(self):
  return self.status

 def __enter__(self):
  return self

 def __exit__(self,exc_type,exc,traceback):
  return None

def http_error(status,payload=None,headers=None):
 message_headers =Message()
 for key,value in (headers or {}).items():
  message_headers[key] =value
 body =json.dumps(payload or {'error':{'code':'fixture'}}).encode('utf-8')
 return urllib.error.HTTPError(TEST_URL,status,'fixture error',message_headers,io.BytesIO(body))

class GraphClientTests(unittest.TestCase):

 def setUp(self):
  gc._fake_path =None
  gc._fake_responses =None
  gc.reset_contacted_hosts()

 @contextlib.contextmanager
 def fake_transport(self,*responses):
  with tempfile.TemporaryDirectory() as temp_dir:
   path =Path(temp_dir) / 'responses.json'
   path.write_text(json.dumps({'responses':[({'method':'GET','url_prefix':TEST_URL,'status':200} | r) for r in responses]}),encoding='utf-8')
   with patch.dict(os.environ,{'GRAPH_FAKE_RESPONSES':str(path),gc._FAKE_TRANSPORT_ACK_VAR:gc._FAKE_TRANSPORT_ACK_VALUE}):
    yield

 credentials ={'GRAPH_TENANT_ID':'tenant-test','GRAPH_CLIENT_ID':'client-test','GRAPH_CLIENT_SECRET':'secret-test'}

 @patch('graph_client._open_url')
 def test_token_success(self,transport):
  transport.return_value =FakeResponse(200,{'access_token':'fixture-token'})
  token =gc.acquire_token(env=self.credentials)
  self.assertEqual(token,'fixture-token')
  request_object =transport.call_args.args[0]
  self.assertEqual(request_object.get_method(),'POST')
  self.assertIn('grant_type=client_credentials',request_object.data.decode())
  self.assertIn('scope=https%3A%2F%2Fgraph.microsoft.com%2F.default',request_object.data.decode())

  transport.side_effect =http_error(401)
  with self.assertRaisesRegex(GraphError,'token acquisition failed with HTTP 401') as raised:
   gc.acquire_token(env=self.credentials)
  self.assertEqual(raised.exception.status,401)
  self.assertNotIn('secret-test',str(raised.exception))

  with self.assertRaisesRegex(gc.CredentialError,'GRAPH_TENANT_ID'):
   gc.load_credentials({})

 @patch('graph_client.time.sleep')
 @patch('graph_client._open_url')
 def test_retries(self,transport,sleep):
  ok,ra0 =(FakeResponse(200,{'value':[]}),{'Retry-After':'0'})
  unknown =gc.WriteOutcomeUnknown
  last ={'error':{'code':'last'}}
  cases =[(False,4,[http_error(429,headers=ra0),ok],200,2,[0.0],{'value':[]}),*((False,1,[http_error(code),ok],200,2,[1.0],None) for code in (429,507,501)),(False,1,[http_error(503,headers=ra0),ok],200,2,[0.0],None),(False,4,[http_error(403,{'error':{'code':'denied'}})],403,1,[],{'error':{'code':'denied'}}),(False,2,[http_error(503),http_error(503),http_error(503,last)],503,3,[1.0,2.0],last),(True,4,[http_error(429,headers=ra0),FakeResponse(201,{})],201,2,[0.0],None),(True,4,[http_error(503,headers=ra0),ok],200,2,[0.0],None),*((True,4,[http_error(code),ok],unknown,1,[],None) for code in (500,503)),(True,4,[http_error(429,headers=ra0),http_error(409)],unknown,2,[0.0],None),(True,4,[http_error(412)],412,1,[],None),(True,4,[urllib.error.URLError('reset')],unknown,1,[],None),(True,4,[TimeoutError('timed out')],unknown,1,[],None)]
  url =TEST_URL
  for index,(write,retries,responses,expected,calls,sleeps,body) in enumerate(cases):
   with self.subTest(case=index):
    transport.reset_mock()
    sleep.reset_mock()
    transport.side_effect =responses
    method ='PUT' if write else 'GET'
    if expected is unknown:
     with self.assertRaisesRegex(unknown,'outcome unknown'):
      gc.request(method,url,max_retries=retries,write=write)
    else:
     status,payload,headers =gc.request(method,url,max_retries=retries,write=write)
     self.assertEqual(status,expected)
     if index ==0:
      self.assertEqual(headers,{})
     if body is not None:
      self.assertEqual(payload,body)
    self.assertEqual(transport.call_count,calls)
    self.assertEqual([c.args[0] for c in sleep.call_args_list],sleeps)

 def test_environment_proxy(self):
  import subprocess
  probe ="import graph_client as gc, urllib.request as u; print([h.proxies['https'] for h in gc._DESTINATION_OPENER.handlers if type(h) is u.ProxyHandler])"
  env ={k:v for k,v in os.environ.items() if 'proxy' not in k.lower()}
  env['HTTPS_PROXY'] ='http://127.0.0.1:9'
  result =subprocess.run([sys.executable,'-c',probe],cwd=PACKAGE_ROOT / 'scripts',env=env,capture_output=True,text=True,check=True,timeout=10)
  self.assertEqual(result.stdout.strip(),"['http://127.0.0.1:9']")

 def test_encoding(self):
  self.assertEqual(gc.encode_path_segments('folder name/hash#100%/plain'),'folder%20name/hash%23100%25/plain')
  self.assertEqual(gc.encode_path_segments(['folder/name','hash# percent% space']),'folder%2Fname/hash%23%20percent%25%20space')

 def test_config(self):
  with tempfile.TemporaryDirectory() as folder:
   path =Path(folder) / 'config.json'
   path.write_text(json.dumps({'drive_id':'config-drive','site_id':'config-site','root_path':'config-root','max_retries':2}))
   config =gc.load_config(path,overrides={'drive_id':'flag-drive','root_path':None,'max_retries':7},env={})
   for key,value in {'drive_id':'flag-drive','site_id':'config-site','root_path':'config-root','max_retries':7,'request_timeout_seconds':60}.items():
    self.assertEqual(config[key],value)
   self.assertEqual(gc.load_config(None,env={'SHAREPOINT_DATAROOM_CONFIG':str(path)})['drive_id'],'config-drive')
   path.write_text('{')
   for target,message in ((path,'invalid JSON'),(path.with_name('missing'),'config file not found'),(None,'neither drive_id nor site_id')):
    with self.subTest(target=target),self.assertRaisesRegex(gc.ConfigError,message):
     gc.load_config(target,env={})

 @patch('graph_client._open_url')
 def test_site_drive(self,transport):
  transport.return_value =FakeResponse(200,{'id':'resolved-drive'})
  resolved =gc.resolve_drive_id('fixture-token',site_id='site,value')
  self.assertEqual(resolved,'resolved-drive')
  request_object =transport.call_args.args[0]
  self.assertIn('/sites/site%2Cvalue/drive',request_object.full_url)
  self.assertEqual(request_object.get_header('Authorization'),'Bearer fixture-token')
  self.assertEqual(request_object.unredirected_hdrs['Authorization'],'Bearer fixture-token')

 def test_upload_boundary(self):
  threshold =4 * 1024 * 1024
  self.assertFalse(gc.uses_chunked_upload(threshold,threshold))
  self.assertTrue(gc.uses_chunked_upload(threshold + 1,threshold))
  self.assertEqual(gc.upload_method(threshold,threshold),'direct')
  self.assertEqual(gc.upload_method(threshold + 1,threshold),'chunked')

 def test_fake_retries(self):
  with self.fake_transport({'status':429,'headers':{'Retry-After':'0'},'once':True},{'json':{'value':[]}}),patch('graph_client.time.sleep') as sleep:
   status,payload,_ =gc.request('GET',TEST_URL,max_retries=1)
  self.assertEqual((status,payload),(200,{'value':[]}))
  sleep.assert_called_once_with(0.0)

 def test_fake_unmatched(self):
  with self.fake_transport(),self.assertRaisesRegex(GraphError,'no fake response matched GET https://graph.microsoft.com/v1.0/missing'):
   gc.request('GET','https://graph.microsoft.com/v1.0/missing')

 def test_host_before_transport(self):
  blocked ='https://blocked.example.invalid/'
  with self.fake_transport({'url_prefix':blocked,'json':{}}),self.assertRaisesRegex(GraphError,'request destination host is not allowed: blocked.example.invalid'):
   gc.request('GET',blocked)
  self.assertEqual(gc.contacted_hosts(),[])
  self.assertIsNone(gc._fake_responses)

 def test_tracked_hosts(self):
  cases =[('GET',TEST_URL,False,200),('POST','https://login.microsoftonline.com/tenant-placeholder/token',False,200),('PUT','https://contoso.sharepoint.com/upload-session',True,202)]
  for method,url,issued,expected in cases:
   with self.subTest(url=url),self.fake_transport({'method':method,'url_prefix':url,'status':expected,'json':{}}):
    self.assertEqual(gc.request(method,url,graph_issued_url=issued)[0],expected)
  self.assertEqual(gc.contacted_hosts(),['contoso.sharepoint.com','graph.microsoft.com','login.microsoftonline.com'])

 def test_tenant_hosts(self):
  for url in ('https://contoso.sharepoint.com/sites/room/_api/upload','https://contoso-my.sharepoint.com/personal/upload','https://CONTOSO.SharePoint.com/download','https://contoso.sharepoint.com:443/download','https://graph.microsoft.com/v1.0/drives/d/items/i/content'):
   with self.subTest(url=url):
    host =gc._guard_destination(url,graph_issued_url=True)
    self.assertTrue(gc.is_graph_issued_host(host))

 @patch('graph_client._open_url')
 def test_host_refusals(self,transport):
  rejected ={'https://sharepoint.com/upload':'not allowed','https://a.b.sharepoint.com/upload':'not allowed','https://contoso.sharepoint.com.evil.example/upload':'not allowed','https://contoso-sharepoint.com/upload':'not allowed','https://contososharepoint.com/upload':'not allowed','https://-contoso.sharepoint.com/upload':'not allowed','https://content.example.invalid/upload':'not allowed','https://login.microsoftonline.com/upload':'not allowed','http://contoso.sharepoint.com/upload':'must use HTTPS','https://user!@contoso.sharepoint.com/upload':'user information','https://user:pass!@contoso.sharepoint.com/upload':'user information','https://evil.example!@contoso.sharepoint.com/upload':'user information','https://contoso.sharepoint.com:8443/upload':'default HTTPS port','https://graph.microsoft.com:444/v1.0/x':'default HTTPS port'}
  for url,message in rejected.items():
   with self.subTest(url=url):
    with self.assertRaisesRegex(GraphError,message):
     gc.request('PUT',url,graph_issued_url=True)
  transport.assert_not_called()
  self.assertEqual(gc.contacted_hosts(),[])

 def test_unflagged_hosts(self):
  with self.assertRaisesRegex(GraphError,'request destination host is not allowed: contoso.sharepoint.com'):
   gc._guard_destination('https://contoso.sharepoint.com/x',graph_issued_url=False)
  with self.assertRaisesRegex(GraphError,'user information'):
   gc._guard_destination('https://user!@login.microsoftonline.com/x',graph_issued_url=False)

 def test_redirects(self):
  handler =gc._DestinationRedirectHandler()
  for flagged in (False,True):
   first =gc.urllib.request.Request(TEST_URL)
   setattr(first,'_zwl_graph_issued_url',flagged)
   for target in ('https://blocked.example.invalid/content','https://content.example.invalid/download','https://contoso.sharepoint.com.evil.example/download','http://contoso.sharepoint.com/download','https://user!@contoso.sharepoint.com/download','https://contoso.sharepoint.com:8443/download'):
    with self.subTest(flagged=flagged,target=target),self.assertRaises(GraphError):
     handler.redirect_request(first,None,302,'redirect',{},target)
   if flagged:
    second =handler.redirect_request(first,None,302,'redirect',{},'https://contoso.sharepoint.com/download')
    self.assertIsNotNone(second)
    self.assertTrue(getattr(second,'_zwl_graph_issued_url'))
    with self.assertRaisesRegex(GraphError,'not allowed'):
     handler.redirect_request(second,None,302,'redirect',{},'https://content.example.invalid/download')
    self.assertEqual(gc.contacted_hosts(),['contoso.sharepoint.com'])
   else:
    self.assertEqual(gc.contacted_hosts(),[])
    with self.assertRaises(GraphError):
     handler.redirect_request(first,None,302,'redirect',{},'https://contoso.sharepoint.com/download')

 def test_fake_matching(self):
  body ={'item':{'mode':'fail'}}
  no_match =(GraphError,'no fake')
  with self.fake_transport({'method':'PUT','match_headers':{'If-Match':'etag-1'},'json':{'matched':'header'}},{'method':'POST','match_absent_headers':['If-Match'],'match_json':body,'json':{'matched':'body'}}):
   _,payload,_ =gc.request('PUT',TEST_URL,headers={'if-match':'etag-1'})
   self.assertEqual(payload,{'matched':'header'})
   with self.assertRaisesRegex(*no_match):
    gc.request('PUT',TEST_URL)
   _,payload,_ =gc.request('POST',TEST_URL,json_body=body)
   self.assertEqual(payload,{'matched':'body'})
   with self.assertRaisesRegex(*no_match):
    gc.request('POST',TEST_URL,headers={'If-Match':'etag-1'},json_body=body)

 def test_https_required(self):
  with self.assertRaisesRegex(GraphError,'request destination URL must use HTTPS'):
   gc.request('GET','http://graph.microsoft.com/v1.0/test')

 def test_redaction(self):
  self.assertEqual(gc._safe_url('https://graph.microsoft.com/v1.0/drives/private-drive/items/private-item'),'https://graph.microsoft.com/v1.0/drives/<redacted>/items/<redacted>')
  self.assertEqual(gc._safe_url('https://login.microsoftonline.com/tenant-value/oauth2/v2.0/token'),'https://login.microsoftonline.com/<redacted>/oauth2/v2.0/token')
  self.assertEqual(gc._safe_url('https://upload.example.invalid/private-session?token=value'),'https://upload.example.invalid/<redacted>')
if __name__ =='__main__':
 unittest.main()

License: MIT (read it).

Back to the listing