#!/usr/bin/env python3
"""Test script to verify the add stock response handling."""

def handle_add_stock_response(response_data):
    """
    Simulate the add stock response handling logic.

    Args:
        response_data: The response from POST /api/stock/

    Returns:
        Stock item ID
    """
    if isinstance(response_data, list):
        # If it's a list, get the first item
        stock_item = response_data[0] if response_data else {}
    else:
        # If it's a dict, use it directly
        stock_item = response_data

    sid = stock_item.get('pk', stock_item.get('id'))
    return sid


print("Testing add stock response handling:\n")

# Test 1: Response is a dict with 'pk' key
print("Test 1: API returns dict with 'pk'")
response = {'pk': 123, 'quantity': 150, 'part': 456, 'location': 476}
sid = handle_add_stock_response(response)
print(f"  Response: {response}")
print(f"  Stock ID: {sid}")
print(f"  Expected: 123, Got: {sid} - {'PASS' if sid == 123 else 'FAIL'}\n")

# Test 2: Response is a dict with 'id' key
print("Test 2: API returns dict with 'id'")
response = {'id': 456, 'quantity': 150, 'part': 789, 'location': 476}
sid = handle_add_stock_response(response)
print(f"  Response: {response}")
print(f"  Stock ID: {sid}")
print(f"  Expected: 456, Got: {sid} - {'PASS' if sid == 456 else 'FAIL'}\n")

# Test 3: Response is a list with dict containing 'pk'
print("Test 3: API returns list with dict containing 'pk'")
response = [{'pk': 789, 'quantity': 150, 'part': 123, 'location': 476}]
sid = handle_add_stock_response(response)
print(f"  Response: {response}")
print(f"  Stock ID: {sid}")
print(f"  Expected: 789, Got: {sid} - {'PASS' if sid == 789 else 'FAIL'}\n")

# Test 4: Response is a list with dict containing 'id'
print("Test 4: API returns list with dict containing 'id'")
response = [{'id': 999, 'quantity': 150, 'part': 123, 'location': 476}]
sid = handle_add_stock_response(response)
print(f"  Response: {response}")
print(f"  Stock ID: {sid}")
print(f"  Expected: 999, Got: {sid} - {'PASS' if sid == 999 else 'FAIL'}\n")

# Test 5: Empty response
print("Test 5: API returns empty dict")
response = {}
sid = handle_add_stock_response(response)
print(f"  Response: {response}")
print(f"  Stock ID: {sid}")
print(f"  Expected: None, Got: {sid} - {'PASS' if sid is None else 'FAIL'}\n")

print("All tests completed!")