1+ import unittest
2+ import os
3+ import sys
4+ from unittest .mock import patch , MagicMock
5+
6+ # Add the contentstack_management module to the path
7+ sys .path .insert (0 , os .path .join (os .path .dirname (__file__ ), '..' , '..' , '..' , '..' ))
8+
9+ import contentstack_management
10+ from contentstack_management .contentstack import Client
11+
12+
13+ class TOTPLoginTests (unittest .TestCase ):
14+ """Unit tests for TOTP login functionality in Contentstack Management Python SDK"""
15+
16+ def setUp (self ):
17+ """Set up test fixtures before each test method"""
18+ self .client = Client ()
19+ self .test_email = "test@example.com"
20+ self .test_password = "test_password"
21+ self .test_secret = "JBSWY3DPEHPK3PXP" # Standard test secret for TOTP
22+ self .test_tfa_token = "123456"
23+
24+ def tearDown (self ):
25+ """Clean up after each test method"""
26+ # Clean up environment variables
27+ if 'MFA_SECRET' in os .environ :
28+ del os .environ ['MFA_SECRET' ]
29+
30+ def test_login_method_signature_with_totp (self ):
31+ """Test that login method accepts TOTP parameters"""
32+ client = contentstack_management .Client ()
33+ # Test that the method exists and can be called with the expected parameters
34+ self .assertTrue (hasattr (client , 'login' ))
35+ self .assertTrue (callable (client .login ))
36+
37+ # Test that the method accepts TOTP parameters without error
38+ try :
39+ client .login (self .test_email , self .test_password , tfa_token = self .test_tfa_token )
40+ client .login (self .test_email , self .test_password , mfa_secret = self .test_secret )
41+ client .login (self .test_email , self .test_password , tfa_token = self .test_tfa_token , mfa_secret = self .test_secret )
42+ except Exception as e :
43+ self .fail (f"Login method should accept TOTP parameters without error: { e } )
44+
45+ def test_generate_totp_method (self ):
46+ """Test the _generate_totp method generates correct TOTP codes"""
47+ # Test with a known secret and verify the TOTP generation
48+ totp_code = self .client ._generate_totp (self .test_secret )
49+
50+ # Verify the TOTP code is a 6-digit string
51+ self .assertIsInstance (totp_code , str )
52+ self .assertEqual (len (totp_code ), 6 )
53+ self .assertTrue (totp_code .isdigit ())
54+
55+ def test_login_with_mfa_secret_generates_totp (self ):
56+ """Test that login with mfa_secret generates TOTP automatically"""
57+ with patch .object (self .client , 'client' ) as mock_client :
58+ mock_response = MagicMock ()
59+ mock_response .status_code = 200
60+ mock_response .json .return_value = {'user' : {'authtoken' : 'test_token' }}
61+ mock_client .post .return_value = mock_response
62+
63+ # Mock the UserSession class
64+ with patch ('contentstack_management.user_session.user_session.UserSession' ) as mock_user_session :
65+ mock_session_instance = MagicMock ()
66+ mock_session_instance .login .return_value = mock_response
67+ mock_user_session .return_value = mock_session_instance
68+
69+ # Mock the _generate_totp method to return a predictable value
70+ with patch .object (self .client , '_generate_totp' , return_value = '654321' ) as mock_generate_totp :
71+ result = self .client .login (
72+ self .test_email ,
73+ self .test_password ,
74+ mfa_secret = self .test_secret
75+ )
76+
77+ # Verify _generate_totp was called with the secret
78+ mock_generate_totp .assert_called_once_with (self .test_secret )
79+
80+ # Verify UserSession was called with generated TOTP
81+ mock_session_instance .login .assert_called_once_with (
82+ self .test_email ,
83+ self .test_password ,
84+ '654321'
85+ )
86+ self .assertEqual (result , mock_response )
87+
88+ def test_login_with_environment_variable (self ):
89+ """Test that login uses MFA_SECRET environment variable when mfa_secret is not provided"""
90+ # Set environment variable
91+ os .environ ['MFA_SECRET' ] = self .test_secret
92+
93+ with patch .object (self .client , 'client' ) as mock_client :
94+ mock_response = MagicMock ()
95+ mock_response .status_code = 200
96+ mock_response .json .return_value = {'user' : {'authtoken' : 'test_token' }}
97+ mock_client .post .return_value = mock_response
98+
99+ # Mock the UserSession class
100+ with patch ('contentstack_management.user_session.user_session.UserSession' ) as mock_user_session :
101+ mock_session_instance = MagicMock ()
102+ mock_session_instance .login .return_value = mock_response
103+ mock_user_session .return_value = mock_session_instance
104+
105+ # Mock the _generate_totp method
106+ with patch .object (self .client , '_generate_totp' , return_value = '789012' ) as mock_generate_totp :
107+ result = self .client .login (self .test_email , self .test_password )
108+
109+ # Verify _generate_totp was called with the environment secret
110+ mock_generate_totp .assert_called_once_with (self .test_secret )
111+
112+ # Verify UserSession was called with generated TOTP
113+ mock_session_instance .login .assert_called_once_with (
114+ self .test_email ,
115+ self .test_password ,
116+ '789012'
117+ )
118+ self .assertEqual (result , mock_response )
119+
120+ def test_backward_compatibility (self ):
121+ """Test that existing login patterns continue to work (backward compatibility)"""
122+ with patch .object (self .client , 'client' ) as mock_client :
123+ mock_response = MagicMock ()
124+ mock_response .status_code = 200
125+ mock_response .json .return_value = {'user' : {'authtoken' : 'test_token' }}
126+ mock_client .post .return_value = mock_response
127+
128+ # Mock the UserSession class
129+ with patch ('contentstack_management.user_session.user_session.UserSession' ) as mock_user_session :
130+ mock_session_instance = MagicMock ()
131+ mock_session_instance .login .return_value = mock_response
132+ mock_user_session .return_value = mock_session_instance
133+
134+ # Test old pattern: client.login(email, password)
135+ result1 = self .client .login (self .test_email , self .test_password )
136+
137+ # Test old pattern: client.login(email, password, tfa_token)
138+ result2 = self .client .login (self .test_email , self .test_password , self .test_tfa_token )
139+
140+ # Both should work without errors
141+ self .assertEqual (result1 , mock_response )
142+ self .assertEqual (result2 , mock_response )
143+
144+
145+ if __name__ == '__main__' :
146+ unittest .main ()
0 commit comments